diff --git a/Dockerfile b/Dockerfile index cb6724de..3a1774ef 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,5 @@ # Build sqlbot +FROM ghcr.io/1panel-dev/maxkb-vector-model:v1.0.1 AS vector-model FROM registry.cn-qingdao.aliyuncs.com/dataease/sqlbot-base:latest AS sqlbot-builder # Set build environment variables @@ -58,15 +59,16 @@ COPY start.sh /opt/sqlbot/app/start.sh COPY g2-ssr/*.ttf /usr/share/fonts/truetype/liberation/ COPY --from=sqlbot-builder ${SQLBOT_HOME} ${SQLBOT_HOME} COPY --from=ssr-builder /app /opt/sqlbot/g2-ssr +COPY --from=vector-model /opt/maxkb/app/model /opt/sqlbot/models WORKDIR ${SQLBOT_HOME}/app RUN mkdir -p /opt/sqlbot/images /opt/sqlbot/g2-ssr -EXPOSE 3000 8000 +EXPOSE 3000 8000 8001 # Add health check HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \ CMD curl -f http://localhost:8000 || exit 1 -ENTRYPOINT ["sh", "start.sh"] \ No newline at end of file +ENTRYPOINT ["sh", "start.sh"] diff --git a/backend/apps/ai_model/embedding.py b/backend/apps/ai_model/embedding.py index dea3ffbe..c458e3ff 100644 --- a/backend/apps/ai_model/embedding.py +++ b/backend/apps/ai_model/embedding.py @@ -1,3 +1,4 @@ +import os.path import threading from typing import Optional @@ -14,7 +15,9 @@ class EmbeddingModelInfo(BaseModel): device: str = 'cpu' -local_embedding_model = EmbeddingModelInfo(folder=settings.LOCAL_MODEL_PATH, name=settings.DEFAULT_EMBEDDING_MODEL) +local_embedding_model = EmbeddingModelInfo(folder=settings.LOCAL_MODEL_PATH, + name=os.path.join(settings.LOCAL_MODEL_PATH, 'embedding', + "shibing624_text2vec-base-chinese")) _lock = threading.Lock() locks = {} diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 3cc59380..5cc0e125 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -186,7 +186,7 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch detail=f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it" ) - request_question = ChatQuestion(chat_id=record.chat_id, question='') + request_question = ChatQuestion(chat_id=record.chat_id, question=record.question) try: llm_service = LLMService(current_user, request_question, current_assistant) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index d1cc6670..a8f3e9c8 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -170,10 +170,11 @@ class AiModelQuestion(BaseModel): lang: str = "简体中文" filter: str = [] sub_query: Optional[list[dict]] = None + terminologies: str = "" def sql_sys_question(self): return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, - lang=self.lang) + lang=self.lang, terminologies=self.terminologies) def sql_user_question(self): return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question, @@ -186,7 +187,7 @@ class AiModelQuestion(BaseModel): return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule) def analysis_sys_question(self): - return get_analysis_template()['system'].format(lang=self.lang) + return get_analysis_template()['system'].format(lang=self.lang, terminologies=self.terminologies) def analysis_user_question(self): return get_analysis_template()['user'].format(fields=self.fields, data=self.data) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 55279269..00c9e618 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -31,6 +31,7 @@ from apps.datasource.models.datasource import CoreDatasource from apps.db.db import exec_sql, get_version from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds from apps.system.schemas.system_schema import AssistantOutDsSchema +from apps.terminology.curd.terminology import get_terminology_template from common.core.config import settings from common.core.deps import CurrentAssistant, CurrentUser from common.error import SingleMessageError @@ -124,8 +125,6 @@ class LLMService: llm_instance = LLMFactory.create_llm(self.config) self.llm = llm_instance.llm - self.init_messages() - def is_running(self, timeout=0.5): try: r = concurrent.futures.wait([self.future], timeout) @@ -210,6 +209,9 @@ class LLMService: data = get_chat_chart_data(self.session, self.record.id) self.chat_question.data = orjson.dumps(data.get('data')).decode() analysis_msg: List[Union[BaseMessage, dict[str, Any]]] = [] + + self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question) + analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question())) analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question())) @@ -860,6 +862,9 @@ class LLMService: def run_task(self, in_chat: bool = True): try: + self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question) + self.init_messages() + # return id if in_chat: yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n' diff --git a/backend/apps/template/generate_chart/generator.py b/backend/apps/template/generate_chart/generator.py index 0962cba7..2c886277 100644 --- a/backend/apps/template/generate_chart/generator.py +++ b/backend/apps/template/generate_chart/generator.py @@ -4,3 +4,7 @@ from apps.template.template import get_base_template def get_chart_template(): template = get_base_template() return template['template']['chart'] + +def get_base_terminology_template(): + template = get_base_template() + return template['template']['terminology'] diff --git a/backend/apps/terminology/curd/terminology.py b/backend/apps/terminology/curd/terminology.py index 835ad3d8..97dc564b 100644 --- a/backend/apps/terminology/curd/terminology.py +++ b/backend/apps/terminology/curd/terminology.py @@ -1,12 +1,24 @@ import datetime +import logging +import traceback +from concurrent.futures import ThreadPoolExecutor from typing import List, Optional +from xml.dom.minidom import parseString -from sqlalchemy import and_, or_, select, func, delete, update +import dicttoxml +from sqlalchemy import and_, or_, select, func, delete, update, union +from sqlalchemy import create_engine, text from sqlalchemy.orm import aliased +from sqlalchemy.orm import sessionmaker +from apps.ai_model.embedding import EmbeddingModelCache +from apps.template.generate_chart.generator import get_base_terminology_template from apps.terminology.models.terminology_model import Terminology, TerminologyInfo +from common.core.config import settings from common.core.deps import SessionDep +executor = ThreadPoolExecutor(max_workers=200) + def page_terminology(session: SessionDep, current_page: int = 1, page_size: int = 10, name: Optional[str] = None): _list: List[TerminologyInfo] = [] @@ -24,7 +36,7 @@ def page_terminology(session: SessionDep, current_page: int = 1, page_size: int # 步骤1:先找到所有匹配的节点ID(无论是父节点还是子节点) matched_ids_subquery = ( select(Terminology.id) - .where(Terminology.word.like(keyword_pattern)) # LIKE查询条件 + .where(Terminology.word.ilike(keyword_pattern)) # LIKE查询条件 .subquery() ) @@ -82,7 +94,6 @@ def page_terminology(session: SessionDep, current_page: int = 1, page_size: int .where(Terminology.id.in_(paginated_parent_ids)) .order_by(Terminology.create_time.desc()) ) - print(str(stmt)) else: parent_ids_subquery = ( select(Terminology.id) @@ -113,7 +124,6 @@ def page_terminology(session: SessionDep, current_page: int = 1, page_size: int .group_by(Terminology.id, Terminology.word) .order_by(Terminology.create_time.desc()) ) - print(str(stmt)) result = session.execute(stmt) @@ -145,13 +155,16 @@ def create_terminology(session: SessionDep, info: TerminologyInfo): _list: List[Terminology] = [] if info.other_words: for other_word in info.other_words: + if other_word.strip() == "": + continue _list.append( Terminology(pid=result.id, word=other_word, create_time=create_time)) session.bulk_save_objects(_list) session.flush() session.commit() - # todo embedding + # embedding + run_save_embeddings([result.id]) return result.id @@ -172,13 +185,16 @@ def update_terminology(session: SessionDep, info: TerminologyInfo): _list: List[Terminology] = [] if info.other_words: for other_word in info.other_words: + if other_word.strip() == "": + continue _list.append( Terminology(pid=info.id, word=other_word, create_time=create_time)) session.bulk_save_objects(_list) session.flush() session.commit() - # todo embedding + # embedding + run_save_embeddings([info.id]) return info.id @@ -187,3 +203,172 @@ def delete_terminology(session: SessionDep, ids: list[int]): stmt = delete(Terminology).where(or_(Terminology.id.in_(ids), Terminology.pid.in_(ids))) session.execute(stmt) 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(): + if not settings.EMBEDDING_ENABLED: + return + engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI)) + session_maker = sessionmaker(bind=engine) + session = session_maker() + stmt1 = select(Terminology.id).where(and_(Terminology.embedding.is_(None), Terminology.pid.is_(None))) + stmt2 = select(Terminology.pid).where(and_(Terminology.embedding.is_(None), Terminology.pid.isnot(None))).distinct() + combined_stmt = union(stmt1, stmt2) + results = session.execute(combined_stmt).scalars().all() + save_embeddings(results) + + +def save_embeddings(ids: List[int]): + if not settings.EMBEDDING_ENABLED: + return + + if not ids or len(ids) == 0: + return + try: + engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI)) + session_maker = sessionmaker(bind=engine) + session = session_maker() + + _list = session.query(Terminology).filter(or_(Terminology.id.in_(ids), Terminology.pid.in_(ids))).all() + + _words_list = [item.word for item in _list] + + model = EmbeddingModelCache.get_model() + + results = model.embed_documents(_words_list) + + for index in range(len(results)): + item = results[index] + stmt = update(Terminology).where(and_(Terminology.id == _list[index].id)).values(embedding=item) + session.execute(stmt) + session.commit() + + except Exception: + traceback.print_exc() + + +embedding_sql = f""" +SELECT id, pid, word, description, similarity +FROM +(SELECT id, pid, word, +COALESCE( + description, + (SELECT description FROM terminology AS parent WHERE parent.id = child.pid) + ) AS description, +( 1 - (embedding <=> :embedding_array) ) AS similarity +FROM terminology AS child +) TEMP +WHERE similarity > {settings.EMBEDDING_SIMILARITY} +ORDER BY similarity DESC +LIMIT {settings.EMBEDDING_TOP_COUNT} +""" + + +def select_terminology_by_word(session: SessionDep, word: str): + if word.strip() == "": + return [] + + _list: List[Terminology] = [] + + stmt = ( + select( + Terminology.id, + Terminology.pid, + Terminology.word, + func.coalesce( + Terminology.description, + select(Terminology.description) + .where(and_(Terminology.id == Terminology.pid)) + .scalar_subquery() + ).label('description') + ) + .where( + text(":sentence ILIKE '%' || word || '%'") + ) + ) + + results = session.execute(stmt, {'sentence': word}).fetchall() + + for row in results: + _list.append(Terminology(id=row.id, word=row.word, pid=row.pid, description=row.description)) + + if settings.EMBEDDING_ENABLED: + try: + model = EmbeddingModelCache.get_model() + + embedding = model.embed_query(word) + + print(embedding_sql) + results = session.execute(text(embedding_sql), {'embedding_array': str(embedding)}) + + for row in results: + _list.append(Terminology(id=row.id, word=row.word, pid=row.pid, description=row.description)) + + except Exception: + traceback.print_exc() + + _map: dict = {} + _ids: set[int] = set() + for row in _list: + if row.id in _ids: + continue + _ids.add(row.id) + if row.pid: + pid = str(row.pid) + else: + pid = str(row.id) + if _map.get(pid) is None: + _map[pid] = {'words': [], 'description': row.description} + _map[pid]['words'].append(row.word) + + _results: list[dict] = [] + for key in _map.keys(): + _results.append(_map.get(key)) + + return _results + + +def get_example(): + _obj = { + 'terminologies': [ + {'words': ['GDP', '国内生产总值'], + 'description': '指在一个季度或一年,一个国家或地区的经济中所生产出的全部最终产品和劳务的价值。'}, + ] + } + return to_xml_string(_obj, 'example') + + +def to_xml_string(_dict: list[dict] | dict, root: str = 'terminologies') -> str: + item_name_func = lambda x: 'terminology' if x == 'terminologies' else 'word' if x == 'words' else 'item' + dicttoxml.LOG.setLevel(logging.ERROR) + xml = dicttoxml.dicttoxml(_dict, + custom_root=root, + item_func=item_name_func, + xml_declaration=False, + encoding='utf-8', + attr_type=False).decode('utf-8') + pretty_xml = parseString(xml).toprettyxml() + + if pretty_xml.startswith('') + 1 + pretty_xml = pretty_xml[end_index:].lstrip() + + return pretty_xml + + +def get_terminology_template(session: SessionDep, question: str) -> str: + _results = select_terminology_by_word(session, question) + if _results and len(_results) > 0: + terminology = to_xml_string(_results) + template = get_base_terminology_template().format(terminologies=terminology) + return template + else: + return '' diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 793bd0be..920ee73f 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -88,9 +88,9 @@ class Settings(BaseSettings): LOCAL_MODEL_PATH: str = '/opt/sqlbot/models' DEFAULT_EMBEDDING_MODEL: str = 'shibing624/text2vec-base-chinese' - - EMBEDDING_SIMILARITY: float = 0.6 - EMBEDDING_TOP_COUNT: int = 3 + EMBEDDING_ENABLED: bool = True + EMBEDDING_SIMILARITY: float = 0.4 + EMBEDDING_TOP_COUNT: int = 5 settings = Settings() # type: ignore diff --git a/backend/main.py b/backend/main.py index 7a639e5a..8ca5a69e 100644 --- a/backend/main.py +++ b/backend/main.py @@ -12,6 +12,7 @@ from alembic import command from apps.api import api_router from apps.system.crud.assistant import init_dynamic_cors from apps.system.middleware.auth import TokenMiddleware +from apps.terminology.curd.terminology import fill_empty_embeddings from common.core.config import settings from common.core.response_middleware import ResponseMiddleware, exception_handler from common.core.sqlbot_cache import init_sqlbot_cache @@ -23,11 +24,16 @@ def run_migrations(): command.upgrade(alembic_cfg, "head") +def init_embedding_data(): + fill_empty_embeddings() + + @asynccontextmanager async def lifespan(app: FastAPI): run_migrations() init_sqlbot_cache() init_dynamic_cors(app) + init_embedding_data() SQLBotLogUtil.info("✅ SQLBot 初始化完成") await sqlbot_xpack.core.clean_xpack_cache() yield diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 5f6273e7..401e7b96 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -46,6 +46,7 @@ dependencies = [ "python-calamine>=0.4.0", "xlrd>=2.0.2", "clickhouse-sqlalchemy>=0.3.2", + "dicttoxml>=1.7.16", "dmpython>=2.5.22; platform_system != 'Darwin'", ] [[tool.uv.index]] diff --git a/backend/template.yaml b/backend/template.yaml index 40a1a134..98e1bf21 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -1,4 +1,16 @@ template: + terminology: | + ### 提问或数据可能包含有专业术语,下面会提供给你可以参考的专业术语及其术语描述。 + + #### 专业术语格式描述: + 有一组专业术语 + 每一个就是专业术语 + 其中同一个内的多个代表术语的多种叫法,也就是术语与它的同义词 + 即该专业术语对应的专业术语的描述 + + #### 专业术语: + {terminologies} + sql: system: | ### 请使用语言:{lang} 回答,若有深度思考过程,则思考过程也需要使用 {lang} 输出 @@ -60,6 +72,8 @@ template: ### 提供表结构如下: {schema} + {terminologies} + ### 响应, 请直接返回JSON结果: ```json @@ -151,6 +165,8 @@ template: ### 说明: 你是一个数据分析师,你的任务是根据给定的数据分析数据,并给出你的分析结果。 + + {terminologies} user: | ### 字段(字段别名): {fields}