Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

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"]
+ENTRYPOINT ["sh", "start.sh"]

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 = {}

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)

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 @@ def chart_user_question(self):
         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)

backend/apps/chat/task/llm.py

@@ -31,6 +31,7 @@
 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 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
         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 @@ def generate_analysis(self):
         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 @@ def run_task_cache(self, in_chat: bool = True):
 
     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'

backend/apps/template/generate_chart/generator.py

@@ -4,3 +4,7 @@
 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']

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('<?xml'):
+        end_index = pretty_xml.find('>') + 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 ''

backend/common/core/config.py

@@ -88,9 +88,9 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
 
     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