Skip to content

Commit 204a3df

Browse files
ulleoulleo
committed
feat: add chat log step display
1 parent bda2c66 commit 204a3df

File tree

2 files changed

+41
-17
lines changed

2 files changed

+41
-17
lines changed

‎backend/apps/chat/curd/chat.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -843,6 +843,17 @@ def end_log(session: SessionDep, log: ChatLog, full_message: Union[list[dict], d
843843
return log
844844

845845

846+
def trigger_log_error(session: SessionDep, log: ChatLog) -> ChatLog:
847+
log.error = True
848+
stmt = update(ChatLog).where(and_(ChatLog.id == log.id)).values(
849+
error=True
850+
)
851+
session.execute(stmt)
852+
session.commit()
853+
854+
return log
855+
856+
846857
def save_sql_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord:
847858
if not record_id:
848859
raise Exception("Record id cannot be None")

‎backend/apps/chat/task/llm.py‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \
3232
get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \
3333
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate, get_chat_predict_data, \
34-
get_chat_chart_config
34+
get_chat_chart_config, trigger_log_error
3535
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \
3636
ChatFinishStep, AxisObj
3737
from apps.data_training.curd.data_training import get_training_template
@@ -283,7 +283,7 @@ def filter_terminology_template(self, _session: Session, oid: int = None, ds_id:
283283
self.current_logs[OperationEnum.FILTER_TERMS] = start_log(session=_session,
284284
operate=OperationEnum.FILTER_TERMS,
285285
record_id=self.record.id, local_operation=True)
286-
286+
287287
self.chat_question.terminologies, term_list = get_terminology_template(_session, self.chat_question.question,
288288
calculate_oid, calculate_ds_id)
289289
self.current_logs[OperationEnum.FILTER_TERMS] = end_log(session=_session,
@@ -303,7 +303,8 @@ def filter_custom_prompts(self, _session: Session, custom_prompt_type: CustomPro
303303
operate=OperationEnum.FILTER_CUSTOM_PROMPT,
304304
record_id=self.record.id,
305305
local_operation=True)
306-
self.chat_question.custom_prompt, prompt_list = find_custom_prompts(_session, custom_prompt_type, calculate_oid,
306+
self.chat_question.custom_prompt, prompt_list = find_custom_prompts(_session, custom_prompt_type,
307+
calculate_oid,
307308
calculate_ds_id)
308309
self.current_logs[OperationEnum.FILTER_CUSTOM_PROMPT] = end_log(session=_session,
309310
log=self.current_logs[
@@ -323,11 +324,13 @@ def filter_training_template(self, _session: Session, oid: int = None, ds_id: in
323324
calculate_ds_id = None
324325
if self.current_assistant and self.current_assistant.type == 1:
325326
self.chat_question.data_training, example_list = get_training_template(_session,
326-
self.chat_question.question, calculate_oid,
327+
self.chat_question.question,
328+
calculate_oid,
327329
None, self.current_assistant.id)
328330
else:
329331
self.chat_question.data_training, example_list = get_training_template(_session,
330-
self.chat_question.question, calculate_oid,
332+
self.chat_question.question,
333+
calculate_oid,
331334
calculate_ds_id)
332335
self.current_logs[OperationEnum.FILTER_SQL_EXAMPLE] = end_log(session=_session,
333336
log=self.current_logs[
@@ -833,12 +836,15 @@ def generate_chart(self, _session: Session, chart_type: Optional[str] = '', sche
833836
reasoning_content=full_thinking_text,
834837
token_usage=token_usage)
835838

836-
@staticmethod
837-
def check_sql(res: str) -> tuple[str, Optional[list]]:
839+
def check_sql(self, session: Session, res: str, operate: OperationEnum) -> tuple[str, Optional[list]]:
838840
json_str = extract_nested_json(res)
841+
842+
log = self.current_logs[operate]
843+
839844
if json_str is None:
840-
raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer',
841-
'traceback': "Cannot parse sql from answer:\n" + res}).decode())
845+
trigger_log_error(session, log)
846+
raise SingleMessageError(orjson.dumps({'message': 'SQL answer is not a valid json object',
847+
'traceback': "SQL answer is not a valid json object:\n" + res}).decode())
842848
sql: str
843849
data: dict
844850
try:
@@ -850,12 +856,15 @@ def check_sql(res: str) -> tuple[str, Optional[list]]:
850856
message = data['message']
851857
raise SingleMessageError(message)
852858
except SingleMessageError as e:
859+
trigger_log_error(session, log)
853860
raise e
854861
except Exception:
862+
trigger_log_error(session, log)
855863
raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer',
856864
'traceback': "Cannot parse sql from answer:\n" + res}).decode())
857865

858866
if sql.strip() == '':
867+
trigger_log_error(session, log)
859868
raise SingleMessageError("SQL query is empty")
860869
return sql, data.get('tables')
861870

@@ -899,8 +908,8 @@ def get_brief_from_sql_answer(res: str) -> Optional[str]:
899908

900909
return brief
901910

902-
def check_save_sql(self, session: Session, res: str) -> str:
903-
sql, *_ = self.check_sql(res=res)
911+
def check_save_sql(self, session: Session, res: str, operate: OperationEnum) -> str:
912+
sql, *_ = self.check_sql(session=session, res=res, operate=operate)
904913
save_sql(session=session, sql=sql, record_id=self.record.id)
905914

906915
self.chat_question.sql = sql
@@ -1149,7 +1158,9 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
11491158
sqlbot_temp_sql_text = None
11501159
assistant_dynamic_sql = None
11511160
# row permission
1152-
sql, tables = self.check_sql(res=full_sql_text)
1161+
1162+
sql_operate = OperationEnum.GENERATE_SQL
1163+
sql, tables = self.check_sql(session=_session, res=full_sql_text, operate=sql_operate)
11531164
if ((not self.current_assistant or is_page_embedded) and is_normal_user(
11541165
self.current_user)) or use_dynamic_ds:
11551166
sql_result = None
@@ -1158,19 +1169,21 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
11581169
dynamic_sql_result = self.generate_assistant_dynamic_sql(_session, sql, tables)
11591170
sqlbot_temp_sql_text = dynamic_sql_result.get(
11601171
'sqlbot_temp_sql_text') if dynamic_sql_result else None
1161-
# sql_result = self.generate_assistant_filter(sql, tables)
11621172
else:
11631173
sql_result = self.generate_filter(_session, sql, tables) # maybe no sql and tables
11641174

11651175
if sql_result:
11661176
SQLBotLogUtil.info(sql_result)
1167-
sql = self.check_save_sql(session=_session, res=sql_result)
1177+
sql_operate = OperationEnum.GENERATE_SQL_WITH_PERMISSIONS
1178+
sql = self.check_save_sql(session=_session, res=sql_result, operate=sql_operate)
11681179
elif dynamic_sql_result and sqlbot_temp_sql_text:
1169-
assistant_dynamic_sql = self.check_save_sql(session=_session, res=sqlbot_temp_sql_text)
1180+
sql_operate = OperationEnum.GENERATE_DYNAMIC_SQL
1181+
assistant_dynamic_sql = self.check_save_sql(session=_session, res=sqlbot_temp_sql_text,
1182+
operate=sql_operate)
11701183
else:
1171-
sql = self.check_save_sql(session=_session, res=full_sql_text)
1184+
sql = self.check_save_sql(session=_session, res=full_sql_text, operate=sql_operate)
11721185
else:
1173-
sql = self.check_save_sql(session=_session, res=full_sql_text)
1186+
sql = self.check_save_sql(session=_session, res=full_sql_text, operate=sql_operate)
11741187

11751188
SQLBotLogUtil.info('sql: ' + sql)
11761189

0 commit comments

Comments
 (0)