diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 2ff94400..7aef6812 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -42,7 +42,8 @@ from apps.datasource.crud.datasource import get_table_schema, get_tables_sample_data from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user from apps.datasource.embedding.ds_embedding import get_ds_embedding -from apps.datasource.models.datasource import CoreDatasource +from apps.datasource.models.datasource import CoreDatasource, DatasourceConf +from apps.datasource.utils.utils import aes_decrypt from apps.db.db import exec_sql, get_version, check_connection, get_sqlglot_dialect from apps.system.crud.aimodel_manage import get_ai_model_list_by_workspace from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds @@ -64,6 +65,59 @@ executor = ThreadPoolExecutor(max_workers=200) + +# ==================================================================== +# 【改进点1】嵌入态数据源身份翻译:连接指纹映射 +# 嵌入态 self.ds 为外部系统(DataEase)侧 AssistantOutDsSchema,其 id 与 +# SQLBot 侧 core_datasource.id 属不同 id 空间,无法直连。以连接指纹 +# (type, host, port, database, schema) —— 数据源的物理属性 —— 做匹配, +# 不受两侧人工命名差异影响(名称匹配不可靠:异名同库会漏、同名异库会错命中)。 +# 规则:host 归一(localhost/::1/0.0.0.0 → 127.0.0.1,写法漂移实测存在); +# database 必须参与(同实例多库仅靠 ip+port 不可分); +# SSL 不参与匹配(DataEase→SQLBot 传输 VO 无 SSL 字段,两侧无对齐值)。 +# 映射失败返回 None → 检索回退为仅按高级应用锚(安全侧失效:少过滤而非错过滤)。 +# ==================================================================== +def _normalize_host(host) -> str: + """host 归一化:localhost/::1/0.0.0.0 与 127.0.0.1 视为同一地址。""" + h = (host or '').strip().lower() + return '127.0.0.1' if h in ('localhost', '::1', '0.0.0.0') else h + + +def _ds_fingerprint(ds_type, host, port, database, schema='') -> tuple: + """连接指纹:数据源的物理身份五元组(schema 仅 PG 类有值,MySQL 恒为空串)。""" + return ((ds_type or '').strip().lower(), _normalize_host(host), + int(port or 0), (database or '').strip(), (schema or '').strip()) + + +def resolve_embedded_ds_id_by_fingerprint(session: Session, oid: int, out_ds) -> Optional[int]: + """嵌入态:按连接指纹把外部系统侧数据源映射到 SQLBot 侧 core_datasource.id。 + out_ds 为 AssistantOutDsSchema(convert2schema 已解密,host/port/dataBase/db_schema 为明文)。 + SQLBot 侧 configuration 经 aes_decrypt 解密为 DatasourceConf 后取同构字段比对。 + 映射失败返回 None,检索回退为仅按高级应用锚。""" + if out_ds is None: + return None + ds_type = (getattr(out_ds, 'type', '') or '').strip().lower() + if ds_type in ('excel', 'api', 'es'): # 无物理连接语义的类型不参与指纹 + return None + fp = _ds_fingerprint(ds_type, getattr(out_ds, 'host', None), getattr(out_ds, 'port', None), + getattr(out_ds, 'dataBase', None), getattr(out_ds, 'db_schema', None)) + try: + rows = session.execute( + select(CoreDatasource.id, CoreDatasource.type, CoreDatasource.configuration) + .where(CoreDatasource.oid == oid) + ).all() + for ds_id, ds_t, conf_enc in rows: + try: + conf = DatasourceConf(**json.loads(aes_decrypt(conf_enc))) + except Exception: + continue # excel 等类型 configuration 结构不同,跳过 + if _ds_fingerprint(ds_t, conf.host, conf.port, conf.database, conf.dbSchema) == fp: + return ds_id + return None + except Exception: + SQLBotLogUtil.error("resolve_embedded_ds_id_by_fingerprint failed") + return None + dynamic_ds_types = [1, 3] dynamic_subsql_prefix = 'select * from sqlbot_dynamic_temp_table_' @@ -383,15 +437,15 @@ def filter_terminology_template(self, _session: Session, oid: int = None, ds_id: calculate_ds_id = ds_id if self.current_assistant: calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid - if self.current_assistant.type == 1: - calculate_ds_id = None + # 【改进点3】高级应用态不再丢弃 ds_id:保留用于数据源过滤,应用锚仍并行传递(双锚叠加) # 使用原始提取的关键词(未扩展)进行术语匹配,避免同义词重复查询 match_text = self.chat_question.extracted_keywords or self.chat_question.question if self.current_assistant and self.current_assistant.type == 1: self.chat_question.terminologies, term_list = get_terminology_template(_session, match_text, calculate_oid, - None, self.current_assistant.id, config=self.embedding_config) + calculate_ds_id, + self.current_assistant.id, config=self.embedding_config) else: self.chat_question.terminologies, term_list = get_terminology_template(_session, match_text, @@ -445,13 +499,12 @@ def filter_training_template(self, _session: Session, oid: int = None, ds_id: in calculate_ds_id = ds_id if self.current_assistant: calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid - if self.current_assistant.type == 1: - calculate_ds_id = None + # 【改进点4】高级应用态不再丢弃 ds_id:保留用于数据源过滤,应用锚仍并行传递(双锚叠加) if self.current_assistant and self.current_assistant.type == 1: self.chat_question.data_training, example_list = get_training_template(_session, self.chat_question.question, calculate_oid, - None, self.current_assistant.id, config=self.embedding_config) + calculate_ds_id, self.current_assistant.id, config=self.embedding_config) else: self.chat_question.data_training, example_list = get_training_template(_session, self.chat_question.question, @@ -918,6 +971,11 @@ def select_datasource(self, _session: Session): if self.ds: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None + # 【改进点2】嵌入态:外部系统(DataEase)侧数据源 id 与 SQLBot 侧不同空间, + # 按连接指纹 (type, host, port, database, schema) 映射出 SQLBot 侧 ds_id, + # 供术语/示例检索做数据源过滤(映射失败回退 None) + if ds_id is None: + ds_id = resolve_embedded_ds_id_by_fingerprint(_session, oid, self.ds) # 提取关键词(前置步骤,结果复用于术语匹配和表匹配) if self.embedding_config.keyword_enabled: @@ -1400,6 +1458,11 @@ def run_task(self, in_chat: bool = True, stream: bool = True, if self.ds: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None + # 【改进点2】嵌入态:外部系统(DataEase)侧数据源 id 与 SQLBot 侧不同空间, + # 按连接指纹 (type, host, port, database, schema) 映射出 SQLBot 侧 ds_id, + # 供术语/示例检索做数据源过滤(映射失败回退 None) + if ds_id is None: + ds_id = resolve_embedded_ds_id_by_fingerprint(_session, oid, self.ds) # 提取关键词(前置步骤,结果复用于术语匹配和表匹配) if self.embedding_config.keyword_enabled: diff --git a/backend/apps/data_training/curd/data_training.py b/backend/apps/data_training/curd/data_training.py index f66c0575..e4d1f4b6 100644 --- a/backend/apps/data_training/curd/data_training.py +++ b/backend/apps/data_training/curd/data_training.py @@ -481,6 +481,20 @@ def save_embeddings(session_maker, ids: List[int]): LIMIT :top_count """ +# 【改进点7】高级应用/嵌入态专用:应用锚 + 数据源锚叠加的向量检索(示例侧单值绑定,精确等值) +embedding_sql_in_advanced_application_and_ds = """ +SELECT id, advanced_application, question, similarity +FROM +(SELECT id, advanced_application, datasource, question, oid, enabled, +( 1 - (embedding <=> :embedding_array) ) AS similarity +FROM data_training AS child +) TEMP +WHERE similarity > :similarity and oid = :oid and advanced_application = :advanced_application and enabled = true +and datasource = :datasource +ORDER BY similarity DESC +LIMIT :top_count +""" + def select_training_by_question(session: SessionDep, question: str, oid: int, datasource: Optional[int] = None, advanced_application_id: Optional[int] = None, *, config: EmbeddingConfig | None = None): @@ -504,6 +518,9 @@ def select_training_by_question(session: SessionDep, question: str, oid: int, da ) if advanced_application_id is not None: stmt = stmt.where(and_(DataTraining.advanced_application == advanced_application_id)) + # 【改进点8】高级应用/嵌入态叠加数据源过滤(示例侧单值绑定,精确等值) + if datasource is not None: + stmt = stmt.where(DataTraining.datasource == datasource) else: stmt = stmt.where(and_(DataTraining.datasource == datasource)) @@ -519,7 +536,13 @@ def select_training_by_question(session: SessionDep, question: str, oid: int, da embedding = parse_valid_embedding(model.embed_query(question)) - if advanced_application_id is not None: + if advanced_application_id is not None and datasource is not None: + # 【改进点8】双锚齐备:应用锚 + 数据源过滤叠加 + results = session.execute(text(embedding_sql_in_advanced_application_and_ds), + {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.data_training_similarity, 'top_count': config.data_training_top_count, + 'advanced_application': advanced_application_id, + 'datasource': datasource}) + elif advanced_application_id is not None: results = session.execute(text(embedding_sql_in_advanced_application), {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.data_training_similarity, 'top_count': config.data_training_top_count, 'advanced_application': advanced_application_id}) diff --git a/backend/apps/terminology/curd/terminology.py b/backend/apps/terminology/curd/terminology.py index 0cb50548..3de101d5 100644 --- a/backend/apps/terminology/curd/terminology.py +++ b/backend/apps/terminology/curd/terminology.py @@ -801,6 +801,26 @@ def save_embeddings(session_maker, ids: List[int]): LIMIT :top_count """ +# 【改进点5】高级应用/嵌入态专用:应用锚 + 数据源锚叠加的向量检索 +# 语义:命中应用锚,且(所有数据源生效 或 绑定数据源数组包含当前所选 ds) +embedding_sql_with_advanced_application_and_ds = """ +SELECT id, pid, word, similarity +FROM +(SELECT id, pid, word, oid, specific_ds, advanced_application, datasource_ids, enabled, +( 1 - (embedding <=> :embedding_array) ) AS similarity +FROM terminology AS child +) TEMP +WHERE similarity > :similarity AND oid = :oid AND enabled = true +AND advanced_application = :advanced_application_id +AND ( + (specific_ds = false OR specific_ds IS NULL) + OR + (specific_ds = true AND datasource_ids IS NOT NULL AND datasource_ids @> jsonb_build_array(:datasource)) +) +ORDER BY similarity DESC +LIMIT :top_count +""" + def select_terminology_by_word(session: SessionDep, word: str, oid: int, datasource: int = None, advanced_application_id: Optional[int] = None, *, config: EmbeddingConfig | None = None): @@ -823,6 +843,18 @@ def select_terminology_by_word(session: SessionDep, word: str, oid: int, datasou if advanced_application_id is not None: stmt = stmt.where(Terminology.advanced_application == advanced_application_id) + # 【改进点6】高级应用/嵌入态叠加数据源过滤:放行"所有数据源",过滤"绑定不含当前 ds"的术语 + if datasource is not None: + stmt = stmt.where( + or_( + or_(Terminology.specific_ds == False, Terminology.specific_ds.is_(None)), + and_( + Terminology.specific_ds == True, + Terminology.datasource_ids.isnot(None), + text("datasource_ids @> jsonb_build_array(:datasource)") + ) + ) + ) elif datasource is not None: stmt = stmt.where( or_( @@ -854,7 +886,13 @@ def select_terminology_by_word(session: SessionDep, word: str, oid: int, datasou embedding = parse_valid_embedding(model.embed_query(word)) - if advanced_application_id is not None: + if advanced_application_id is not None and datasource is not None: + # 【改进点6】双锚齐备:应用锚 + 数据源过滤叠加 + results = session.execute(text(embedding_sql_with_advanced_application_and_ds), + {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.terminology_similarity, 'top_count': config.terminology_top_count, + 'advanced_application_id': advanced_application_id, + 'datasource': datasource}).fetchall() + elif advanced_application_id is not None: results = session.execute(text(embedding_sql_with_advanced_application), {'embedding_array': str(embedding), 'oid': oid, 'similarity': config.terminology_similarity, 'top_count': config.terminology_top_count, 'advanced_application_id': advanced_application_id}).fetchall()