diff --git a/README.md b/README.md
index fc7d4b30..55d13ee1 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-
+
基于大模型和 RAG 的智能问数系统
@@ -16,7 +16,7 @@ TBD
SQLBot 的优势包括：
-- **开箱即用**: 只需配置大模型和数据源即可开启问数之旅，通过大模型和 RAG 的结合来实现高质量的 NL2SQL 和 Text2SQL；
+- **开箱即用**: 只需配置大模型和数据源即可开启问数之旅，通过大模型和 RAG 的结合来实现高质量的 text2sql；
- **易于集成**: 支持快速嵌入到第三方业务系统，也支持被 n8n、MaxKB、Dify 、Coze 等 AI 应用开发平台集成调用，让各类应用快速拥有智能问数能力；
- **安全可控**: 提供基于工作空间的资源隔离机制，能够实现细粒度的数据权限控制。
diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py
index e86fc11f..c1f431d7 100644
--- a/backend/apps/chat/models/chat_model.py
+++ b/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='用户提问')
diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py
index 6671291f..3311c9f8 100644
--- a/backend/apps/chat/task/llm.py
+++ b/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:
diff --git a/backend/apps/system/api/assistant.py b/backend/apps/system/api/assistant.py
index 0b1b7cd7..962a1f04 100644
--- a/backend/apps/system/api/assistant.py
+++ b/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()
diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py
index 090c4f1b..8238ab45 100644
--- a/backend/apps/system/crud/assistant.py
+++ b/backend/apps/system/crud/assistant.py
@@ -29,9 +29,8 @@ def get_assistant_user(*, id: int):
email="sqlbot-inner-assistant@sqlbot.com")
-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
diff --git a/backend/apps/template/generate_dynamic/__init__.py b/backend/apps/template/generate_dynamic/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/template/generate_dynamic/generator.py b/backend/apps/template/generate_dynamic/generator.py
new file mode 100644
index 00000000..4c69e0e4
--- /dev/null
+++ b/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']
diff --git a/backend/common/utils/utils.py b/backend/common/utils/utils.py
index 9d35dd4b..a0281df5 100644
--- a/backend/common/utils/utils.py
+++ b/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
+
+
diff --git a/backend/template.yaml b/backend/template.yaml
index 3b9d5c07..76f65a77 100644
--- a/backend/template.yaml
+++ b/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}
diff --git a/frontend/public/assistant.js b/frontend/public/assistant.js
index e2421176..cc0d7b2b 100644
--- a/frontend/public/assistant.js
+++ b/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)
diff --git a/frontend/src/views/embedded/index.vue b/frontend/src/views/embedded/index.vue
index 6c09a965..04530faf 100644
--- a/frontend/src/views/embedded/index.vue
+++ b/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()