import asyncio
import concurrent
import json
import os
import re
import traceback
import urllib.parse
import warnings
from concurrent.futures import ThreadPoolExecutor, Future
from datetime import datetime
from typing import Any, List, Optional, Union, Dict, Iterator
import orjson
import pandas as pd
import requests
import sqlglot
import sqlparse
from langchain.chat_models.base import BaseChatModel
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, BaseMessageChunk
from sqlalchemy import and_, select
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlbot_xpack.config.model import SysArgModel
from sqlbot_xpack.custom_prompt.curd.custom_prompt import find_custom_prompts
from sqlbot_xpack.custom_prompt.models.custom_prompt_model import CustomPromptTypeEnum
from sqlbot_xpack.license.license_manage import SQLBotLicenseUtil
from sqlglot import exp
from sqlmodel import Session
from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config
from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
save_error_message, save_sql_exec_data, save_chart_answer, save_chart, \
finish_record, save_analysis_answer, save_predict_answer, save_predict_data, \
save_select_datasource_answer, save_recommend_question_answer, save_extracted_keywords, \
get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \
get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate, get_chat_predict_data, \
get_chat_chart_config, trigger_log_error
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \
ChatFinishStep, AxisObj, SystemPromptMessage, HumanPromptMessage, AIPromptMessage
from apps.data_training.curd.data_training import get_training_template
from apps.datasource.crud.datasource import get_table_schema, get_tables_sample_data
from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user
from apps.datasource.embedding.ds_embedding import get_ds_embedding
from apps.datasource.models.datasource import CoreDatasource
from apps.db.db import exec_sql, get_version, check_connection, get_sqlglot_dialect
from apps.system.crud.aimodel_manage import get_ai_model_list_by_workspace
from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds
from apps.system.crud.parameter_manage import get_groups
from apps.system.crud.user import user_ws_list
from apps.system.schemas.system_schema import AssistantOutDsSchema
from apps.terminology.curd.terminology import expand_with_terminology
from apps.terminology.curd.terminology import get_terminology_template
from common.core.config import settings
from common.core.db import engine
from common.core.deps import CurrentAssistant, CurrentUser
from common.error import SingleMessageError, SQLBotDBError, ParseSQLResultError, SQLBotDBConnectionError
from common.utils.data_format import DataFormat
from common.utils.locale import I18n, I18nHelper
from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
warnings.filterwarnings("ignore")
executor = ThreadPoolExecutor(max_workers=200)
dynamic_ds_types = [1, 3]
dynamic_subsql_prefix = 'select * from sqlbot_dynamic_temp_table_'
session_maker = scoped_session(sessionmaker(bind=engine, class_=Session))
i18n = I18n()
def extract_tables_from_sql(sql: str, ds_type: str = None) -> set:
"""从 SQL 中提取真实表名（排除 CTE 别名）"""
tables = set()
dialect = get_sqlglot_dialect(ds_type)
try:
statements = sqlglot.parse(sql, dialect=dialect)
for stmt in statements:
if stmt:
# 收集 CTE 别名，排除嵌套 CTE
cte_names = set()
for cte in stmt.find_all(exp.CTE):
if cte.alias:
cte_names.add(cte.alias)
for table in stmt.find_all(exp.Table):
if table.name and table.name not in cte_names:
tables.add(table.name)
except Exception:
pass
return tables
class LLMService:
ds: CoreDatasource
chat_question: ChatQuestion
oid: int
record: ChatRecord
config: LLMConfig
llm: BaseChatModel
sql_message: List[Union[BaseMessage, dict[str, Any]]]
chart_message: List[Union[BaseMessage, dict[str, Any]]]
# session: Session = db_session
current_user: CurrentUser
current_assistant: Optional[CurrentAssistant] = None
out_ds_instance: Optional[AssistantOutDs] = None
change_title: bool = False
generate_sql_logs: List[ChatLog]
generate_chart_logs: List[ChatLog]
current_logs: dict[OperationEnum, ChatLog]
chunk_list: List[str]
future: Future
trans: I18nHelper = None
last_execute_sql_error: str = None
articles_number: int = 4
enable_sql_row_limit: bool = settings.GENERATE_SQL_QUERY_LIMIT_ENABLED
base_message_round_count_limit: int = settings.GENERATE_SQL_QUERY_HISTORY_ROUND_COUNT
def __init__(self, session: Session, current_user: CurrentUser, chat_question: ChatQuestion,
current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False,
embedding: bool = False, config: LLMConfig = None):
self.sql_message = []
self.chart_message = []
self.generate_sql_logs = []
self.generate_chart_logs = []
self.current_logs = {}
self.chunk_list = []
self.current_user = current_user
self.current_assistant = current_assistant
self.table_name_list = []
chat_question.lang = get_lang_name(current_user.language)
self.trans = i18n(lang=current_user.language)
chat_id = chat_question.chat_id
chat: Chat | None = session.get(Chat, chat_id)
if not chat:
raise SingleMessageError(f"Chat with id {chat_id} not found")
self.oid = chat.oid
if self.oid and not current_assistant:
w_list = user_ws_list(session, self.current_user.id)
oid_list = [item.id for item in w_list]
if int(self.oid) not in oid_list:
raise SingleMessageError("Current user cannot not access this chat")
if self.oid and current_assistant:
if self.oid != self.current_user.oid:
raise SingleMessageError("Current assistant user cannot not access this chat")
self.current_user.oid = chat.oid
ds: CoreDatasource | AssistantOutDsSchema | None = None
if not chat.datasource and chat_question.datasource_id:
_ds = session.get(CoreDatasource, chat_question.datasource_id)
if _ds:
if _ds.oid != self.oid:
raise SingleMessageError(
f"Datasource with id {chat_question.datasource_id} does not belong to current workspace")
chat.datasource = _ds.id
chat.engine_type = _ds.type_name
# save chat
session.add(chat)
session.flush()
session.refresh(chat)
session.commit()
if chat.datasource:
# Get available datasource
if current_assistant and current_assistant.type in dynamic_ds_types:
self.out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
ds = self.out_ds_instance.get_ds(chat.datasource)
if not ds:
raise SingleMessageError("No available datasource configuration found")
chat_question.engine = ds.type + get_version(ds)
else:
ds = session.get(CoreDatasource, chat.datasource)
if not ds:
raise SingleMessageError("No available datasource configuration found")
chat_question.engine = (ds.type_name if ds.type != 'excel' else 'PostgreSQL') + get_version(ds)
self.generate_sql_logs = list_generate_sql_logs(session=session, chart_id=chat_id)
self.generate_chart_logs = list_generate_chart_logs(session=session, chart_id=chat_id)
self.change_title = not get_chat_brief_generate(session=session, chat_id=chat_id)
self.ds = (
ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None
self.chat_question = chat_question
self.config = config
if no_reasoning:
# only work while using qwen
if self.config.additional_params:
if self.config.additional_params.get('extra_body'):
if self.config.additional_params.get('extra_body').get('enable_thinking'):
del self.config.additional_params['extra_body']['enable_thinking']
self.chat_question.ai_modal_id = self.config.model_id
self.chat_question.ai_modal_name = self.config.model_name
# Create LLM instance through factory
llm_instance = LLMFactory.create_llm(self.config)
self.llm = llm_instance.llm
# get last_execute_sql_error
last_execute_sql_error = get_last_execute_sql_error(session, self.chat_question.chat_id)
if last_execute_sql_error:
self.chat_question.error_msg = f'''
{last_execute_sql_error}
'''
else:
self.chat_question.error_msg = ''
@classmethod
async def create(cls, *args, **kwargs):
specialized_model_id = None
_ai_model_list = []
if args[3]:
if args[1]:
ws_id = args[1].oid
_ai_model_list = get_ai_model_list_by_workspace(args[0], ws_id)
if args[3].enable_custom_model:
if args[3].custom_model:
if any(str(model.id) == str(args[3].custom_model) for model in _ai_model_list):
specialized_model_id = args[3].custom_model
print("use custom model: id[" + specialized_model_id + "]")
# mcp call, custom model id
if args[2] and args[2].custom_model:
if args[1]:
ws_id = args[1].oid
_ai_model_list = get_ai_model_list_by_workspace(args[0], ws_id)
if any(str(model.id) == str(args[2].custom_model) for model in _ai_model_list):
specialized_model_id = args[2].custom_model
print("use custom model: id[" + specialized_model_id + "]")
config: LLMConfig = await get_default_config(specialized_model_id)
# 构造函数可能同步调用宿主 API（type==1 高级应用动态数据源，
# LLMService.__init__ -> AssistantOutDsFactory.get_instance -> requests.get），
# 必须在工作线程执行，避免阻塞事件循环（Issue #1288）
instance = await asyncio.to_thread(cls, *args, **kwargs, config=config)
chat_params: list[SysArgModel] = await get_groups(args[0], "chat")
for config in chat_params:
if config.pkey == 'chat.sqlbot_name':
if config.pval.strip():
instance.chat_question.sqlbot_name = config.pval
if config.pkey == 'chat.limit_rows':
if config.pval.lower().strip() == 'true':
instance.enable_sql_row_limit = True
else:
instance.enable_sql_row_limit = False
if config.pkey == 'chat.context_record_count':
count_value = config.pval
if count_value is None:
count_value = settings.GENERATE_SQL_QUERY_HISTORY_ROUND_COUNT
count_value = int(count_value)
if count_value < 0:
count_value = 0
instance.base_message_round_count_limit = count_value
return instance
def is_running(self, timeout=0.5):
try:
r = concurrent.futures.wait([self.future], timeout)
if len(r.not_done) > 0:
return True
else:
return False
except Exception as e:
return True
def init_messages(self, session: Session):
self.table_name_list = self.choose_table_schema(session)
last_sql_messages: List[dict[str, Any]] = self.generate_sql_logs[-1].messages if len(
self.generate_sql_logs) > 0 else []
if self.chat_question.regenerate_record_id:
# filter record before regenerate_record_id
_temp_log = next(
filter(lambda obj: obj.pid == self.chat_question.regenerate_record_id, self.generate_sql_logs), None)
last_sql_messages: List[dict[str, Any]] = _temp_log.messages if _temp_log else []
# 排除所有的系统提示词
last_sql_messages = [obj for obj in last_sql_messages if obj.get("sqlbot_system") != True]
count_limit = self.base_message_round_count_limit
self.sql_message = []
# add sys prompt
_system_templates = self.chat_question.sql_sys_question(self.ds.type, self.enable_sql_row_limit)
self.sql_message.append(SystemPromptMessage(content=_system_templates['system']))
self.sql_message.append(HumanPromptMessage(content=_system_templates['rules']))
self.sql_message.append(
AIPromptMessage(content='我已掌握所有规则，包括表结构、SQL规范、安全限制和输出格式，我会严格遵守这些规则。'))
self.sql_message.append(HumanPromptMessage(content=_system_templates['schema']))
self.sql_message.append(
AIPromptMessage(content='我已确认您提供的数据库信息与表结构schema，我生成的SQL不会超出您提供的范围。'))
if _system_templates.get('custom_prompt'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['custom_prompt']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的额外信息，我会进行参考。'))
if _system_templates.get('terminologies'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['terminologies']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的术语信息，我会进行参考。'))
if _system_templates.get('data_training'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['data_training']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的SQL示例，我会进行参考。'))
if last_sql_messages is not None and len(last_sql_messages) > 0:
last_rounds = get_last_conversation_rounds(last_sql_messages, rounds=count_limit)
for _msg_dict in last_rounds:
_msg: BaseMessage
if _msg_dict.get('type') == 'human':
_msg = HumanMessage(content=_msg_dict.get('content'))
self.sql_message.append(_msg)
elif _msg_dict.get('type') == 'ai':
_msg = AIMessage(content=_msg_dict.get('content'))
self.sql_message.append(_msg)
last_chart_messages: List[dict[str, Any]] = self.generate_chart_logs[-1].messages if len(
self.generate_chart_logs) > 0 else []
if self.chat_question.regenerate_record_id:
# filter record before regenerate_record_id
_temp_log = next(
filter(lambda obj: obj.pid == self.chat_question.regenerate_record_id, self.generate_chart_logs), None)
last_chart_messages: List[dict[str, Any]] = _temp_log.messages if _temp_log else []
# 排除所有的系统提示词
last_chart_messages = [obj for obj in last_chart_messages if obj.get("sqlbot_system") != True]
count_chart_limit = self.base_message_round_count_limit
self.chart_message = []
# add sys prompt
_chart_system_templates = self.chat_question.chart_sys_question()
self.chart_message.append(SystemPromptMessage(content=_chart_system_templates['system']))
self.chart_message.append(HumanPromptMessage(content=_chart_system_templates['rules']))
self.chart_message.append(AIPromptMessage(content='我已掌握所有规则，我会严格遵守这些规则来生成符合要求的JSON。'))
if last_chart_messages is not None and len(last_chart_messages) > 0:
last_rounds = get_last_conversation_rounds(last_chart_messages, rounds=count_chart_limit)
for _msg_dict in last_rounds:
_msg: BaseMessage
if _msg_dict.get('type') == 'human':
_msg = HumanMessage(content=_msg_dict.get('content'))
self.chart_message.append(_msg)
elif _msg_dict.get('type') == 'ai':
_msg = AIMessage(content=_msg_dict.get('content'))
self.chart_message.append(_msg)
def init_record(self, session: Session) -> ChatRecord:
self.record = save_question(session=session, current_user=self.current_user, question=self.chat_question)
return self.record
def get_record(self):
return self.record
def set_record(self, record: ChatRecord):
self.record = record
def set_articles_number(self, articles_number: int):
self.articles_number = articles_number
def get_fields_from_chart(self, _session: Session):
chart_info = get_chart_config(_session, self.record.id)
return format_chart_fields(chart_info)
def filter_terminology_template(self, _session: Session, oid: int = None, ds_id: int = None):
self.current_logs[OperationEnum.FILTER_TERMS] = start_log(session=_session,
operate=OperationEnum.FILTER_TERMS,
record_id=self.record.id, local_operation=True)
calculate_oid = oid
calculate_ds_id = ds_id
if self.current_assistant:
calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid
if self.current_assistant.type == 1:
calculate_ds_id = None
# 使用原始提取的关键词（未扩展）进行术语匹配，避免同义词重复查询
match_text = self.chat_question.extracted_keywords or self.chat_question.question
if self.current_assistant and self.current_assistant.type == 1:
self.chat_question.terminologies, term_list = get_terminology_template(_session,
match_text,
calculate_oid,
None, self.current_assistant.id)
else:
self.chat_question.terminologies, term_list = get_terminology_template(_session,
match_text,
calculate_oid,
calculate_ds_id)
self.current_logs[OperationEnum.FILTER_TERMS] = end_log(session=_session,
log=self.current_logs[OperationEnum.FILTER_TERMS],
full_message=term_list)
def filter_custom_prompts(self, _session: Session, custom_prompt_type: CustomPromptTypeEnum, oid: int = None,
ds_id: int = None):
if SQLBotLicenseUtil.valid():
self.current_logs[OperationEnum.FILTER_CUSTOM_PROMPT] = start_log(session=_session,
operate=OperationEnum.FILTER_CUSTOM_PROMPT,
record_id=self.record.id,
local_operation=True)
calculate_oid = oid
calculate_ds_id = ds_id
if self.current_assistant:
calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid
if self.current_assistant.type == 1:
calculate_ds_id = None
if self.current_assistant and self.current_assistant.type == 1:
self.chat_question.custom_prompt, prompt_list = find_custom_prompts(_session,
custom_prompt_type,
calculate_oid,
None, self.current_assistant.id)
else:
self.chat_question.custom_prompt, prompt_list = find_custom_prompts(_session,
custom_prompt_type,
calculate_oid,
calculate_ds_id)
self.current_logs[OperationEnum.FILTER_CUSTOM_PROMPT] = end_log(session=_session,
log=self.current_logs[
OperationEnum.FILTER_CUSTOM_PROMPT],
full_message=prompt_list)
def filter_training_template(self, _session: Session, oid: int = None, ds_id: int = None):
self.current_logs[OperationEnum.FILTER_SQL_EXAMPLE] = start_log(session=_session,
operate=OperationEnum.FILTER_SQL_EXAMPLE,
record_id=self.record.id,
local_operation=True)
calculate_oid = oid
calculate_ds_id = ds_id
if self.current_assistant:
calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid
if self.current_assistant.type == 1:
calculate_ds_id = None
if self.current_assistant and self.current_assistant.type == 1:
self.chat_question.data_training, example_list = get_training_template(_session,
self.chat_question.question,
calculate_oid,
None, self.current_assistant.id)
else:
self.chat_question.data_training, example_list = get_training_template(_session,
self.chat_question.question,
calculate_oid,
calculate_ds_id)
self.current_logs[OperationEnum.FILTER_SQL_EXAMPLE] = end_log(session=_session,
log=self.current_logs[
OperationEnum.FILTER_SQL_EXAMPLE],
full_message=example_list)
def extract_keywords(self, _session: Session) -> tuple[bool, str]:
"""从用户问题中提取核心业务实体词。
使用 LLM 提取名词性业务实体，忽略时间词、动作词、修饰词。
支持上下文历史，可根据对话上下文理解用户意图。
提取结果可复用于表匹配、术语扩展等环节。
Returns:
(is_error, result) 元组：
- (False, keywords) 成功，关键词逗号分隔
- (True, error_message) 错误（LLM 失败或非查数据意图）
"""
system_prompt = self.chat_question.extract_keywords_sys_prompt()
keywords_msg: List[Union[BaseMessage, dict[str, Any]]] = []
keywords_msg.append(SystemPromptMessage(content=system_prompt))
# 加载上下文历史（仅取用户的历史提问，不含 AI 生成的 SQL 等回复）
last_sql_messages: List[dict[str, Any]] = self.generate_sql_logs[-1].messages if len(
self.generate_sql_logs) > 0 else []
if self.chat_question.regenerate_record_id:
_temp_log = next(
filter(lambda obj: obj.pid == self.chat_question.regenerate_record_id, self.generate_sql_logs), None)
last_sql_messages: List[dict[str, Any]] = _temp_log.messages if _temp_log else []
# 排除系统提示词和 AI 回复，只保留用户的历史提问
last_user_messages = [
obj for obj in last_sql_messages
if obj.get("sqlbot_system") != True and obj.get('type') == 'human'
]
if last_user_messages:
last_rounds = get_last_conversation_rounds(last_user_messages, rounds=self.base_message_round_count_limit)
for _msg_dict in last_rounds:
content = _msg_dict.get('content', '')
# 提取 标签内的用户原始提问，排除 error-msg 等干扰信息
match = re.search(r'(.*?)', content, re.DOTALL)
if match:
question_text = match.group(1).strip()
if question_text:
keywords_msg.append(HumanMessage(content=question_text))
# 当前问题
keywords_msg.append(HumanMessage(content=self.chat_question.question))
self.current_logs[OperationEnum.EXTRACT_KEYWORDS] = start_log(
session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.EXTRACT_KEYWORDS,
record_id=self.record.id,
full_message=[{'type': msg.type, 'content': msg.content} for msg in keywords_msg]
)
full_thinking_text = ''
full_text = ''
token_usage = {}
res = process_stream(self.llm.stream(keywords_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
result_text = full_text.strip()
keywords_msg.append(AIMessage(result_text))
self.current_logs[OperationEnum.EXTRACT_KEYWORDS] = end_log(
session=_session,
log=self.current_logs[OperationEnum.EXTRACT_KEYWORDS],
full_message=[{'type': msg.type, 'content': msg.content} for msg in keywords_msg],
reasoning_content=full_thinking_text if full_thinking_text else None,
token_usage=token_usage
)
if result_text.startswith('ERROR:NOT_DATA_QUERY'):
# return True, '您的问题似乎不是查询数据的问题，请提出与数据查询相关的问题。'
return False, ''
if result_text == 'EMPTY' or not result_text:
self.chat_question.extracted_keywords = self.chat_question.question
return False, ''
self.chat_question.extracted_keywords = result_text
return False, result_text
def choose_table_schema(self, _session: Session):
self.current_logs[OperationEnum.CHOOSE_TABLE] = start_log(session=_session,
operate=OperationEnum.CHOOSE_TABLE,
record_id=self.record.id,
local_operation=True)
# 使用扩展后的关键词进行表匹配（包含同义词，匹配更全面）
keywords = self.chat_question.expanded_keywords or self.chat_question.question
self.chat_question.db_schema, tables = self.out_ds_instance.get_db_schema(
self.ds.id, self.chat_question.question) if self.out_ds_instance else get_table_schema(
session=_session,
current_user=self.current_user,
ds=self.ds,
question=self.chat_question.question,
keywords=keywords)
# Get sample data for all tables
if not self.out_ds_instance and settings.SAMPLE_DATA_ENABLED:
self.chat_question.sample_data = get_tables_sample_data(
session=_session,
current_user=self.current_user,
ds=self.ds,
table_list=tables)
self.current_logs[OperationEnum.CHOOSE_TABLE] = end_log(session=_session,
log=self.current_logs[OperationEnum.CHOOSE_TABLE],
full_message=self.chat_question.db_schema)
return tables
def generate_analysis(self, _session: Session):
fields = self.get_fields_from_chart(_session)
self.chat_question.fields = orjson.dumps(fields).decode()
data = get_chat_chart_data(_session, self.record.id)
self.chat_question.data = orjson.dumps(data.get('data')).decode()
analysis_msg: List[Union[BaseMessage, dict[str, Any]]] = []
# 从 record 加载已提取的关键词（SQL 生成阶段已持久化）
if self.record.extracted_keywords:
self.chat_question.extracted_keywords = self.record.extracted_keywords
if self.record.expanded_keywords:
self.chat_question.expanded_keywords = self.record.expanded_keywords
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
self.filter_terminology_template(_session, self.oid, ds_id)
self.filter_custom_prompts(_session, CustomPromptTypeEnum.ANALYSIS, self.oid, ds_id)
analysis_msg.append(SystemPromptMessage(content=self.chat_question.analysis_sys_question()))
analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question()))
self.current_logs[OperationEnum.ANALYSIS] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.ANALYSIS,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in analysis_msg])
full_thinking_text = ''
full_analysis_text = ''
token_usage = {}
res = process_stream(self.llm.stream(analysis_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_analysis_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
analysis_msg.append(AIMessage(full_analysis_text))
self.current_logs[OperationEnum.ANALYSIS] = end_log(session=_session,
log=self.current_logs[
OperationEnum.ANALYSIS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in analysis_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
self.record = save_analysis_answer(session=_session, record_id=self.record.id,
answer=orjson.dumps({'content': full_analysis_text}).decode())
def generate_predict(self, _session: Session):
fields = self.get_fields_from_chart(_session)
self.chat_question.fields = orjson.dumps(fields).decode()
data = get_chat_chart_data(_session, self.record.id)
self.chat_question.data = orjson.dumps(data.get('data')).decode()
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
self.filter_custom_prompts(_session, CustomPromptTypeEnum.PREDICT_DATA, self.oid, ds_id)
predict_msg: List[Union[BaseMessage, dict[str, Any]]] = []
predict_msg.append(SystemPromptMessage(content=self.chat_question.predict_sys_question()))
predict_msg.append(HumanMessage(content=self.chat_question.predict_user_question()))
self.current_logs[OperationEnum.PREDICT_DATA] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.PREDICT_DATA,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in predict_msg])
full_thinking_text = ''
full_predict_text = ''
token_usage = {}
res = process_stream(self.llm.stream(predict_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_predict_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
predict_msg.append(AIMessage(full_predict_text))
self.record = save_predict_answer(session=_session, record_id=self.record.id,
answer=orjson.dumps({'content': full_predict_text}).decode())
self.current_logs[OperationEnum.PREDICT_DATA] = end_log(session=_session,
log=self.current_logs[
OperationEnum.PREDICT_DATA],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in predict_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
def generate_recommend_questions_task(self, _session: Session):
# get schema
if self.ds and not self.chat_question.db_schema:
self.chat_question.db_schema, tables = self.out_ds_instance.get_db_schema(
self.ds.id, self.chat_question.question) if self.out_ds_instance else get_table_schema(
session=_session,
current_user=self.current_user, ds=self.ds,
question=self.chat_question.question,
embedding=False)
# Get sample data for all tables
# if not self.out_ds_instance:
# self.chat_question.sample_data = get_tables_sample_data(
# session=_session,
# current_user=self.current_user,
# ds=self.ds)
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
guess_msg.append(SystemPromptMessage(content=self.chat_question.guess_sys_question(self.articles_number)))
old_questions = list(map(lambda q: q.strip(), get_old_questions(_session, self.record.datasource)))
guess_msg.append(
HumanMessage(content=self.chat_question.guess_user_question(orjson.dumps(old_questions).decode())))
self.current_logs[OperationEnum.GENERATE_RECOMMENDED_QUESTIONS] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.GENERATE_RECOMMENDED_QUESTIONS,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in guess_msg])
full_thinking_text = ''
full_guess_text = ''
token_usage = {}
res = process_stream(self.llm.stream(guess_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_guess_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
guess_msg.append(AIMessage(full_guess_text))
self.current_logs[OperationEnum.GENERATE_RECOMMENDED_QUESTIONS] = end_log(session=_session,
log=self.current_logs[
OperationEnum.GENERATE_RECOMMENDED_QUESTIONS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in guess_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
self.record = save_recommend_question_answer(session=_session, record_id=self.record.id,
answer={'content': full_guess_text},
articles_number=self.articles_number)
yield {'recommended_question': self.record.recommended_question}
def select_datasource(self, _session: Session):
datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
datasource_msg.append(SystemPromptMessage(self.chat_question.datasource_sys_question()))
if self.current_assistant and self.current_assistant.type != 4:
_ds_list = get_assistant_ds(session=_session, llm_service=self)
else:
stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(
and_(CoreDatasource.oid == self.oid))
_ds_list = [
{
"id": ds.id,
"name": ds.name,
"description": ds.description
}
for ds in _session.exec(stmt)
]
if not _ds_list:
raise SingleMessageError('No available datasource configuration found')
ignore_auto_select = _ds_list and len(_ds_list) == 1
# ignore auto select ds
full_thinking_text = ''
full_text = ''
if not ignore_auto_select:
if settings.TABLE_EMBEDDING_ENABLED and (
not self.current_assistant or (self.current_assistant and self.current_assistant.type != 1)):
_ds_list = get_ds_embedding(_session, _ds_list, self.out_ds_instance,
self.chat_question.question, self.current_assistant)
# yield {'content': '{"id":' + str(ds.get('id')) + '}'}
_ds_list_dict = []
for _ds in _ds_list:
_ds_list_dict.append(_ds)
datasource_msg.append(
HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
self.current_logs[OperationEnum.CHOOSE_DATASOURCE] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.CHOOSE_DATASOURCE,
record_id=self.record.id,
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for
msg in datasource_msg])
token_usage = {}
res = process_stream(self.llm.stream(datasource_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
datasource_msg.append(AIMessage(full_text))
self.current_logs[OperationEnum.CHOOSE_DATASOURCE] = end_log(session=_session,
log=self.current_logs[
OperationEnum.CHOOSE_DATASOURCE],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in datasource_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
json_str = extract_nested_json(full_text)
if json_str is None:
raise SingleMessageError(f'Cannot parse datasource from answer: {full_text}')
ds = orjson.loads(json_str)
_error: Exception | None = None
_datasource: int | None = None
_engine_type: str | None = None
try:
data: dict = _ds_list[0] if ignore_auto_select else ds
if data.get('id') and data.get('id') != 0:
_datasource = data['id']
_chat = _session.get(Chat, self.record.chat_id)
_chat.datasource = _datasource
if self.current_assistant and self.current_assistant.type in dynamic_ds_types:
_ds = self.out_ds_instance.get_ds(data['id'])
self.ds = _ds
self.chat_question.engine = _ds.type + get_version(self.ds)
_engine_type = self.chat_question.engine
_chat.engine_type = _ds.type
else:
_ds = _session.get(CoreDatasource, _datasource)
if not _ds:
_datasource = None
raise SingleMessageError(f"Datasource configuration with id {_datasource} not found")
self.ds = CoreDatasource(**_ds.model_dump())
self.chat_question.engine = (_ds.type_name if _ds.type != 'excel' else 'PostgreSQL') + get_version(
self.ds)
_engine_type = self.chat_question.engine
_chat.engine_type = _ds.type_name
# save chat
with _session.begin_nested():
# 为了能继续记日志，先单独处理下事务
try:
_session.add(_chat)
_session.flush()
_session.refresh(_chat)
_session.commit()
except Exception as e:
_session.rollback()
raise e
elif data['fail']:
raise SingleMessageError(data['fail'])
else:
raise SingleMessageError('No available datasource configuration found')
except Exception as e:
_error = e
if not ignore_auto_select and not settings.TABLE_EMBEDDING_ENABLED:
self.record = save_select_datasource_answer(session=_session, record_id=self.record.id,
answer=orjson.dumps({'content': full_text}).decode(),
datasource=_datasource,
engine_type=_engine_type)
if self.ds:
oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
# 提取关键词（前置步骤，结果复用于术语匹配和表匹配）
if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
is_error, kw_result = self.extract_keywords(_session)
if is_error:
raise SingleMessageError(kw_result)
if kw_result:
self.chat_question.extracted_keywords = kw_result
SQLBotLogUtil.info(f"提取的关键词: {kw_result}")
# 术语同义词扩展（用于表匹配，不用于术语模板查询）
expanded = expand_with_terminology(kw_result, _session, oid)
self.chat_question.expanded_keywords = expanded
SQLBotLogUtil.info(f"扩展后的关键词: {expanded}")
# 持久化到 record，供后续独立接口（如 generate_analysis）使用
save_extracted_keywords(_session, self.record.id,
expanded_keywords=expanded, extracted_keywords=kw_result)
self.filter_terminology_template(_session, oid, ds_id)
self.filter_training_template(_session, oid, ds_id)
self.filter_custom_prompts(_session, CustomPromptTypeEnum.GENERATE_SQL, oid, ds_id)
self.init_messages(_session)
if _error:
raise _error
def generate_sql(self, _session: Session):
# append current question
self.sql_message.append(HumanMessage(
self.chat_question.sql_user_question(current_time=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
change_title=self.change_title)))
self.current_logs[OperationEnum.GENERATE_SQL] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.GENERATE_SQL,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for msg
in self.sql_message])
full_thinking_text = ''
full_sql_text = ''
token_usage = {}
res = process_stream(self.llm.stream(self.sql_message), token_usage)
for chunk in res:
if chunk.get('content'):
full_sql_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
self.sql_message.append(AIMessage(full_sql_text))
self.current_logs[OperationEnum.GENERATE_SQL] = end_log(session=_session,
log=self.current_logs[OperationEnum.GENERATE_SQL],
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in self.sql_message],
reasoning_content=full_thinking_text,
token_usage=token_usage)
self.record = save_sql_answer(session=_session, record_id=self.record.id,
answer=orjson.dumps({'content': full_sql_text}).decode())
def generate_with_sub_sql(self, session: Session, 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
dynamic_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = []
dynamic_sql_msg.append(SystemPromptMessage(content=self.chat_question.dynamic_sys_question()))
dynamic_sql_msg.append(HumanMessage(content=self.chat_question.dynamic_user_question()))
self.current_logs[OperationEnum.GENERATE_DYNAMIC_SQL] = start_log(session=session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.GENERATE_DYNAMIC_SQL,
record_id=self.record.id,
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for
msg in dynamic_sql_msg])
full_thinking_text = ''
full_dynamic_text = ''
token_usage = {}
res = process_stream(self.llm.stream(dynamic_sql_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_dynamic_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
dynamic_sql_msg.append(AIMessage(full_dynamic_text))
self.current_logs[OperationEnum.GENERATE_DYNAMIC_SQL] = end_log(session=session,
log=self.current_logs[
OperationEnum.GENERATE_DYNAMIC_SQL],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in dynamic_sql_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
SQLBotLogUtil.info(full_dynamic_text)
return full_dynamic_text
def generate_assistant_dynamic_sql(self, _session: Session, sql, tables: List):
ds: AssistantOutDsSchema = self.ds
sub_query = []
result_dict = {}
for table in ds.tables:
if table.name in tables and table.sql:
# sub_query.append({"table": table.name, "query": table.sql})
result_dict[table.name] = table.sql
sub_query.append({"table": table.name, "query": f'{dynamic_subsql_prefix}{table.name}'})
if not sub_query:
return None
temp_sql_text = self.generate_with_sub_sql(session=_session, sql=sql, sub_mappings=sub_query)
result_dict['sqlbot_temp_sql_text'] = temp_sql_text
return result_dict
def build_table_filter(self, session: Session, sql: str, filters: list):
filter = json.dumps(filters, ensure_ascii=False)
self.chat_question.sql = sql
self.chat_question.filter = filter
permission_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = []
permission_sql_msg.append(SystemPromptMessage(content=self.chat_question.filter_sys_question()))
permission_sql_msg.append(HumanMessage(content=self.chat_question.filter_user_question()))
self.current_logs[OperationEnum.GENERATE_SQL_WITH_PERMISSIONS] = start_log(session=session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.GENERATE_SQL_WITH_PERMISSIONS,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in permission_sql_msg])
full_thinking_text = ''
full_filter_text = ''
token_usage = {}
res = process_stream(self.llm.stream(permission_sql_msg), token_usage)
for chunk in res:
if chunk.get('content'):
full_filter_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
permission_sql_msg.append(AIMessage(full_filter_text))
self.current_logs[OperationEnum.GENERATE_SQL_WITH_PERMISSIONS] = end_log(session=session,
log=self.current_logs[
OperationEnum.GENERATE_SQL_WITH_PERMISSIONS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in permission_sql_msg],
reasoning_content=full_thinking_text,
token_usage=token_usage)
SQLBotLogUtil.info(full_filter_text)
return full_filter_text
def generate_filter(self, _session: Session, sql: str, tables: List):
filters = get_row_permission_filters(session=_session, current_user=self.current_user, ds=self.ds,
tables=tables)
if not filters:
return None
return self.build_table_filter(session=_session, sql=sql, filters=filters)
def generate_assistant_filter(self, _session: Session, sql, tables: List):
ds: AssistantOutDsSchema = self.ds
filters = []
for table in ds.tables:
if table.name in tables and table.rule:
filters.append({"table": table.name, "filter": table.rule})
if not filters:
return None
return self.build_table_filter(session=_session, sql=sql, filters=filters)
def generate_chart(self, _session: Session, chart_type: Optional[str] = '', schema: Optional[str] = ''):
# append current question
self.chart_message.append(HumanMessage(self.chat_question.chart_user_question(chart_type, schema)))
self.current_logs[OperationEnum.GENERATE_CHART] = start_log(session=_session,
ai_modal_id=self.chat_question.ai_modal_id,
ai_modal_name=self.chat_question.ai_modal_name,
operate=OperationEnum.GENERATE_CHART,
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in self.chart_message])
full_thinking_text = ''
full_chart_text = ''
token_usage = {}
res = process_stream(self.llm.stream(self.chart_message), token_usage)
for chunk in res:
if chunk.get('content'):
full_chart_text += chunk.get('content')
if chunk.get('reasoning_content'):
full_thinking_text += chunk.get('reasoning_content')
yield chunk
self.chart_message.append(AIMessage(full_chart_text))
self.record = save_chart_answer(session=_session, record_id=self.record.id,
answer=orjson.dumps({'content': full_chart_text}).decode())
self.current_logs[OperationEnum.GENERATE_CHART] = end_log(session=_session,
log=self.current_logs[OperationEnum.GENERATE_CHART],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in self.chart_message],
reasoning_content=full_thinking_text,
token_usage=token_usage)
def check_sql(self, session: Session, res: str, operate: OperationEnum) -> tuple[str, Optional[list]]:
json_str = extract_nested_json(res)
log = self.current_logs[operate]
if json_str is None:
trigger_log_error(session, log)
raise SingleMessageError(orjson.dumps({'message': 'SQL answer is not a valid json object',
'traceback': "SQL answer is not a valid json object:\n" + res}).decode())
sql: str
data: dict
try:
data = orjson.loads(json_str)
if data['success']:
sql = data['sql']
else:
message = data['message']
raise SingleMessageError(message)
except SingleMessageError as e:
trigger_log_error(session, log)
raise e
except Exception:
trigger_log_error(session, log)
raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer',
'traceback': "Cannot parse sql from answer:\n" + res}).decode())
if sql.strip() == '':
trigger_log_error(session, log)
raise SingleMessageError("SQL query is empty")
return sql, data.get('tables')
@staticmethod
def get_chart_type_from_sql_answer(res: str) -> Optional[str]:
json_str = extract_nested_json(res)
if json_str is None:
return None
chart_type: Optional[str]
data: dict
try:
data = orjson.loads(json_str)
if data['success']:
chart_type = data['chart-type']
else:
return None
except Exception:
return None
return chart_type
@staticmethod
def get_brief_from_sql_answer(res: str) -> Optional[str]:
json_str = extract_nested_json(res)
if json_str is None:
return None
brief: Optional[str]
data: dict
try:
data = orjson.loads(json_str)
if data['success']:
brief = data['brief']
else:
return None
except Exception:
return None
return brief
def check_save_sql(self, session: Session, res: str, operate: OperationEnum) -> str:
sql, *_ = self.check_sql(session=session, res=res, operate=operate)
save_sql(session=session, sql=sql, record_id=self.record.id)
self.chat_question.sql = sql
return sql
def check_save_chart(self, session: Session, res: str) -> Dict[str, Any]:
json_str = extract_nested_json(res)
if json_str is None:
raise SingleMessageError(orjson.dumps({'message': 'Cannot parse chart config from answer',
'traceback': "Cannot parse chart config from answer:\n" + res}).decode())
data: dict
chart: Dict[str, Any] = {}
message = ''
error = False
try:
data = orjson.loads(json_str)
if data['type'] and data['type'] != 'error':
# todo type check
chart = data
if chart.get('columns'):
for v in chart.get('columns'):
v['value'] = v.get('value').lower()
if chart.get('axis'):
if chart.get('axis').get('x'):
chart.get('axis').get('x')['value'] = chart.get('axis').get('x').get('value').lower()
y_axis = chart.get('axis').get('y')
if y_axis:
if isinstance(y_axis, list):
# 数组格式: y: [{name, value}, ...]
for item in y_axis:
if item.get('value'):
item['value'] = item['value'].lower()
elif isinstance(y_axis, dict) and y_axis.get('value'):
# 旧格式: y: {name, value}
y_axis['value'] = y_axis['value'].lower()
if chart.get('axis').get('series'):
chart.get('axis').get('series')['value'] = chart.get('axis').get('series').get('value').lower()
if chart.get('axis') and chart['axis'].get('multi-quota'):
multi_quota = chart['axis']['multi-quota']
if multi_quota.get('value'):
if isinstance(multi_quota['value'], list):
# 将数组中的每个值转换为小写
multi_quota['value'] = [v.lower() if v else v for v in multi_quota['value']]
elif isinstance(multi_quota['value'], str):
# 如果是字符串，也转换为小写
multi_quota['value'] = multi_quota['value'].lower()
elif data['type'] == 'error':
message = data['reason']
error = True
else:
raise Exception('Chart is empty')
except Exception:
error = True
message = orjson.dumps({'message': 'Cannot parse chart config from answer',
'traceback': "Cannot parse chart config from answer:\n" + res}).decode()
if error:
raise SingleMessageError(message)
save_chart(session=session, chart=orjson.dumps(chart).decode(), record_id=self.record.id)
return chart
def check_save_predict_data(self, session: Session, res: str) -> bool:
json_str = extract_nested_json(res)
if not json_str:
json_str = ''
save_predict_data(session=session, record_id=self.record.id, data=json_str)
if json_str == '':
return False
return True
def save_error(self, session: Session, message: str):
return save_error_message(session=session, record_id=self.record.id, message=message)
def save_sql_data(self, session: Session, data_obj: Dict[str, Any]):
try:
data_result = data_obj.get('data')
limit = 1000
if data_result:
data_result = prepare_for_orjson(data_result)
if data_result and len(data_result) > limit and self.enable_sql_row_limit:
data_obj['data'] = data_result[:limit]
data_obj['limit'] = limit
else:
data_obj['data'] = data_result
data_obj['datasource'] = self.ds.id
return save_sql_exec_data(session=session, record_id=self.record.id,
data=orjson.dumps(data_obj).decode())
except Exception as e:
raise e
def finish(self, session: Session):
return finish_record(session=session, record_id=self.record.id)
def execute_sql(self, sql: str):
"""Execute SQL query
Args:
ds: Data source instance
sql: SQL query statement
Returns:
Query results
"""
SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}")
try:
return exec_sql(ds=self.ds, sql=sql, origin_column=False)
except Exception as e:
if isinstance(e, ParseSQLResultError):
raise e
else:
err = traceback.format_exc(limit=1, chain=True)
raise SQLBotDBError(err)
def pop_chunk(self):
try:
chunk = self.chunk_list.pop(0)
return chunk
except IndexError as e:
return None
def await_result(self):
while self.is_running():
while True:
chunk = self.pop_chunk()
if chunk is not None:
yield chunk
else:
break
while True:
chunk = self.pop_chunk()
if chunk is None:
break
yield chunk
def run_task_async(self, in_chat: bool = True, stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
if in_chat:
stream = True
self.future = executor.submit(self.run_task_cache, in_chat, stream, finish_step, return_img)
def run_task_cache(self, in_chat: bool = True, stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
for chunk in self.run_task(in_chat, stream, finish_step, return_img):
self.chunk_list.append(chunk)
def run_task(self, in_chat: bool = True, stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
json_result: Dict[str, Any] = {'success': True}
_session = None
try:
_session = session_maker()
if self.ds:
oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
# 提取关键词（前置步骤，结果复用于术语匹配和表匹配）
if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
is_error, kw_result = self.extract_keywords(_session)
if is_error:
raise SingleMessageError(kw_result)
if kw_result:
self.chat_question.extracted_keywords = kw_result
SQLBotLogUtil.info(f"提取的关键词: {kw_result}")
# 术语同义词扩展（用于表匹配，不用于术语模板查询）
expanded = expand_with_terminology(kw_result, _session, oid)
self.chat_question.expanded_keywords = expanded
SQLBotLogUtil.info(f"扩展后的关键词: {expanded}")
# 持久化到 record，供后续独立接口（如 generate_analysis）使用
save_extracted_keywords(_session, self.record.id,
expanded_keywords=expanded, extracted_keywords=kw_result)
self.filter_terminology_template(_session, oid, ds_id)
self.filter_training_template(_session, oid, ds_id)
self.filter_custom_prompts(_session, CustomPromptTypeEnum.GENERATE_SQL, oid, ds_id)
self.init_messages(_session)
# return id
if in_chat:
yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
if self.get_record().regenerate_record_id:
yield 'data:' + orjson.dumps({'type': 'regenerate_record_id',
'regenerate_record_id': self.get_record().regenerate_record_id}).decode() + '\n\n'
yield 'data:' + orjson.dumps(
{'type': 'question', 'question': self.get_record().question}).decode() + '\n\n'
else:
if stream:
yield '> ' + self.trans('i18n_chat.record_id_in_mcp') + str(self.get_record().id) + '\n'
yield '> ' + self.get_record().question + '\n\n'
if not stream:
json_result['record_id'] = self.get_record().id
# select datasource if datasource is none
if not self.ds:
ds_res = self.select_datasource(_session)
for chunk in ds_res:
SQLBotLogUtil.info(chunk)
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'datasource-result'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'id': self.ds.id, 'datasource_name': self.ds.name,
'engine_type': self.ds.type_name or self.ds.type,
'type': 'datasource'}).decode() + '\n\n'
else:
self.validate_history_ds(_session)
# check connection
connected = check_connection(ds=self.ds, trans=None)
if not connected:
raise SQLBotDBConnectionError('Connect DB failed')
# generate sql
sql_res = self.generate_sql(_session)
full_sql_text = ''
for chunk in sql_res:
full_sql_text += chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'sql-result'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
# filter sql
SQLBotLogUtil.info(full_sql_text)
chart_type = self.get_chart_type_from_sql_answer(full_sql_text)
# return title
if self.change_title:
llm_brief = self.get_brief_from_sql_answer(full_sql_text)
llm_brief_generated = bool(llm_brief)
if llm_brief_generated or (self.chat_question.question and self.chat_question.question.strip() != ''):
save_brief = llm_brief if (llm_brief and llm_brief != '') else self.chat_question.question.strip()[
:20]
brief = rename_chat(session=_session,
rename_object=RenameChat(id=self.get_record().chat_id,
brief=save_brief, brief_generate=llm_brief_generated))
if in_chat:
yield 'data:' + orjson.dumps({'type': 'brief', 'brief': brief}).decode() + '\n\n'
if not stream:
json_result['title'] = brief
use_dynamic_ds: bool = self.current_assistant and self.current_assistant.type in dynamic_ds_types
is_page_embedded: bool = self.current_assistant and self.current_assistant.type == 4
dynamic_sql_result = None
sqlbot_temp_sql_text = None
assistant_dynamic_sql = None
# row permission
sql_operate = OperationEnum.GENERATE_SQL
sql, tables = self.check_sql(session=_session, res=full_sql_text, operate=sql_operate)
# 表名安全检查：用 sqlglot 解析真实 SQL，不信任 AI 返回的 tables
actual_tables = extract_tables_from_sql(sql, ds_type=self.ds.type)
if not actual_tables:
raise SingleMessageError(
"SQL parsing failed: unable to extract table names. "
"This may indicate an unsupported SQL syntax or a security issue."
)
allowed_tables = set(self.table_name_list)
unauthorized_tables = actual_tables - allowed_tables
if unauthorized_tables:
raise SingleMessageError(
f"SQL contains unauthorized tables: {', '.join(unauthorized_tables)}. "
f"Allowed tables: {', '.join(allowed_tables)}"
)
if ((not self.current_assistant or is_page_embedded) and is_normal_user(
self.current_user)) or use_dynamic_ds:
sql_result = None
if use_dynamic_ds:
dynamic_sql_result = self.generate_assistant_dynamic_sql(_session, sql, tables)
sqlbot_temp_sql_text = dynamic_sql_result.get(
'sqlbot_temp_sql_text') if dynamic_sql_result else None
else:
sql_result = self.generate_filter(_session, sql, tables) # maybe no sql and tables
if sql_result:
SQLBotLogUtil.info(sql_result)
sql_operate = OperationEnum.GENERATE_SQL_WITH_PERMISSIONS
sql = self.check_save_sql(session=_session, res=sql_result, operate=sql_operate)
elif dynamic_sql_result and sqlbot_temp_sql_text:
sql_operate = OperationEnum.GENERATE_DYNAMIC_SQL
assistant_dynamic_sql = self.check_save_sql(session=_session, res=sqlbot_temp_sql_text,
operate=sql_operate)
else:
sql = self.check_save_sql(session=_session, res=full_sql_text, operate=sql_operate)
else:
sql = self.check_save_sql(session=_session, res=full_sql_text, operate=sql_operate)
SQLBotLogUtil.info('sql: ' + sql)
if not stream:
json_result['sql'] = sql
format_sql = sqlparse.format(sql, reindent=True)
if in_chat:
yield 'data:' + orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n'
else:
if stream:
yield f'```sql\n{format_sql}\n```\n\n'
# execute sql
real_execute_sql = sql
if sqlbot_temp_sql_text and assistant_dynamic_sql:
dynamic_sql_result.pop('sqlbot_temp_sql_text')
for origin_table, subsql in dynamic_sql_result.items():
assistant_dynamic_sql = assistant_dynamic_sql.replace(f'{dynamic_subsql_prefix}{origin_table}',
subsql)
real_execute_sql = assistant_dynamic_sql
if finish_step.value <= ChatFinishStep.GENERATE_SQL.value:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n'
if not stream:
yield json_result
return
self.current_logs[OperationEnum.EXECUTE_SQL] = start_log(session=_session,
operate=OperationEnum.EXECUTE_SQL,
record_id=self.record.id, local_operation=True)
result = self.execute_sql(sql=real_execute_sql)
self.current_logs[OperationEnum.EXECUTE_SQL] = end_log(session=_session,
log=self.current_logs[OperationEnum.EXECUTE_SQL],
full_message={'sql': real_execute_sql,
'count': len(result.get('data'))})
_data = DataFormat.convert_large_numbers_in_object_array(result.get('data'))
_data = DataFormat.normalize_qualified_sql_column_keys_in_object_array(_data)
result["data"] = _data
self.save_sql_data(session=_session, data_obj=result)
if in_chat:
yield 'data:' + orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n'
if not stream:
json_result['data'] = get_chat_chart_data(_session, self.record.id)
if finish_step.value <= ChatFinishStep.QUERY_DATA.value:
if stream:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n'
else:
_column_list = []
for field in result.get('fields'):
_column_list.append(AxisObj(name=field, value=field))
md_data, _fields_list = DataFormat.convert_object_array_for_pandas(_column_list,
result.get('data'))
# data, _fields_list, col_formats = self.format_pd_data(_column_list, result.get('data'))
if not _data or not _fields_list:
yield 'The SQL execution result is empty.\n\n'
else:
df = pd.DataFrame(_data, columns=_fields_list)
df_safe = DataFormat.safe_convert_to_string(df)
markdown_table = df_safe.to_markdown(index=False)
yield markdown_table + '\n\n'
else:
yield json_result
return
# generate chart
used_tables_schema, used_tables = self.out_ds_instance.get_db_schema(
self.ds.id, self.chat_question.question, embedding=False,
table_list=tables) if self.out_ds_instance else get_table_schema(
session=_session,
current_user=self.current_user,
ds=self.ds,
question=self.chat_question.question,
embedding=False, table_list=tables)
SQLBotLogUtil.info('used_tables_schema: \n' + used_tables_schema)
chart_res = self.generate_chart(_session, chart_type, used_tables_schema)
full_chart_text = ''
for chunk in chart_res:
full_chart_text += chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'chart-result'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
# filter chart
SQLBotLogUtil.info(full_chart_text)
chart = self.check_save_chart(session=_session, res=full_chart_text)
SQLBotLogUtil.info(chart)
if not stream:
json_result['chart'] = chart
if in_chat:
yield 'data:' + orjson.dumps(
{'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
else:
if stream:
md_data, _fields_list = DataFormat.convert_data_fields_for_pandas(chart, result.get('fields'),
result.get('data'))
# data, _fields_list, col_formats = self.format_pd_data(_column_list, result.get('data'))
if not md_data or not _fields_list:
yield 'The SQL execution result is empty.\n\n'
else:
df = pd.DataFrame(md_data, columns=_fields_list)
df_safe = DataFormat.safe_convert_to_string(df)
markdown_table = df_safe.to_markdown(index=False)
yield markdown_table + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n'
else:
# generate picture
try:
if chart.get('type') != 'table' and return_img:
# yield '### generated chart picture\n\n'
self.current_logs[OperationEnum.GENERATE_PICTURE] = start_log(session=_session,
operate=OperationEnum.GENERATE_PICTURE,
record_id=self.record.id,
local_operation=True)
image_url, error = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(result))
SQLBotLogUtil.info(image_url)
if stream:
yield f''
else:
json_result['image_url'] = image_url
if error is not None:
raise error
self.current_logs[OperationEnum.GENERATE_PICTURE] = end_log(session=_session,
log=self.current_logs[
OperationEnum.GENERATE_PICTURE],
full_message=image_url)
except Exception as e:
if stream:
if chart.get('type') != 'table':
yield 'generate or fetch chart picture error.\n\n'
raise e
if not stream:
yield json_result
except Exception as e:
traceback.print_exc()
error_msg: str
if isinstance(e, SingleMessageError):
error_msg = str(e)
elif isinstance(e, SQLBotDBConnectionError):
error_msg = orjson.dumps(
{'message': str(e), 'type': 'db-connection-err'}).decode()
elif isinstance(e, SQLBotDBError):
error_msg = orjson.dumps(
{'message': 'Execute SQL Failed', 'traceback': str(e), 'type': 'exec-sql-err'}).decode()
else:
error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
if _session:
self.save_error(session=_session, message=error_msg)
if in_chat:
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
else:
if stream:
yield f'❌ **ERROR:**\n'
yield f'> {error_msg}\n'
else:
json_result['success'] = False
json_result['message'] = error_msg
yield json_result
finally:
self.finish(_session)
session_maker.remove()
def run_recommend_questions_task_async(self):
self.future = executor.submit(self.run_recommend_questions_task_cache)
def run_recommend_questions_task_cache(self):
for chunk in self.run_recommend_questions_task():
self.chunk_list.append(chunk)
def run_recommend_questions_task(self):
try:
_session = session_maker()
res = self.generate_recommend_questions_task(_session)
for chunk in res:
if chunk.get('recommended_question'):
yield 'data:' + orjson.dumps(
{'content': chunk.get('recommended_question'),
'type': 'recommended_question'}).decode() + '\n\n'
else:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'recommended_question_result'}).decode() + '\n\n'
except Exception:
traceback.print_exc()
finally:
session_maker.remove()
def run_analysis_or_predict_task_async(self, session: Session, action_type: str, base_record: ChatRecord,
in_chat: bool = True, stream: bool = True):
self.set_record(save_analysis_predict_record(session, base_record, action_type))
self.future = executor.submit(self.run_analysis_or_predict_task_cache, action_type, in_chat, stream)
def run_analysis_or_predict_task_cache(self, action_type: str, in_chat: bool = True, stream: bool = True):
for chunk in self.run_analysis_or_predict_task(action_type, in_chat, stream):
self.chunk_list.append(chunk)
def run_analysis_or_predict_task(self, action_type: str, in_chat: bool = True, stream: bool = True):
json_result: Dict[str, Any] = {'success': True}
_session = None
try:
_session = session_maker()
if in_chat:
yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
else:
if stream:
yield '> ' + self.trans('i18n_chat.record_id_in_mcp') + str(self.get_record().id) + '\n'
yield '> ' + self.get_record().question + '\n\n'
if not stream:
json_result['record_id'] = self.get_record().id
if action_type == 'analysis':
# generate analysis
analysis_res = self.generate_analysis(_session)
full_text = ''
for chunk in analysis_res:
full_text += chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'analysis-result'}).decode() + '\n\n'
else:
if stream:
yield chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
yield 'data:' + orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
else:
if stream:
yield '\n\n'
if not stream:
json_result['content'] = full_text
elif action_type == 'predict':
# generate predict
analysis_res = self.generate_predict(_session)
full_text = ''
for chunk in analysis_res:
full_text += chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'predict-result'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
has_data = self.check_save_predict_data(session=_session, res=full_text)
if has_data:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict-success'}).decode() + '\n\n'
else:
chart = get_chat_chart_config(_session, self.record.id)
origin_data = get_chat_chart_data(_session, self.record.id)
predict_data = get_chat_predict_data(_session, self.record.id)
if stream:
md_data, _fields_list = DataFormat.convert_data_fields_for_pandas(chart,
origin_data.get('fields'),
predict_data)
if not md_data or not _fields_list:
yield 'Predict data result is empty.\n\n'
else:
df = pd.DataFrame(md_data, columns=_fields_list)
df_safe = DataFormat.safe_convert_to_string(df)
markdown_table = df_safe.to_markdown(index=False)
yield markdown_table + '\n\n'
else:
json_result['origin_data'] = origin_data
json_result['predict_data'] = predict_data
# generate picture
try:
if chart.get('type') != 'table':
# yield '### generated chart picture\n\n'
_data = get_chat_chart_data(_session, self.record.id)
_data['data'] = _data.get('data') + predict_data
image_url, error = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(_data))
SQLBotLogUtil.info(image_url)
if stream:
yield f''
else:
json_result['image_url'] = image_url
if error is not None:
raise error
except Exception as e:
if stream:
if chart.get('type') != 'table':
yield 'generate or fetch chart picture error.\n\n'
raise e
else:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict-failed'}).decode() + '\n\n'
else:
if stream:
yield full_text + '\n\n'
if not stream:
json_result['success'] = False
json_result['message'] = full_text
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
self.finish(_session)
if not stream:
yield json_result
except Exception as e:
traceback.print_exc()
error_msg: str
if isinstance(e, SingleMessageError):
error_msg = str(e)
else:
error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
if _session:
self.save_error(session=_session, message=error_msg)
if in_chat:
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
else:
if stream:
yield f'❌ **ERROR:**\n'
yield f'> {error_msg}\n'
else:
json_result['success'] = False
json_result['message'] = error_msg
yield json_result
finally:
# end
session_maker.remove()
def validate_history_ds(self, session: Session):
_ds = self.ds
if not self.current_assistant or self.current_assistant.type == 4:
try:
current_ds = session.get(CoreDatasource, _ds.id)
if not current_ds:
raise SingleMessageError('chat.ds_is_invalid')
except Exception as e:
raise SingleMessageError("chat.ds_is_invalid")
else:
try:
_ds_list: list[dict] = get_assistant_ds(session=session, llm_service=self)
match_ds = any(item.get("id") == _ds.id for item in _ds_list)
if not match_ds:
type = self.current_assistant.type
msg = f"[please check ds list and public ds list]" if type == 0 else f"[please check ds api]"
raise SingleMessageError(msg)
except Exception as e:
raise SingleMessageError(f"ds is invalid [{str(e)}]")
def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
"""Execute SQL query using SQLDatabase
Args:
db: SQLDatabase instance
sql: SQL query statement
Returns:
str: Query results formatted as string
"""
try:
# Execute query
result = db.run(sql)
if not result:
return "Query executed successfully but returned no results."
# Format results
return str(result)
except Exception as e:
error_msg = f"SQL execution failed: {str(e)}"
SQLBotLogUtil.exception(error_msg)
raise RuntimeError(error_msg)
def request_picture(chat_id: int, record_id: int, chart: dict, data: dict):
file_name = f'c_{chat_id}_r_{record_id}'
columns = chart.get('columns') if chart.get('columns') else []
x = None
y = None
series = None
multi_quota_fields = []
multi_quota_name = None
if chart.get('axis'):
axis_data = chart.get('axis')
x = axis_data.get('x')
y = axis_data.get('y')
series = axis_data.get('series')
# 获取multi-quota字段列表
if axis_data.get('multi-quota') and 'value' in axis_data.get('multi-quota'):
multi_quota_fields = axis_data.get('multi-quota').get('value', [])
multi_quota_name = axis_data.get('multi-quota').get('name')
axis = []
for v in columns:
axis.append({'name': v.get('name'), 'value': v.get('value')})
if x:
axis.append({'name': x.get('name'), 'value': x.get('value'), 'type': 'x'})
if y:
y_list = y if isinstance(y, list) else [y]
for y_item in y_list:
if isinstance(y_item, dict) and 'value' in y_item:
y_obj = {
'name': y_item.get('name'),
'value': y_item.get('value'),
'type': 'y'
}
# 如果是multi-quota字段，添加标志
if y_item.get('value') in multi_quota_fields:
y_obj['multi-quota'] = True
axis.append(y_obj)
if series:
axis.append({'name': series.get('name'), 'value': series.get('value'), 'type': 'series'})
if multi_quota_name:
axis.append({'name': multi_quota_name, 'value': multi_quota_name, 'type': 'other-info'})
request_obj = {
"path": os.path.join(settings.MCP_IMAGE_PATH, file_name),
"type": chart.get('type'),
"data": orjson.dumps(data.get('data') if data.get('data') else []).decode(),
"axis": orjson.dumps(axis).decode(),
}
_error = None
try:
requests.post(url=settings.MCP_IMAGE_HOST, json=request_obj, timeout=settings.SERVER_IMAGE_TIMEOUT)
except Exception as e:
_error = e
request_path = urllib.parse.urljoin(settings.SERVER_IMAGE_HOST, f"{file_name}.png")
return request_path, _error
def get_token_usage(chunk: BaseMessageChunk, token_usage: dict = None):
try:
if chunk.usage_metadata:
if token_usage is None:
token_usage = {}
token_usage['input_tokens'] = chunk.usage_metadata.get('input_tokens')
token_usage['output_tokens'] = chunk.usage_metadata.get('output_tokens')
token_usage['total_tokens'] = chunk.usage_metadata.get('total_tokens')
except Exception:
pass
def process_stream(res: Iterator[BaseMessageChunk],
token_usage: Dict[str, Any] = None,
enable_tag_parsing: bool = settings.PARSE_REASONING_BLOCK_ENABLED,
start_tag: str = settings.DEFAULT_REASONING_CONTENT_START,
end_tag: str = settings.DEFAULT_REASONING_CONTENT_END
):
if token_usage is None:
token_usage = {}
in_thinking_block = False # 标记是否在思考过程块中
current_thinking = '' # 当前收集的思考过程内容
pending_start_tag = '' # 用于缓存可能被截断的开始标签部分
for chunk in res:
SQLBotLogUtil.info(chunk)
reasoning_content_chunk = ''
content = chunk.content
output_content = '' # 实际要输出的内容
# 检查additional_kwargs中的reasoning_content
if 'reasoning_content' in chunk.additional_kwargs:
reasoning_content = chunk.additional_kwargs.get('reasoning_content', '')
if reasoning_content is None:
reasoning_content = ''
# 累积additional_kwargs中的思考内容到current_thinking
current_thinking += reasoning_content
reasoning_content_chunk = reasoning_content
# 只有当current_thinking不是空字符串时才跳过标签解析
if not in_thinking_block and current_thinking.strip() != '':
output_content = content # 正常输出content
yield {
'content': output_content,
'reasoning_content': reasoning_content_chunk
}
get_token_usage(chunk, token_usage)
continue # 跳过后续的标签解析逻辑
# 如果没有有效的思考内容，并且启用了标签解析，才执行标签解析逻辑
# 如果有缓存的开始标签部分，先拼接当前内容
if pending_start_tag:
content = pending_start_tag + content
pending_start_tag = ''
# 检查是否开始思考过程块（处理可能被截断的开始标签）
if enable_tag_parsing and not in_thinking_block and start_tag:
if start_tag in content:
start_idx = content.index(start_tag)
# 只有当开始标签前面没有其他文本时才认为是真正的思考块开始
if start_idx == 0 or content[:start_idx].strip() == '':
# 完整标签存在且前面没有其他文本
output_content += content[:start_idx] # 输出开始标签之前的内容
content = content[start_idx + len(start_tag):] # 移除开始标签
in_thinking_block = True
else:
# 开始标签前面有其他文本，不认为是思考块开始
output_content += content
content = ''
else:
# 检查是否可能有部分开始标签
for i in range(1, len(start_tag)):
if content.endswith(start_tag[:i]):
# 只有当当前内容全是空白时才缓存部分标签
if content[:-i].strip() == '':
pending_start_tag = start_tag[:i]
content = content[:-i] # 移除可能的部分标签
output_content += content
content = ''
break
# 处理思考块内容
if enable_tag_parsing and in_thinking_block and end_tag:
if end_tag in content:
# 找到结束标签
end_idx = content.index(end_tag)
current_thinking += content[:end_idx] # 收集思考内容
reasoning_content_chunk += current_thinking # 添加到当前块的思考内容
content = content[end_idx + len(end_tag):] # 移除结束标签后的内容
current_thinking = '' # 重置当前思考内容
in_thinking_block = False
output_content += content # 输出结束标签之后的内容
else:
# 在遇到结束标签前，持续收集思考内容
current_thinking += content
reasoning_content_chunk += content
content = ''
else:
# 不在思考块中或标签解析未启用，正常输出
output_content += content
yield {
'content': output_content,
'reasoning_content': reasoning_content_chunk
}
get_token_usage(chunk, token_usage)
def get_lang_name(lang: str):
if not lang:
return '简体中文'
normalized = lang.lower()
if normalized.startswith('zh-tw'):
return '繁体中文'
if normalized.startswith('en'):
return '英文'
if normalized.startswith('ko'):
return '韩语'
if normalized.startswith('th'):
return '泰语'
return '简体中文'
def get_last_conversation_rounds(messages, rounds=settings.GENERATE_SQL_QUERY_HISTORY_ROUND_COUNT):
"""获取最后N轮对话，处理不完整对话的情况"""
if not messages or rounds <= 0:
return []
# 找到所有用户消息的位置
human_indices = []
for index, msg in enumerate(messages):
if msg.get('type') == 'human':
human_indices.append(index)
# 如果没有用户消息，返回空
if not human_indices:
return []
# 计算从哪个索引开始
if len(human_indices) <= rounds:
# 如果用户消息数少于等于需要的轮数，从第一个用户消息开始
start_index = human_indices[0]
else:
# 否则，从倒数第N个用户消息开始
start_index = human_indices[-rounds]
return messages[start_index:]