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

README.md

@@ -1,4 +1,4 @@
-<p align="center"><img src= "TBD" alt="SQLBot" width="300" /></p>
+<p align="center"><img src="https://resource-fit2cloud-com.oss-cn-hangzhou.aliyuncs.com/sqlbot/sqlbot.png" alt="SQLBot" width="300" /></p>
 <h3 align="center">基于大模型和 RAG 的智能问数系统</h3>
 <p align="center">
   <a href="https://www.gnu.org/licenses/gpl-3.0.html#license-text"><img src="https://img.shields.io/github/license/1Panel-dev/SQLBot?color=%231890FF" alt="License: GPL v3"></a>
@@ -16,7 +16,7 @@ TBD
 
 SQLBot 的优势包括：
 
-- **开箱即用**: 只需配置大模型和数据源即可开启问数之旅，通过大模型和 RAG 的结合来实现高质量的 NL2SQL 和 Text2SQL；
+- **开箱即用**: 只需配置大模型和数据源即可开启问数之旅，通过大模型和 RAG 的结合来实现高质量的 text2sql；
 - **易于集成**: 支持快速嵌入到第三方业务系统，也支持被 n8n、MaxKB、Dify 、Coze 等 AI 应用开发平台集成调用，让各类应用快速拥有智能问数能力；
 - **安全可控**: 提供基于工作空间的资源隔离机制，能够实现细粒度的数据权限控制。
 

backend/apps/chat/models/chat_model.py

@@ -9,6 +9,7 @@ from sqlmodel import SQLModel, Field
 from apps.template.filter.generator import get_permissions_template
 from apps.template.generate_analysis.generator import get_analysis_template
 from apps.template.generate_chart.generator import get_chart_template
+from apps.template.generate_dynamic.generator import get_dynamic_template
 from apps.template.generate_guess_question.generator import get_guess_question_template
 from apps.template.generate_predict.generator import get_predict_template
 from apps.template.generate_sql.generator import get_sql_template
@@ -107,6 +108,7 @@ class AiModelQuestion(BaseModel):
     data: str = ""
     lang: str = "简体中文"
     filter: str = []
+    sub_query: Optional[list[dict]] = None
 
     def sql_sys_question(self):
         return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, lang=self.lang)
@@ -151,7 +153,12 @@ class AiModelQuestion(BaseModel):
 
     def filter_user_question(self):
         return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter)
-
+    
+    def dynamic_sys_question(self): 
+        return get_dynamic_template()['system'].format(lang=self.lang, engine=self.engine)
+    
+    def dynamic_user_question(self): 
+        return get_dynamic_template()['user'].format(sql=self.sql, sub_query=self.sub_query)
 
 class ChatQuestion(AiModelQuestion):
     question: str = Body(description='用户提问')

backend/apps/chat/task/llm.py

@@ -38,7 +38,7 @@ from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, ge
 from apps.system.schemas.system_schema import AssistantOutDsSchema
 from common.core.config import settings
 from common.core.deps import CurrentAssistant, CurrentUser
-from common.utils.utils import SQLBotLogUtil, extract_nested_json
+from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
 
 warnings.filterwarnings("ignore")
 
@@ -71,7 +71,7 @@ class LLMService:
         engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
         session_maker = sessionmaker(bind=engine)
         self.session = session_maker()
-
+        self.session.exec = self.session.exec if hasattr(self.session, "exec") else self.session.execute
         self.current_user = current_user
         self.current_assistant = current_assistant
         # chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
@@ -365,7 +365,7 @@ class LLMService:
         datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
         datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
         if self.current_assistant:
-            _ds_list = get_assistant_ds(llm_service=self)
+            _ds_list = get_assistant_ds(session=self.session, llm_service=self)
         else:
             oid: str = self.current_user.oid
             stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(
@@ -516,6 +516,41 @@ class LLMService:
                                                            [{'type': msg.type, 'content': msg.content} for msg in
                                                             self.sql_message]).decode())
 
+    def generate_with_sub_sql(self, sql, sub_mappings: list):
+        sub_query = json.dumps(sub_mappings, ensure_ascii=False)
+        self.chat_question.sql = sql
+        self.chat_question.sub_query = sub_query
+        msg: List[Union[BaseMessage, dict[str, Any]]] = []
+        msg.append(SystemMessage(content=self.chat_question.dynamic_sys_question()))
+        msg.append(HumanMessage(content=self.chat_question.dynamic_user_question()))
+        full_thinking_text = ''
+        full_dynamic_text = ''
+        res = self.llm.stream(msg)
+        token_usage = {}
+        for chunk in res:
+            SQLBotLogUtil.info(chunk)
+            reasoning_content_chunk = ''
+            if 'reasoning_content' in chunk.additional_kwargs:
+                reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
+            if reasoning_content_chunk is None:
+                reasoning_content_chunk = ''
+            full_thinking_text += reasoning_content_chunk
+            full_dynamic_text += chunk.content
+            get_token_usage(chunk, token_usage)
+
+        SQLBotLogUtil.info(full_dynamic_text)
+        return full_dynamic_text
+    
+    def generate_assistant_dynamic_sql(self, sql, tables: List):
+        ds: AssistantOutDsSchema = self.ds
+        sub_query = []
+        for table in ds.tables:
+            if table.name in tables and table.sql:
+                sub_query.append({"table": table.name, "query": table.sql})
+        if not sub_query:
+            return None
+        return self.generate_with_sub_sql(sql=sql, sub_mappings=sub_query)
+        
     def build_table_filter(self, sql: str, filters: list):
         filter = json.dumps(filters, ensure_ascii=False)
         self.chat_question.sql = sql
