feat: configure embedding retrieval strategies at runtime

Author: ulleo

AGENTS.md

@@ -62,3 +62,20 @@ SQLBot 让业务用户用自然语言提问，基于已配置的数据源生成
 - 按验收条件选择验证层级，不把构建或单元测试等同于产品验收。普通任务不默认构建镜像或启动完整运行栈；已要求真实接口、浏览器或部署验证时，可启动必要的隔离服务。环境范围、外部成本或数据用途不明确时先确认。
 - 除非用户明确要求，不要上传、发布或推送镜像 / wheel / 包。
 - 变更涉及本文件或 `docs/agents/` 描述的约定（目录职责、命令、流程）时，同步更新对应文档。
+
+## Agent skills
+
+### Issue tracker
+
+处理 Issue 或发布规格、任务时，使用 dataease/SQLBot 的 GitHub Issues；
+操作约定见 `docs/agents/issue-tracker.md`。
+
+### Triage labels
+
+执行 triage 时，使用五个默认角色标签；
+映射见 `docs/agents/triage-labels.md`。
+
+### Domain docs
+
+工程技能采用 single-context 布局；探索相关领域前，
+按 `docs/agents/domain.md` 读取词汇表和架构决策。

CONTEXT.md

@@ -54,6 +54,18 @@ _Avoid_: Terminology、SQL example
 术语、SQL 示例和自定义提示词生效的边界：工作空间、数据源或高级应用。高级应用是独立知识域，不继承工作空间级知识。
 _Avoid_: 权限——作用域决定向模型注入哪些知识，权限决定用户能访问哪些数据
 
+**Knowledge Retrieval / 知识召回**:
+根据问题匹配相关术语和 SQL 示例，为问数提供业务解释与查询示例。字面匹配与向量匹配是补充关系。
+_Avoid_: 自定义提示词注入、全系统向量化
+
+**Table Selection / 表筛选**:
+从当前可访问的候选表中选择与问题相关的表，为生成 SQL 提供表结构上下文。
+_Avoid_: 表权限——筛选不授予访问权限
+
+**Datasource Preselection / 数据源预筛选**:
+在自动选择数据源之前，缩小与问题相关的数据源候选范围。
+_Avoid_: 数据源选择结果——预筛选结果仍是候选集合
+
 **Row Permission / 行权限**（xpack）:
 对表追加的行级过滤条件；同一表命中多条时取 AND，经改写 SQL 执行。
 _Avoid_: Column Permission

backend/apps/chat/task/llm.py

@@ -52,6 +52,7 @@
 from apps.terminology.curd.terminology import expand_with_terminology
 from apps.terminology.curd.terminology import get_terminology_template
 from common.core.config import settings
+from common.core.embedding_config import load_embedding_config
 from common.core.db import engine
 from common.core.deps import CurrentAssistant, CurrentUser
 from common.error import SingleMessageError, SQLBotDBError, ParseSQLResultError, SQLBotDBConnectionError
@@ -125,6 +126,7 @@ class LLMService:
     def __init__(self, session: Session, current_user: CurrentUser, chat_question: ChatQuestion,
                  current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False,
                  embedding: bool = False, config: LLMConfig = None):
+        self.embedding_config = load_embedding_config(session)
         self.sql_message = []
         self.chart_message = []
         self.generate_sql_logs = []
@@ -155,6 +157,9 @@ def __init__(self, session: Session, current_user: CurrentUser, chat_question: C
                 raise SingleMessageError("Current assistant user cannot not access this chat")
 
         self.current_user.oid = chat.oid
+        SQLBotLogUtil.info(
+            f"[embedding] stage=config chat_id={chat_id} config={self.embedding_config}"
+        )
         ds: CoreDatasource | AssistantOutDsSchema | None = None
         if not chat.datasource and chat_question.datasource_id:
             _ds = session.get(CoreDatasource, chat_question.datasource_id)
@@ -386,13 +391,19 @@ def filter_terminology_template(self, _session: Session, oid: int = None, ds_id:
             self.chat_question.terminologies, term_list = get_terminology_template(_session,
                                                                                    match_text,
                                                                                    calculate_oid,
-                                                                                   None, self.current_assistant.id)
+                                                                                   None, self.current_assistant.id, config=self.embedding_config)
         else:
             self.chat_question.terminologies, term_list = get_terminology_template(_session,
                                                                                    match_text,
                                                                                    calculate_oid,
-                                                                                   calculate_ds_id)
+                                                                                   calculate_ds_id, config=self.embedding_config)
 
