From 947e8c3826929e3c2496de1b19ba906aeda3fef6 Mon Sep 17 00:00:00 2001 From: Chih-Yu Yeh Date: Mon, 17 Feb 2025 15:48:05 +0800 Subject: [PATCH] chore(wren-ai-service): support deleting sql pairs using project id (#1307) --- wren-ai-service/src/pipelines/indexing/sql_pairs.py | 5 ++++- .../src/web/v1/services/semantics_preparation.py | 6 ++++++ 2 files changed, 10 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 0afab4bf0..a5e8e6832 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -127,9 +127,10 @@ async def clean( sql_pairs: List[SqlPair], embedding: Dict[str, Any], project_id: Optional[str] = "", + delete_all: bool = False, ) -> Dict[str, Any]: sql_pair_ids = [sql_pair.id for sql_pair in sql_pairs] - if sql_pair_ids: + if sql_pair_ids or delete_all: await cleaner.run(sql_pair_ids=sql_pair_ids, project_id=project_id) return embedding @@ -216,12 +217,14 @@ async def clean( self, sql_pairs: List[SqlPair], project_id: Optional[str] = None, + delete_all: bool = False, ) -> None: await clean( sql_pairs=sql_pairs, embedding={"documents": []}, cleaner=self._components["cleaner"], project_id=project_id, + delete_all=delete_all, ) diff --git a/wren-ai-service/src/web/v1/services/semantics_preparation.py b/wren-ai-service/src/web/v1/services/semantics_preparation.py index 0be60a599..ffd444936 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -141,6 +141,12 @@ async def delete_semantics(self, project_id: str): tasks = [ self._pipelines[name].clean(project_id=project_id) for name in ["db_schema", "historical_question", "table_description"] + ] + [ + self._pipelines["sql_pairs"].clean( + sql_pairs=[], + project_id=project_id, + delete_all=True, + ) ] await asyncio.gather(*tasks)