@@ -635,27 +670,23 @@ class LLMService:
                                                          full_message=orjson.dumps(
                                                              [{'type': msg.type, 'content': msg.content} for msg in
                                                               self.chart_message]).decode())
-
-    def check_save_sql(self, res: str) -> str:
-
+    def check_sql(self, res: str) -> tuple[any]:
         json_str = extract_nested_json(res)
-        data = orjson.loads(json_str)
-
+        data: dict = orjson.loads(json_str)
         sql = ''
         message = ''
-        error = False
-
         if data['success']:
             sql = data['sql']
         else:
             message = data['message']
-            error = True
-
-        if error:
             raise Exception(message)
+        
         if sql.strip() == '':
             raise Exception("SQL query is empty")
-
+        return sql, data.get('tables')
+        
+    def check_save_sql(self, res: str) -> str:
+        sql, *_ = self.check_sql(res=res)
         save_sql(session=self.session, sql=sql, record_id=self.record.id)
 
         self.chat_question.sql = sql
@@ -716,6 +747,10 @@ class LLMService:
         return save_error_message(session=self.session, record_id=self.record.id, message=message)
 
     def save_sql_data(self, data_obj: Dict[str, Any]):
+        data_result = data_obj.get('data')
+        if data_result:
+            data_result = prepare_for_orjson(data_result)
+            data_obj['data'] = data_result
         return save_sql_exec_data(session=self.session, record_id=self.record.id,
                                   data=orjson.dumps(data_obj).decode())
 
@@ -812,34 +847,27 @@ class LLMService:
 
             # filter sql
             SQLBotLogUtil.info(full_sql_text)
+            use_dynamic_ds: bool = self.current_assistant and self.current_assistant.type == 1
 
             # todo row permission
-            if (not self.current_assistant and is_normal_user(self.current_user)) or (
-                    self.current_assistant and self.current_assistant.type == 1):
-                sql_json_str = extract_nested_json(full_sql_text)
-                data = orjson.loads(sql_json_str)
-
-                sql = ''
-                message = ''
-                error = False
-                if data['success']:
-                    sql = data['sql']
-                else:
-                    message = data['message']
-                    error = True
-                if error:
-                    raise Exception(message)
-                if sql.strip() == '':
-                    raise Exception("SQL query is empty")
+            if (not self.current_assistant and is_normal_user(self.current_user)) or use_dynamic_ds:
+                sql, tables = self.check_sql(res=full_sql_text)
 
                 if self.current_assistant:
-                    sql_result = self.generate_assistant_filter(data.get('sql'), data.get('tables'))
+                    dynamic_sql_result = self.generate_assistant_dynamic_sql(sql, tables)
+                    if dynamic_sql_result:
+                        SQLBotLogUtil.info(dynamic_sql_result)
+                        sql, *_ = self.check_sql(res=dynamic_sql_result)
+                        
+                    sql_result = self.generate_assistant_filter(sql, tables)
                 else:
-                    sql_result = self.generate_filter(data.get('sql'), data.get('tables'))  # maybe no sql and tables
+                    sql_result = self.generate_filter(sql, tables)  # maybe no sql and tables
 
                 if sql_result:
                     SQLBotLogUtil.info(sql_result)
                     sql = self.check_save_sql(res=sql_result)
+                elif dynamic_sql_result:
+                    sql = self.check_save_sql(res=dynamic_sql_result)
                 else:
                     sql = self.check_save_sql(res=full_sql_text)
             else:

backend/apps/system/api/assistant.py

@@ -15,7 +15,9 @@ from common.core.config import settings
 router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
 
 @router.get("/info/{id}") 
-async def info(request: Request, response: Response, session: SessionDep, id: int) -> dict:
+async def info(request: Request, response: Response, session: SessionDep, id: int) -> AssistantModel:
+    if not id:
+        raise Exception('miss assistant id')
     db_model = await get_assistant_info(session=session, assistant_id=id)
     if not db_model:
         raise RuntimeError(f"assistant application not exist")