+        SQLBotLogUtil.info(
+            f"[embedding] stage=terminology chat_id={self.chat_question.chat_id} "
+            f"record_id={self.record.id} "
+            f"retrieval_mode={'literal_and_vector' if self.embedding_config.enabled else 'literal_only'} "
+            f"result_count={len(term_list)}"
+        )
         self.current_logs[OperationEnum.FILTER_TERMS] = end_log(session=_session,
                                                                 log=self.current_logs[OperationEnum.FILTER_TERMS],
                                                                 full_message=term_list)
@@ -440,12 +451,18 @@ def filter_training_template(self, _session: Session, oid: int = None, ds_id: in
             self.chat_question.data_training, example_list = get_training_template(_session,
                                                                                    self.chat_question.question,
                                                                                    calculate_oid,
-                                                                                   None, self.current_assistant.id)
+                                                                                   None, self.current_assistant.id, config=self.embedding_config)
         else:
             self.chat_question.data_training, example_list = get_training_template(_session,
                                                                                    self.chat_question.question,
                                                                                    calculate_oid,
-                                                                                   calculate_ds_id)
+                                                                                   calculate_ds_id, config=self.embedding_config)
+        SQLBotLogUtil.info(
+            f"[embedding] stage=sql_example chat_id={self.chat_question.chat_id} "
+            f"record_id={self.record.id} "
+            f"retrieval_mode={'literal_and_vector' if self.embedding_config.enabled else 'literal_only'} "
+            f"result_count={len(example_list)}"
+        )
         self.current_logs[OperationEnum.FILTER_SQL_EXAMPLE] = end_log(session=_session,
                                                                       log=self.current_logs[
                                                                           OperationEnum.FILTER_SQL_EXAMPLE],
@@ -552,7 +569,15 @@ def choose_table_schema(self, _session: Session):
             current_user=self.current_user,
             ds=self.ds,
             question=self.chat_question.question,
-            keywords=keywords)
+            keywords=keywords, config=self.embedding_config)
+
+        SQLBotLogUtil.info(
+            f"[embedding] stage=table_context chat_id={self.chat_question.chat_id} "
+            f"record_id={self.record.id} ds_id={self.ds.id} "
+            f"source={'dynamic_datasource' if self.out_ds_instance else 'configured_datasource'} "
+            f"vector_enabled={self.embedding_config.table_enabled and not self.out_ds_instance} "
+            f"result_count={len(tables)}"
+        )
 
         # Get sample data for all tables
         if not self.out_ds_instance and settings.SAMPLE_DATA_ENABLED:
@@ -767,16 +792,28 @@ def select_datasource(self, _session: Session):
         if not _ds_list:
             raise SingleMessageError('No available datasource configuration found')
         ignore_auto_select = _ds_list and len(_ds_list) == 1
+        if ignore_auto_select:
+            SQLBotLogUtil.info(
+                f"[embedding] stage=datasource chat_id={self.chat_question.chat_id} "
+                f"record_id={self.record.id} action=skip reason=single_candidate candidate_count=1"
+            )
         # ignore auto select ds
 
         full_thinking_text = ''
         full_text = ''
         if not ignore_auto_select:
-            if settings.TABLE_EMBEDDING_ENABLED and (
-                    not self.current_assistant or (self.current_assistant and self.current_assistant.type != 1)):
+            if self.embedding_config.table_enabled and (
+                    not self.current_assistant or (self.current_assistant and self.current_assistant.type not in dynamic_ds_types)):
                 _ds_list = get_ds_embedding(_session, _ds_list, self.out_ds_instance,
-                                            self.chat_question.question, self.current_assistant)
+                                            self.chat_question.question, self.current_assistant, config=self.embedding_config)
                 # yield {'content': '{"id":' + str(ds.get('id')) + '}'}
+            else:
+                reason = 'disabled' if not self.embedding_config.table_enabled else 'dynamic_datasource'
+                SQLBotLogUtil.info(
+                    f"[embedding] stage=datasource chat_id={self.chat_question.chat_id} "
+                    f"record_id={self.record.id} action=skip reason={reason} "
+                    f"candidate_count={len(_ds_list)} result_count={len(_ds_list)}"
+                )
 
             _ds_list_dict = []
             for _ds in _ds_list:
