From e5b9dc9e4c7c262fe9fa52bfb4be4dba0fba3324 Mon Sep 17 00:00:00 2001 From: ulleo Date: Thu, 25 Sep 2025 17:30:53 +0800 Subject: [PATCH] feat: terminology settings add datasource #127 --- backend/apps/chat/task/llm.py | 8 ++-- backend/apps/terminology/curd/terminology.py | 42 ++++++++++++++++---- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 89fb52a1..f36bca31 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -241,8 +241,9 @@ class LLMService: self.chat_question.data = orjson.dumps(data.get('data')).decode() analysis_msg: List[Union[BaseMessage, dict[str, Any]]] = [] + ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question, - self.current_user.oid) + self.current_user.oid, ds_id) analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question())) analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question())) @@ -504,7 +505,8 @@ class LLMService: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None - self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question, oid) + self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question, oid, + ds_id) self.chat_question.data_training = get_training_template(self.session, self.chat_question.question, ds_id, oid) @@ -897,7 +899,7 @@ class LLMService: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question, - oid) + oid, ds_id) self.chat_question.data_training = get_training_template(self.session, self.chat_question.question, ds_id, oid) diff --git a/backend/apps/terminology/curd/terminology.py b/backend/apps/terminology/curd/terminology.py index 7b8d0a04..f10f3481 100644 --- a/backend/apps/terminology/curd/terminology.py +++ b/backend/apps/terminology/curd/terminology.py @@ -1,7 +1,7 @@ import datetime import logging import traceback -from typing import List, Optional +from typing import List, Optional, Any from xml.dom.minidom import parseString import dicttoxml @@ -367,17 +367,22 @@ def save_embeddings(session: Session, ids: List[int]): embedding_sql = f""" SELECT id, pid, word, similarity FROM -(SELECT id, pid, word, oid, +(SELECT id, pid, word, oid, specific_ds, datasource_ids, ( 1 - (embedding <=> :embedding_array) ) AS similarity FROM terminology AS child ) TEMP -WHERE similarity > {settings.EMBEDDING_TERMINOLOGY_SIMILARITY} and oid = :oid +WHERE similarity > {settings.EMBEDDING_TERMINOLOGY_SIMILARITY} AND oid = :oid +AND ( + (:datasource IS NULL AND (specific_ds = false OR specific_ds IS NULL)) + OR + (:datasource IS NOT NULL 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 {settings.EMBEDDING_TERMINOLOGY_TOP_COUNT} """ -def select_terminology_by_word(session: SessionDep, word: str, oid: int): +def select_terminology_by_word(session: SessionDep, word: str, oid: int, datasource: int = None): if word.strip() == "": return [] @@ -394,7 +399,26 @@ def select_terminology_by_word(session: SessionDep, word: str, oid: int): ) ) - results = session.execute(stmt, {'sentence': word}).fetchall() + 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)") + ) + ) + ) + else: + stmt = stmt.where(or_(Terminology.specific_ds == False, Terminology.specific_ds.is_(None))) + + # 执行查询 + params: dict[str, Any] = {'sentence': word} + if datasource is not None: + params['datasource'] = datasource + + results = session.execute(stmt, params).fetchall() for row in results: _list.append(Terminology(id=row.id, word=row.word, pid=row.pid)) @@ -405,7 +429,8 @@ def select_terminology_by_word(session: SessionDep, word: str, oid: int): embedding = model.embed_query(word) - results = session.execute(text(embedding_sql), {'embedding_array': str(embedding), 'oid': oid}) + results = session.execute(text(embedding_sql), {'embedding_array': str(embedding), 'oid': oid, + 'datasource': datasource}).fetchall() for row in results: _list.append(Terminology(id=row.id, word=row.word, pid=row.pid)) @@ -481,10 +506,11 @@ def to_xml_string(_dict: list[dict] | dict, root: str = 'terminologies') -> str: return pretty_xml -def get_terminology_template(session: SessionDep, question: str, oid: Optional[int] = 1) -> str: +def get_terminology_template(session: SessionDep, question: str, oid: Optional[int] = 1, + datasource: Optional[int] = None) -> str: if not oid: oid = 1 - _results = select_terminology_by_word(session, question, oid) + _results = select_terminology_by_word(session, question, oid, datasource) if _results and len(_results) > 0: terminology = to_xml_string(_results) template = get_base_terminology_template().format(terminologies=terminology)