@@ -23,12 +25,15 @@ async def info(request: Request, response: Response, session: SessionDep, id: in
     response.headers["Access-Control-Allow-Origin"] = db_model.domain
     origin = request.headers.get("origin") or request.headers.get("referer")
     origin = origin.rstrip('/')
-    """ if origin != db_model.domain:
-        raise RuntimeError("invalid domain [{origin}]") """
-    return db_model.model_dump()
+    if origin != db_model.domain:
+        raise RuntimeError("invalid domain [{origin}]")
+    return db_model
 
 @router.get("/validator", response_model=AssistantValidator) 
-async def info(session: SessionDep, id: str, virtual: Optional[int] = Query(None), online: Optional[bool] = Query(default=False)):
+async def validator(session: SessionDep, id: int, virtual: Optional[int] = Query(None)):
+    if not id:
+        raise Exception('miss assistant id')
+    
     db_model = await get_assistant_info(session=session, assistant_id=id)
     if not db_model:
         return AssistantValidator()

backend/apps/system/crud/assistant.py

@@ -29,9 +29,8 @@ def get_assistant_user(*, id: int):
                        email="[email protected]")
 
 
-def get_assistant_ds(llm_service) -> list[dict]:
+def get_assistant_ds(session: Session, llm_service) -> list[dict]:
     assistant: AssistantHeader = llm_service.current_assistant
-    session: Session = llm_service.session
     type = assistant.type
     if type == 0:
         configuration = assistant.configuration

backend/apps/template/generate_dynamic/generator.py

@@ -0,0 +1,6 @@
+from apps.template.template import get_base_template
+
+
+def get_dynamic_template():
+    template = get_base_template()
+    return template['template']['dynamic_sql']

backend/common/utils/utils.py

@@ -1,3 +1,4 @@
+import base64
 import hashlib
 import inspect
 import logging
@@ -215,6 +216,20 @@ class SQLBotLogUtil:
         logger = SQLBotLogUtil._get_logger()
         if logger.isEnabledFor(logging.CRITICAL):
             logger._log(logging.CRITICAL, msg, args, **kwargs)
+            
+def prepare_for_orjson(data):
+    if not data:
+        return data
+    if isinstance(data, bytes):
+        return base64.b64encode(data).decode('utf-8')
+    elif isinstance(data, dict):
+        return {k: prepare_for_orjson(v) for k, v in data.items()}
+    elif isinstance(data, (list, tuple)):
+        return [prepare_for_orjson(item) for item in data]
+    else:
+        return data
+        
+    
 
 
 

backend/template.yaml

@@ -208,3 +208,31 @@ template:
       
       ### 过滤条件:
       {filter}
+  dynamic_sql:
+    system: |
+      ### 请使用语言：{lang} 回答
+      
+      ### 说明：
+      提供给你一句SQL和一组子查询映射表，你需要将给定的SQL查询中的表名替换为对应的子查询。请严格保持原始SQL的结构不变，只替换表引用部分，生成符合{engine}数据库引擎规范的新SQL语句。
+      - 子查询映射表标记为sub_query，格式为[{{"table":"表名","query":"子查询语句"}},...]
+      你必须遵守以下规则:
+      - 生成的SQL必须符合{engine}的规范。
+      - 不要替换原来SQL中的过滤条件。
+      - 完全匹配表名（注意大小写敏感）。
+      - 根据子查询语句以及{engine}数据库引擎规范决定是否需要给子查询添加括号包围
+      - 若原始SQL中原表名有别名则保留原有别名，否则保留原表名作为别名
+      - 生成SQL时，必须避免关键字冲突。
+      - 生成的SQL使用JSON格式返回：
+      {{"success":true,"sql":"生成的SQL语句"}}
+      - 如果不能生成SQL，回答：
+      {{"success":false,"message":"无法生成SQL的原因"}}
+
+      ### 响应, 请直接返回JSON结果:
+      ```json
+
+    user: |
+      ### sql:
+      {sql}
+      
+      ### 子查询映射表:
+      {sub_query}

frontend/public/assistant.js

@@ -1,4 +1,4 @@
-; (function () {
+;(function () {
   window.sqlbot_assistant_handler = window.sqlbot_assistant_handler || {}
   const defaultData = {
     id: '1',
@@ -461,8 +461,10 @@
         let tempData = Object.assign(defaultData, { id, domain_url })
         if (config_json) {
           const config = JSON.parse(config_json)
-          Object.assign(tempData, config)
-          tempData = Object.assign(tempData, config)
+          if (config) {
+            delete config.id
+            tempData = Object.assign(tempData, config)
+          }
         }
         tempData['online'] = online && online.toString().toLowerCase() == 'true'
         initsqlbot_assistant(tempData)

frontend/src/views/embedded/index.vue

@@ -93,6 +93,10 @@ onMounted(() => {
 })
 onBeforeMount(async () => {
   const assistantId = route.query.id
+  if (!assistantId) {
+    ElMessage.error('Miss assistant id, please check assistant url')
+    return
+  }
   const online = route.query.online
   setFormatOnline(online)
   const now = Date.now()