@@ -873,7 +910,7 @@ def select_datasource(self, _session: Session):
         except Exception as e:
             _error = e
 
-        if not ignore_auto_select and not settings.TABLE_EMBEDDING_ENABLED:
+        if not ignore_auto_select and not self.embedding_config.table_enabled:
             self.record = save_select_datasource_answer(session=_session, record_id=self.record.id,
                                                         answer=orjson.dumps({'content': full_text}).decode(),
                                                         datasource=_datasource,
@@ -883,7 +920,7 @@ def select_datasource(self, _session: Session):
             ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
 
             # 提取关键词（前置步骤，结果复用于术语匹配和表匹配）
-            if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
+            if self.embedding_config.keyword_enabled:
                 is_error, kw_result = self.extract_keywords(_session)
                 if is_error:
                     raise SingleMessageError(kw_result)
@@ -1365,7 +1402,7 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
                 ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
 
                 # 提取关键词（前置步骤，结果复用于术语匹配和表匹配）
-                if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
+                if self.embedding_config.keyword_enabled:
                     is_error, kw_result = self.extract_keywords(_session)
                     if is_error:
                         raise SingleMessageError(kw_result)

backend/apps/data_training/curd/data_training.py

@@ -3,17 +3,19 @@
 from typing import List, Optional
 from xml.dom.minidom import parseString
 
-from sqlalchemy import and_, select, func, delete, update, or_
+from sqlalchemy import case, and_, select, func, delete, update, or_
 from sqlalchemy import text
 
 from apps.ai_model.embedding import EmbeddingModelCache
+from apps.datasource.embedding.utils import parse_valid_embedding
 from apps.data_training.models.data_training_model import DataTrainingInfo, DataTraining, DataTrainingInfoResult
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AssistantModel
 from apps.template.generate_chart.generator import get_base_data_training_template
-from common.core.config import settings
+from common.core.embedding_config import EmbeddingConfig, load_embedding_config
 from common.core.deps import SessionDep, Trans
 from common.utils.dict_to_xml import dict_to_xml
+from common.utils.embedding_lifecycle import generate_missing_embeddings
 from common.utils.embedding_threads import run_save_data_training_embeddings
 
 
@@ -266,6 +268,7 @@ def update_training(session: SessionDep, info: DataTrainingInfo, oid: int, trans
         raise Exception(trans("i18n_data_training.exists_in_db"))
 
     stmt = update(DataTraining).where(and_(DataTraining.id == info.id)).values(
+        embedding=case((DataTraining.question.is_distinct_from(info.question.strip()), None), else_=DataTraining.embedding),
         question=info.question.strip(),
         description=info.description.strip(),
         datasource=info.datasource,
@@ -442,84 +445,46 @@ def enable_training(session: SessionDep, id: int, enabled: bool, trans: Trans, o
     session.commit()
 
 
-# def run_save_embeddings(ids: List[int]):
-#     executor.submit(save_embeddings, ids)
-#
-#
-# def fill_empty_embeddings():
-#     executor.submit(run_fill_empty_embeddings)
-
-
 def run_fill_empty_embeddings(session_maker):
-    try:
-        if not settings.EMBEDDING_ENABLED:
-            return
-
-        session = session_maker()
-        stmt = select(DataTraining.id).where(and_(DataTraining.embedding.is_(None)))
-        results = session.execute(stmt).scalars().all()
-
-        save_embeddings(session_maker, results)
-    except Exception:
-        traceback.print_exc()
-    finally:
-        session_maker.remove()
+    return generate_missing_embeddings(session_maker, DataTraining, DataTraining.question)
 
 
 def save_embeddings(session_maker, ids: List[int]):
-    if not settings.EMBEDDING_ENABLED:
-        return
-
-    if not ids or len(ids) == 0:
+    if not ids:
         return
-    try:
-        session = session_maker()
-        _list = session.query(DataTraining).filter(and_(DataTraining.id.in_(ids))).all()
-
-        _question_list = [item.question for item in _list]
-
-        model = EmbeddingModelCache.get_model()
-
-        results = model.embed_documents(_question_list)
-
-        for index in range(len(results)):
-            item = results[index]
-            stmt = update(DataTraining).where(and_(DataTraining.id == _list[index].id)).values(embedding=item)
-            session.execute(stmt)
-            session.commit()
-
-    except Exception:
-        traceback.print_exc()
-    finally:
-        session_maker.remove()
+    return generate_missing_embeddings(
+        session_maker, DataTraining, DataTraining.question,
+        DataTraining.id.in_(ids),
+    )
 
 
-embedding_sql = f"""
+embedding_sql = """
 SELECT id, datasource, question, similarity
 FROM
 (SELECT id, datasource, question, oid, enabled,
 ( 1 - (embedding <=> :embedding_array) ) AS similarity
 FROM data_training AS child
 ) TEMP
-WHERE similarity > {settings.EMBEDDING_DATA_TRAINING_SIMILARITY} and oid = :oid and datasource = :datasource and enabled = true
+WHERE similarity > :similarity and oid = :oid and datasource = :datasource and enabled = true
 ORDER BY similarity DESC
-LIMIT {settings.EMBEDDING_DATA_TRAINING_TOP_COUNT}
+LIMIT :top_count
 """
-embedding_sql_in_advanced_application = f"""
+embedding_sql_in_advanced_application = """
 SELECT id, advanced_application, question, similarity
 FROM
 (SELECT id, advanced_application, question, oid, enabled,
 ( 1 - (embedding <=> :embedding_array) ) AS similarity
 FROM data_training AS child
 ) TEMP
-WHERE similarity > {settings.EMBEDDING_DATA_TRAINING_SIMILARITY} and oid = :oid and advanced_application = :advanced_application and enabled = true
+WHERE similarity > :similarity and oid = :oid and advanced_application = :advanced_application and enabled = true
 ORDER BY similarity DESC
-LIMIT {settings.EMBEDDING_DATA_TRAINING_TOP_COUNT}
+LIMIT :top_count
 """
 
 
 def select_training_by_question(session: SessionDep, question: str, oid: int, datasource: Optional[int] = None,
-                                advanced_application_id: Optional[int] = None):
+                                advanced_application_id: Optional[int] = None, *, config: EmbeddingConfig | None = None):
+    config = config or load_embedding_config(session)
     if question.strip() == "":
         return []
 
@@ -547,20 +512,20 @@ def select_training_by_question(session: SessionDep, question: str, oid: int, da
     for row in results:
         _list.append(DataTraining(id=row.id, question=row.question))
 
-    if settings.EMBEDDING_ENABLED:
+    if config.enabled:
         with session.begin_nested():
             try:
                 model = EmbeddingModelCache.get_model()
 
-                embedding = model.embed_query(question)
+                embedding = parse_valid_embedding(model.embed_query(question))
 
                 if advanced_application_id is not None:
                     results = session.execute(text(embedding_sql_in_advanced_application),
-                                              {'embedding_array': str(embedding), 'oid': oid,
+                                              {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.data_training_similarity, 'top_count': config.data_training_top_count,
                                                'advanced_application': advanced_application_id})
                 else:
                     results = session.execute(text(embedding_sql),
-                                              {'embedding_array': str(embedding), 'oid': oid, 'datasource': datasource})
+                                              {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.data_training_similarity, 'top_count': config.data_training_top_count, 'datasource': datasource})
 
                 for row in results:
                     _list.append(DataTraining(id=row.id, question=row.question))
@@ -617,12 +582,14 @@ def to_xml_string(_dict: list[dict] | dict, root: str = 'sql-examples') -> str:
 
 
 def get_training_template(session: SessionDep, question: str, oid: Optional[int] = 1, datasource: Optional[int] = None,
-                          advanced_application_id: Optional[int] = None) -> tuple[str, list[dict]]:
+                          advanced_application_id: Optional[int] = None, *,
+                          config: EmbeddingConfig | None = None) -> tuple[str, list[dict]]:
+    config = config or load_embedding_config(session)
     if not oid:
         oid = 1
     if not datasource and not advanced_application_id:
         return '', []
-    _results = select_training_by_question(session, question, oid, datasource, advanced_application_id)
+    _results = select_training_by_question(session, question, oid, datasource, advanced_application_id, config=config)
     if _results and len(_results) > 0:
         data_training = to_xml_string(_results)
         template = get_base_data_training_template().format(data_training=data_training)