diff --git a/backend/alembic/versions/037_create_chat_log.py b/backend/alembic/versions/037_create_chat_log.py new file mode 100644 index 00000000..59a7f425 --- /dev/null +++ b/backend/alembic/versions/037_create_chat_log.py @@ -0,0 +1,131 @@ +"""035_create_chat_log + +Revision ID: 68a06302cf70 +Revises: 29559ee607af +Create Date: 2025-08-18 16:02:43.353110 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '68a06302cf70' +down_revision = '646e7ca28e0e' +branch_labels = None +depends_on = None + + +sql=''' +CREATE OR REPLACE FUNCTION safe_jsonb_cast(text) RETURNS jsonb AS +$$ +BEGIN + RETURN $1::jsonb; +EXCEPTION + WHEN others THEN + RETURN to_json($1::text)::jsonb; +END; +$$ LANGUAGE plpgsql; + +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '0', + id, + ai_modal_id, + safe_jsonb_cast(full_sql_message), + create_time, + finish_time, + safe_jsonb_cast(token_sql), + safe_jsonb_cast(sql_answer)->>'reasoning_content' +FROM chat_record +WHERE full_sql_message IS NOT NULL; +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '1', + id, + ai_modal_id, + safe_jsonb_cast(full_chart_message), + create_time, + finish_time, + safe_jsonb_cast(token_chart), + safe_jsonb_cast(chart_answer)->>'reasoning_content' +FROM chat_record +WHERE full_chart_message IS NOT NULL; +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '2', + id, + ai_modal_id, + safe_jsonb_cast(full_analysis_message), + create_time, + finish_time, + safe_jsonb_cast(token_analysis), + safe_jsonb_cast(analysis)->>'reasoning_content' +FROM chat_record +WHERE full_analysis_message IS NOT NULL; +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '3', + id, + ai_modal_id, + safe_jsonb_cast(full_predict_message), + create_time, + finish_time, + safe_jsonb_cast(token_predict), + safe_jsonb_cast(predict)->>'reasoning_content' +FROM chat_record +WHERE full_predict_message IS NOT NULL; +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '4', + id, + ai_modal_id, + safe_jsonb_cast(full_recommended_question_message), + create_time, + finish_time, + safe_jsonb_cast(token_recommended_question), + safe_jsonb_cast(recommended_question_answer)->>'reasoning_content' +FROM chat_record +WHERE full_recommended_question_message IS NOT NULL; +INSERT INTO chat_log(type, operate, pid, ai_modal_id, messages, start_time, finish_time, token_usage, reasoning_content) +SELECT '0', + '6', + id, + ai_modal_id, + safe_jsonb_cast(full_select_datasource_message), + create_time, + finish_time, + safe_jsonb_cast(token_select_datasource_question), + safe_jsonb_cast(datasource_select_answer)->>'reasoning_content' +FROM chat_record +WHERE full_select_datasource_message IS NOT NULL; + +''' + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('chat_log', + sa.Column('id', sa.BigInteger(), sa.Identity(always=True), nullable=False), + sa.Column('type', sa.Enum('0', name='typeenum', native_enum=False, length=3), nullable=True), + sa.Column('operate', sa.Enum('0', '1', '2', '3', '4', '5', '6', name='operationenum', native_enum=False, length=3), nullable=True), + sa.Column('pid', sa.BigInteger(), nullable=True), + sa.Column('ai_modal_id', sa.BigInteger(), nullable=True), + sa.Column('base_modal', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True), + sa.Column('messages', postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column('reasoning_content', sa.Text(), nullable=True), + sa.Column('start_time', sa.DateTime(), nullable=True), + sa.Column('finish_time', sa.DateTime(), nullable=True), + sa.Column('token_usage', postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + + op.execute(sql) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + op.drop_table('chat_log') + # ### end Alembic commands ### diff --git a/backend/alembic/versions/038_remove_chat_record_cloumns.py b/backend/alembic/versions/038_remove_chat_record_cloumns.py new file mode 100644 index 00000000..870c5477 --- /dev/null +++ b/backend/alembic/versions/038_remove_chat_record_cloumns.py @@ -0,0 +1,51 @@ +"""038_remove_chat_record_cloumns + +Revision ID: fc23c4f3e755 +Revises: 68a06302cf70 +Create Date: 2025-08-21 14:34:59.149410 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'fc23c4f3e755' +down_revision = '68a06302cf70' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('chat_record', 'token_predict') + op.drop_column('chat_record', 'token_select_datasource_question') + op.drop_column('chat_record', 'token_sql') + op.drop_column('chat_record', 'full_analysis_message') + op.drop_column('chat_record', 'full_recommended_question_message') + op.drop_column('chat_record', 'token_chart') + op.drop_column('chat_record', 'full_predict_message') + op.drop_column('chat_record', 'full_chart_message') + op.drop_column('chat_record', 'full_sql_message') + op.drop_column('chat_record', 'full_select_datasource_message') + op.drop_column('chat_record', 'token_recommended_question') + op.drop_column('chat_record', 'token_analysis') + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('chat_record', sa.Column('token_analysis', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('token_recommended_question', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_select_datasource_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_sql_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_chart_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_predict_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('token_chart', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_recommended_question_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('full_analysis_message', sa.TEXT(), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('token_sql', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('token_select_datasource_question', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + op.add_column('chat_record', sa.Column('token_predict', sa.VARCHAR(length=256), autoincrement=False, nullable=True)) + # ### end Alembic commands ### diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 0d48669d..5ffbd940 100644 --- a/backend/apps/chat/curd/chat.py +++ b/backend/apps/chat/curd/chat.py @@ -4,8 +4,10 @@ from typing import List import orjson import sqlparse from sqlalchemy import and_, select, update +from sqlalchemy.orm import aliased -from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion +from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion, ChatLog, \ + TypeEnum, OperationEnum, ChatRecordResult from apps.datasource.models.datasource import CoreDatasource from apps.system.crud.assistant import AssistantOutDsFactory from common.core.deps import CurrentAssistant, SessionDep, CurrentUser @@ -95,7 +97,10 @@ def get_chat_with_records_with_data(session: SessionDep, chart_id: int, current_ current_assistant: CurrentAssistant) -> ChatInfo: return get_chat_with_records(session, chart_id, current_user, current_assistant, True) + dynamic_ds_types = [1, 3] + + def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser, current_assistant: CurrentAssistant, with_data: bool = False) -> ChatInfo: chat = session.get(Chat, chart_id) @@ -118,13 +123,36 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr chat_info.datasource_name = ds.name chat_info.ds_type = ds.type - stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time, - ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, - ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict, - ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id, - ChatRecord.recommended_question, ChatRecord.first_chat, - ChatRecord.finish, ChatRecord.error).where( - and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time) + sql_alias_log = aliased(ChatLog) + chart_alias_log = aliased(ChatLog) + analysis_alias_log = aliased(ChatLog) + predict_alias_log = aliased(ChatLog) + + stmt = (select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time, + ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, + ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict, + ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id, + ChatRecord.recommended_question, ChatRecord.first_chat, + ChatRecord.finish, ChatRecord.error, + sql_alias_log.reasoning_content.label('sql_reasoning_content'), + chart_alias_log.reasoning_content.label('chart_reasoning_content'), + analysis_alias_log.reasoning_content.label('analysis_reasoning_content'), + predict_alias_log.reasoning_content.label('predict_reasoning_content') + ) + .outerjoin(sql_alias_log, and_(sql_alias_log.pid == ChatRecord.id, + sql_alias_log.type == TypeEnum.CHAT, + sql_alias_log.operate == OperationEnum.GENERATE_SQL)) + .outerjoin(chart_alias_log, and_(chart_alias_log.pid == ChatRecord.id, + chart_alias_log.type == TypeEnum.CHAT, + chart_alias_log.operate == OperationEnum.GENERATE_CHART)) + .outerjoin(analysis_alias_log, and_(analysis_alias_log.pid == ChatRecord.id, + analysis_alias_log.type == TypeEnum.CHAT, + analysis_alias_log.operate == OperationEnum.ANALYSIS)) + .outerjoin(predict_alias_log, and_(predict_alias_log.pid == ChatRecord.id, + predict_alias_log.type == TypeEnum.CHAT, + predict_alias_log.operate == OperationEnum.PREDICT_DATA)) + .where(and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by( + ChatRecord.create_time)) if with_data: stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time, ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, @@ -136,28 +164,35 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr ChatRecord.create_time) result = session.execute(stmt).all() - record_list: list[ChatRecord] = [] + record_list: list[ChatRecordResult] = [] for row in result: if not with_data: record_list.append( - ChatRecord(id=row.id, chat_id=row.chat_id, create_time=row.create_time, finish_time=row.finish_time, - question=row.question, sql_answer=row.sql_answer, sql=row.sql, - chart_answer=row.chart_answer, chart=row.chart, - analysis=row.analysis, predict=row.predict, - datasource_select_answer=row.datasource_select_answer, - analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id, - recommended_question=row.recommended_question, first_chat=row.first_chat, - finish=row.finish, error=row.error)) + ChatRecordResult(id=row.id, chat_id=row.chat_id, create_time=row.create_time, + finish_time=row.finish_time, + question=row.question, sql_answer=row.sql_answer, sql=row.sql, + chart_answer=row.chart_answer, chart=row.chart, + analysis=row.analysis, predict=row.predict, + datasource_select_answer=row.datasource_select_answer, + analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id, + recommended_question=row.recommended_question, first_chat=row.first_chat, + finish=row.finish, error=row.error, + sql_reasoning_content=row.sql_reasoning_content, + chart_reasoning_content=row.chart_reasoning_content, + analysis_reasoning_content=row.analysis_reasoning_content, + predict_reasoning_content=row.predict_reasoning_content, + )) else: record_list.append( - ChatRecord(id=row.id, chat_id=row.chat_id, create_time=row.create_time, finish_time=row.finish_time, - question=row.question, sql_answer=row.sql_answer, sql=row.sql, - chart_answer=row.chart_answer, chart=row.chart, - analysis=row.analysis, predict=row.predict, - datasource_select_answer=row.datasource_select_answer, - analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id, - recommended_question=row.recommended_question, first_chat=row.first_chat, - finish=row.finish, error=row.error, data=row.data, predict_data=row.predict_data)) + ChatRecordResult(id=row.id, chat_id=row.chat_id, create_time=row.create_time, + finish_time=row.finish_time, + question=row.question, sql_answer=row.sql_answer, sql=row.sql, + chart_answer=row.chart_answer, chart=row.chart, + analysis=row.analysis, predict=row.predict, + datasource_select_answer=row.datasource_select_answer, + analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id, + recommended_question=row.recommended_question, first_chat=row.first_chat, + finish=row.finish, error=row.error, data=row.data, predict_data=row.predict_data)) result = list(map(format_record, record_list)) @@ -166,27 +201,35 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr return chat_info -def format_record(record: ChatRecord): +def format_record(record: ChatRecordResult): _dict = record.model_dump() if record.sql_answer and record.sql_answer.strip() != '' and record.sql_answer.strip()[0] == '{' and \ record.sql_answer.strip()[-1] == '}': _obj = orjson.loads(record.sql_answer) _dict['sql_answer'] = _obj.get('reasoning_content') + if record.sql_reasoning_content and record.sql_reasoning_content.strip() != '': + _dict['sql_answer'] = record.sql_reasoning_content if record.chart_answer and record.chart_answer.strip() != '' and record.chart_answer.strip()[0] == '{' and \ record.chart_answer.strip()[-1] == '}': _obj = orjson.loads(record.chart_answer) _dict['chart_answer'] = _obj.get('reasoning_content') + if record.chart_reasoning_content and record.chart_reasoning_content.strip() != '': + _dict['chart_answer'] = record.chart_reasoning_content if record.analysis and record.analysis.strip() != '' and record.analysis.strip()[0] == '{' and \ record.analysis.strip()[-1] == '}': _obj = orjson.loads(record.analysis) _dict['analysis_thinking'] = _obj.get('reasoning_content') _dict['analysis'] = _obj.get('content') + if record.analysis_reasoning_content and record.analysis_reasoning_content.strip() != '': + _dict['analysis_thinking'] = record.analysis_reasoning_content if record.predict and record.predict.strip() != '' and record.predict.strip()[0] == '{' and record.predict.strip()[ -1] == '}': _obj = orjson.loads(record.predict) _dict['predict'] = _obj.get('reasoning_content') _dict['predict_content'] = _obj.get('content') + if record.predict_reasoning_content and record.predict_reasoning_content.strip() != '': + _dict['predict'] = record.predict_reasoning_content if record.data and record.data.strip() != '': try: _obj = orjson.loads(record.data) @@ -208,19 +251,30 @@ def format_record(record: ChatRecord): return _dict -def list_base_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> List[ChatRecord]: - stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.full_sql_message, ChatRecord.full_chart_message, - ChatRecord.first_chat, ChatRecord.create_time).where( - and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id, - ChatRecord.analysis_record_id.is_(None), ChatRecord.predict_record_id.is_(None))).order_by( - ChatRecord.create_time) +def list_generate_sql_logs(session: SessionDep, chart_id: int) -> List[ChatLog]: + stmt = select(ChatLog).where( + and_(ChatLog.pid.in_(select(ChatRecord.id).where(and_(ChatRecord.chat_id == chart_id))), + ChatLog.type == TypeEnum.CHAT, ChatLog.operate == OperationEnum.GENERATE_SQL)).order_by( + ChatLog.start_time) result = session.execute(stmt).all() - record_list: List[ChatRecord] = [] - for r in result: - record_list.append( - ChatRecord(id=r.id, chat_id=r.chat_id, create_time=r.create_time, full_sql_message=r.full_sql_message, - full_chart_message=r.full_chart_message, first_chat=r.first_chat)) - return record_list + _list = [] + for row in result: + for r in row: + _list.append(ChatLog(**r.model_dump())) + return _list + + +def list_generate_chart_logs(session: SessionDep, chart_id: int) -> List[ChatLog]: + stmt = select(ChatLog).where( + and_(ChatLog.pid.in_(select(ChatRecord.id).where(and_(ChatRecord.chat_id == chart_id))), + ChatLog.type == TypeEnum.CHAT, ChatLog.operate == OperationEnum.GENERATE_CHART)).order_by( + ChatLog.start_time) + result = session.execute(stmt).all() + _list = [] + for row in result: + for r in row: + _list.append(ChatLog(**r.model_dump())) + return _list def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat, @@ -344,125 +398,117 @@ def save_analysis_predict_record(session: SessionDep, base_record: ChatRecord, a return result -def save_full_sql_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord: - return save_full_sql_message_and_answer(session=session, record_id=record_id, full_message=full_message, answer='') +def start_log(session: SessionDep, ai_modal_id: int, ai_modal_name: str, operate: OperationEnum, record_id: int, + full_message: list[dict]) -> ChatLog: + log = ChatLog(type=TypeEnum.CHAT, operate=operate, pid=record_id, ai_modal_id=ai_modal_id, base_modal=ai_modal_name, + messages=full_message, start_time=datetime.datetime.now()) + + result = ChatLog(**log.model_dump()) + + session.add(log) + session.flush() + session.refresh(log) + result.id = log.id + session.commit() + + return result -def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer: str, full_message: str, - token_usage: dict = None) -> ChatRecord: +def end_log(session: SessionDep, log: ChatLog, full_message: list[dict], reasoning_content: str = None, + token_usage=None) -> ChatLog: + if token_usage is None: + token_usage = {} + log.messages = full_message + log.token_usage = token_usage + log.finish_time = datetime.datetime.now() + log.reasoning_content = reasoning_content if reasoning_content and len(reasoning_content.strip()) > 0 else None + + stmt = update(ChatLog).where(and_(ChatLog.id == log.id)).values( + messages=log.messages, + token_usage=log.token_usage, + finish_time=log.finish_time, + reasoning_content=log.reasoning_content + ) + session.execute(stmt) + session.commit() + + return log + + +def save_sql_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord: if not record_id: raise Exception("Record id cannot be None") - record = get_chat_record_by_id(session, record_id) - record.full_sql_message = full_message - record.sql_answer = answer - - if token_usage: - record.token_sql = orjson.dumps(token_usage).decode() - - result = ChatRecord(**record.model_dump()) - - stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_sql_message=record.full_sql_message, - sql_answer=record.sql_answer, - token_sql=record.token_sql, + stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values( + sql_answer=answer, ) session.execute(stmt) session.commit() - return result - - -def save_full_analysis_message_and_answer(session: SessionDep, record_id: int, answer: str, - full_message: str, token_usage: dict = None) -> ChatRecord: - if not record_id: - raise Exception("Record id cannot be None") record = get_chat_record_by_id(session, record_id) - record.full_analysis_message = full_message - record.analysis = answer + return record - if token_usage: - record.token_analysis = orjson.dumps(token_usage).decode() - result = ChatRecord(**record.model_dump()) +def save_analysis_answer(session: SessionDep, record_id: int, answer: str = '') -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") - stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_analysis_message=record.full_analysis_message, - analysis=record.analysis, - token_analysis=record.token_analysis, + stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values( + analysis=answer, ) session.execute(stmt) session.commit() - return result - - -def save_full_predict_message_and_answer(session: SessionDep, record_id: int, answer: str, - full_message: str, data: str, token_usage: dict = None) -> ChatRecord: - if not record_id: - raise Exception("Record id cannot be None") record = get_chat_record_by_id(session, record_id) - record.full_predict_message = full_message - record.predict = answer - record.predict_data = data + return record - if token_usage: - record.token_predict = orjson.dumps(token_usage).decode() - result = ChatRecord(**record.model_dump()) +def save_predict_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") - stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_predict_message=record.full_predict_message, - predict=record.predict, - predict_data=record.predict_data, - token_predict=record.token_predict + stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values( + predict=answer, ) session.execute(stmt) session.commit() - return result + record = get_chat_record_by_id(session, record_id) + + return record -def save_full_select_datasource_message_and_answer(session: SessionDep, record_id: int, answer: str, - full_message: str, datasource: int = None, - engine_type: str = None, token_usage: dict = None) -> ChatRecord: +def save_select_datasource_answer(session: SessionDep, record_id: int, answer: str, + datasource: int = None, engine_type: str = None) -> ChatRecord: if not record_id: raise Exception("Record id cannot be None") record = get_chat_record_by_id(session, record_id) - record.full_select_datasource_message = full_message record.datasource_select_answer = answer if datasource: record.datasource = datasource record.engine_type = engine_type - if token_usage: - record.token_select_datasource_question = orjson.dumps(token_usage).decode() - result = ChatRecord(**record.model_dump()) if datasource: stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_select_datasource_message=record.full_select_datasource_message, datasource_select_answer=record.datasource_select_answer, - token_select_datasource_question=record.token_select_datasource_question, datasource=record.datasource, engine_type=record.engine_type, ) else: stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_select_datasource_message=record.full_select_datasource_message, datasource_select_answer=record.datasource_select_answer, - token_select_datasource_question=record.token_select_datasource_question, ) session.execute(stmt) @@ -472,16 +518,12 @@ def save_full_select_datasource_message_and_answer(session: SessionDep, record_i return result -def save_full_recommend_question_message_and_answer(session: SessionDep, record_id: int, answer: dict = None, - full_message: str = '[]', token_usage: dict = None) -> ChatRecord: +def save_recommend_question_answer(session: SessionDep, record_id: int, + answer: dict = None) -> ChatRecord: if not record_id: raise Exception("Record id cannot be None") - record = get_chat_record_by_id(session, record_id) - record.full_recommended_question_message = full_message - - if answer: - record.recommended_question_answer = orjson.dumps(answer).decode() + recommended_question_answer = orjson.dumps(answer).decode() json_str = '[]' if answer and answer.get('content') and answer.get('content') != '': @@ -492,25 +534,20 @@ def save_full_recommend_question_message_and_answer(session: SessionDep, record_ json_str = '[]' except Exception as e: pass - record.recommended_question = json_str + recommended_question = json_str - if token_usage: - record.token_recommended_question = orjson.dumps(token_usage).decode() - - result = ChatRecord(**record.model_dump()) - - stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_recommended_question_message=record.full_recommended_question_message, - recommended_question_answer=record.recommended_question_answer, - recommended_question=record.recommended_question, - token_recommended_question=record.token_recommended_question + stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values( + recommended_question_answer=recommended_question_answer, + recommended_question=recommended_question, ) session.execute(stmt) session.commit() - return result + record = get_chat_record_by_id(session, record_id) + + return record def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord: @@ -534,36 +571,21 @@ def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord: return result -def save_full_chart_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord: - return save_full_chart_message_and_answer(session=session, record_id=record_id, full_message=full_message, - answer='') - - -def save_full_chart_message_and_answer(session: SessionDep, record_id: int, answer: str, - full_message: str, token_usage: dict = None) -> ChatRecord: +def save_chart_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord: if not record_id: raise Exception("Record id cannot be None") - record = get_chat_record_by_id(session, record_id) - record.full_chart_message = full_message - record.chart_answer = answer - - if token_usage: - record.token_chart = orjson.dumps(token_usage).decode() - - result = ChatRecord(**record.model_dump()) - - stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values( - full_chart_message=record.full_chart_message, - chart_answer=record.chart_answer, - token_chart=record.token_chart + stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values( + chart_answer=answer, ) session.execute(stmt) session.commit() - return result + record = get_chat_record_by_id(session, record_id) + + return record def save_chart(session: SessionDep, record_id: int, chart: str) -> ChatRecord: diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 07dae5e2..d1cc6670 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -1,9 +1,12 @@ from datetime import datetime +from enum import Enum from typing import List, Optional from fastapi import Body from pydantic import BaseModel from sqlalchemy import Column, Integer, Text, BigInteger, DateTime, Identity, Boolean +from sqlalchemy import Enum as SQLAlchemyEnum +from sqlalchemy.dialects.postgresql import JSONB from sqlmodel import SQLModel, Field from apps.template.filter.generator import get_permissions_template @@ -16,6 +19,47 @@ from apps.template.generate_sql.generator import get_sql_template from apps.template.select_datasource.generator import get_datasource_template +def enum_values(enum_class: type[Enum]) -> list: + """Get values for enum.""" + return [status.value for status in enum_class] + + +class TypeEnum(Enum): + CHAT = "0" + + +# TODO other usage + +class OperationEnum(Enum): + GENERATE_SQL = '0' + GENERATE_CHART = '1' + ANALYSIS = '2' + PREDICT_DATA = '3' + GENERATE_RECOMMENDED_QUESTIONS = '4' + GENERATE_SQL_WITH_PERMISSIONS = '5' + CHOOSE_DATASOURCE = '6' + GENERATE_DYNAMIC_SQL = '7' + + +# TODO choose table / check connection / generate description + +class ChatLog(SQLModel, table=True): + __tablename__ = "chat_log" + id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True)) + type: TypeEnum = Field( + sa_column=Column(SQLAlchemyEnum(TypeEnum, native_enum=False, values_callable=enum_values, length=3))) + operate: OperationEnum = Field( + sa_column=Column(SQLAlchemyEnum(OperationEnum, native_enum=False, values_callable=enum_values, length=3))) + pid: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True)) + ai_modal_id: Optional[int] = Field(sa_column=Column(BigInteger)) + base_modal: Optional[str] = Field(max_length=255) + messages: Optional[list[dict]] = Field(sa_column=Column(JSONB)) + reasoning_content: Optional[str | None] = Field(sa_column=Column(Text, nullable=True)) + start_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) + finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) + token_usage: Optional[dict | None | int] = Field(sa_column=Column(JSONB)) + + class Chat(SQLModel, table=True): __tablename__ = "chat" id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True)) @@ -54,24 +98,40 @@ class ChatRecord(SQLModel, table=True): recommended_question_answer: str = Field(sa_column=Column(Text, nullable=True)) recommended_question: str = Field(sa_column=Column(Text, nullable=True)) datasource_select_answer: str = Field(sa_column=Column(Text, nullable=True)) - full_sql_message: str = Field(sa_column=Column(Text, nullable=True)) - token_sql: str = Field(max_length=256, nullable=True) - full_chart_message: str = Field(sa_column=Column(Text, nullable=True)) - token_chart: str = Field(max_length=256, nullable=True) - full_analysis_message: str = Field(sa_column=Column(Text, nullable=True)) - token_analysis: str = Field(max_length=256, nullable=True) - full_predict_message: str = Field(sa_column=Column(Text, nullable=True)) - token_predict: str = Field(max_length=256, nullable=True) - full_recommended_question_message: str = Field(sa_column=Column(Text, nullable=True)) - token_recommended_question: str = Field(max_length=256, nullable=True) - full_select_datasource_message: str = Field(sa_column=Column(Text, nullable=True)) - token_select_datasource_question: str = Field(max_length=256, nullable=True) finish: bool = Field(sa_column=Column(Boolean, nullable=True, default=False)) error: str = Field(sa_column=Column(Text, nullable=True)) analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) +class ChatRecordResult(BaseModel): + id: Optional[int] = None + chat_id: Optional[int] = None + ai_modal_id: Optional[int] = None + first_chat: bool = False + create_time: Optional[datetime] = None + finish_time: Optional[datetime] = None + question: Optional[str] = None + sql_answer: Optional[str] = None + sql: Optional[str] = None + data: Optional[str] = None + chart_answer: Optional[str] = None + chart: Optional[str] = None + analysis: Optional[str] = None + predict: Optional[str] = None + predict_data: Optional[str] = None + recommended_question: Optional[str] = None + datasource_select_answer: Optional[str] = None + finish: Optional[bool] = None + error: Optional[str] = None + analysis_record_id: Optional[int] = None + predict_record_id: Optional[int] = None + sql_reasoning_content: Optional[str] = None + chart_reasoning_content: Optional[str] = None + analysis_reasoning_content: Optional[str] = None + predict_reasoning_content: Optional[str] = None + + class CreateChat(BaseModel): id: int = None question: str = None @@ -100,6 +160,7 @@ class ChatInfo(BaseModel): class AiModelQuestion(BaseModel): ai_modal_id: int = None + ai_modal_name: str = None # Specific model name engine: str = "" db_schema: str = "" sql: str = "" diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 10efb3e7..ec5e9d1d 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -18,13 +18,13 @@ from sqlalchemy.orm import sessionmaker from sqlmodel import create_engine, Session from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config -from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \ - save_error_message, save_sql_exec_data, save_full_chart_message, save_full_chart_message_and_answer, save_chart, \ - finish_record, save_full_analysis_message_and_answer, save_full_predict_message_and_answer, save_predict_data, \ - save_full_select_datasource_message_and_answer, save_full_recommend_question_message_and_answer, \ - get_old_questions, save_analysis_predict_record, list_base_records, rename_chat, get_chart_config, \ - get_chat_chart_data -from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat +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, \ + 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 +from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum from apps.datasource.crud.datasource import get_table_schema from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user from apps.datasource.models.datasource import CoreDatasource @@ -38,12 +38,13 @@ from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_o warnings.filterwarnings("ignore") -base_message_count_limit = 5 +base_message_count_limit = 6 executor = ThreadPoolExecutor(max_workers=200) dynamic_ds_types = [1, 3] + class LLMService: ds: CoreDatasource chat_question: ChatQuestion @@ -52,13 +53,18 @@ class LLMService: llm: BaseChatModel sql_message: List[Union[BaseMessage, dict[str, Any]]] = [] chart_message: List[Union[BaseMessage, dict[str, Any]]] = [] - history_records: List[ChatRecord] = [] + session: 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 @@ -94,12 +100,10 @@ class LLMService: chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL' chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds) - history_records: List[ChatRecord] = list( - map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False, - list_base_records(session=self.session, - current_user=current_user, - chart_id=chat_id)))) - self.change_title = len(history_records) == 0 + self.generate_sql_logs = list_generate_sql_logs(session=self.session, chart_id=chat_id) + self.generate_chart_logs = list_generate_chart_logs(session=self.session, chart_id=chat_id) + + self.change_title = len(self.generate_sql_logs) == 0 chat_question.lang = get_lang_name(current_user.language) @@ -114,13 +118,12 @@ class LLMService: 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 - self.history_records = history_records - self.init_messages() def is_running(self, timeout=0.5): @@ -134,21 +137,8 @@ class LLMService: return True def init_messages(self): - # self.agent_executor = create_react_agent(self.llm) - last_sql_messages = list( - filter(lambda r: True if r.full_sql_message is not None and r.full_sql_message.strip() != '' else False, - self.history_records)) - last_sql_message_str = "[]" if last_sql_messages is None or len(last_sql_messages) == 0 else last_sql_messages[ - -1].full_sql_message - - last_chart_messages = list( - filter( - lambda r: True if r.full_chart_message is not None and r.full_chart_message.strip() != '' else False, - self.history_records)) - last_chart_message_str = "[]" if last_chart_messages is None or len(last_chart_messages) == 0 else \ - last_chart_messages[-1].full_chart_message - - last_sql_messages: List[dict[str, Any]] = orjson.loads(last_sql_message_str) + last_sql_messages: List[dict[str, Any]] = self.generate_sql_logs[-1].messages if len( + self.generate_sql_logs) > 0 else [] # todo maybe can configure count_limit = 0 - base_message_count_limit @@ -167,7 +157,8 @@ class LLMService: _msg = AIMessage(content=last_sql_message['content']) self.sql_message.append(_msg) - last_chart_messages: List[dict[str, Any]] = orjson.loads(last_chart_message_str) + last_chart_messages: List[dict[str, Any]] = self.generate_chart_logs[-1].messages if len( + self.generate_chart_logs) > 0 else [] self.chart_message = [] # add sys prompt @@ -177,11 +168,11 @@ class LLMService: # limit count for last_chart_message in last_chart_messages: _msg: BaseMessage - if last_chart_message['type'] == 'human': - _msg = HumanMessage(content=last_chart_message['content']) + if last_chart_message.get('type') == 'human': + _msg = HumanMessage(content=last_chart_message.get('content')) self.chart_message.append(_msg) - elif last_chart_message['type'] == 'ai': - _msg = AIMessage(content=last_chart_message['content']) + elif last_chart_message.get('type') == 'ai': + _msg = AIMessage(content=last_chart_message.get('content')) self.chart_message.append(_msg) def init_record(self) -> ChatRecord: @@ -222,11 +213,16 @@ class LLMService: analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question())) analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question())) - self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='', - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} for msg - in - analysis_msg]).decode()) + self.current_logs[OperationEnum.ANALYSIS] = start_log(session=self.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, + 'content': msg.content} for + msg + in analysis_msg]) full_thinking_text = '' full_analysis_text = '' res = self.llm.stream(analysis_msg) @@ -247,14 +243,18 @@ class LLMService: get_token_usage(chunk, token_usage) analysis_msg.append(AIMessage(full_analysis_text)) - self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, - token_usage=token_usage, - answer=orjson.dumps({'content': full_analysis_text, - 'reasoning_content': full_thinking_text}).decode(), - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} for msg - in - analysis_msg]).decode()) + + self.current_logs[OperationEnum.ANALYSIS] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.ANALYSIS], + full_message=[ + {'type': msg.type, + 'content': msg.content} + for msg in analysis_msg], + reasoning_content=full_thinking_text, + token_usage=token_usage) + self.record = save_analysis_answer(session=self.session, record_id=self.record.id, + answer=orjson.dumps({'content': full_analysis_text}).decode()) def generate_predict(self): fields = self.get_fields_from_chart() @@ -265,12 +265,16 @@ class LLMService: predict_msg.append(SystemMessage(content=self.chat_question.predict_sys_question())) predict_msg.append(HumanMessage(content=self.chat_question.predict_user_question())) - self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id, answer='', - data='', - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} for msg - in - predict_msg]).decode()) + self.current_logs[OperationEnum.PREDICT_DATA] = start_log(session=self.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, + 'content': msg.content} for + msg + in predict_msg]) full_thinking_text = '' full_predict_text = '' res = self.llm.stream(predict_msg) @@ -291,15 +295,17 @@ class LLMService: get_token_usage(chunk, token_usage) predict_msg.append(AIMessage(full_predict_text)) - self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id, - token_usage=token_usage, - answer=orjson.dumps({'content': full_predict_text, - 'reasoning_content': full_thinking_text}).decode(), - data='', - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} for msg - in - predict_msg]).decode()) + self.record = save_predict_answer(session=self.session, record_id=self.record.id, + answer=orjson.dumps({'content': full_predict_text}).decode()) + self.current_logs[OperationEnum.PREDICT_DATA] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.PREDICT_DATA], + full_message=[ + {'type': msg.type, + 'content': msg.content} + for msg in predict_msg], + reasoning_content=full_thinking_text, + token_usage=token_usage) def generate_recommend_questions_task(self): @@ -316,12 +322,16 @@ class LLMService: guess_msg.append( HumanMessage(content=self.chat_question.guess_user_question(orjson.dumps(old_questions).decode()))) - self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id, - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} - for msg - in - guess_msg]).decode()) + self.current_logs[OperationEnum.GENERATE_RECOMMENDED_QUESTIONS] = start_log(session=self.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, + 'content': msg.content} for + msg + in guess_msg]) full_thinking_text = '' full_guess_text = '' token_usage = {} @@ -342,15 +352,19 @@ class LLMService: get_token_usage(chunk, token_usage) guess_msg.append(AIMessage(full_guess_text)) - self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id, - token_usage=token_usage, - answer={'content': full_guess_text, - 'reasoning_content': full_thinking_text}, - full_message=orjson.dumps([{'type': msg.type, - 'content': msg.content} - for msg - in - guess_msg]).decode()) + + self.current_logs[OperationEnum.GENERATE_RECOMMENDED_QUESTIONS] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.GENERATE_RECOMMENDED_QUESTIONS], + full_message=[ + {'type': msg.type, + 'content': msg.content} + for msg in guess_msg], + reasoning_content=full_thinking_text, + token_usage=token_usage) + self.record = save_recommend_question_answer(session=self.session, record_id=self.record.id, + answer={'content': full_guess_text}) + yield {'recommended_question': self.record.recommended_question} def select_datasource(self): @@ -376,6 +390,9 @@ class LLMService: ignore_auto_select = _ds_list and len(_ds_list) == 1 # ignore auto select ds + full_thinking_text = '' + full_text = '' + if not ignore_auto_select: _ds_list_dict = [] for _ds in _ds_list: @@ -383,20 +400,15 @@ class LLMService: datasource_msg.append( HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode()))) - history_msg = [] - if self.record.full_select_datasource_message and self.record.full_select_datasource_message.strip() != '': - history_msg = orjson.loads(self.record.full_select_datasource_message) + self.current_logs[OperationEnum.CHOOSE_DATASOURCE] = start_log(session=self.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, + 'content': msg.content} for + msg in datasource_msg]) - self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id, - answer='', - full_message=orjson.dumps(history_msg + - [{'type': msg.type, - 'content': msg.content} - for msg - in - datasource_msg]).decode()) - full_thinking_text = '' - full_text = '' token_usage = {} res = self.llm.stream(datasource_msg) for chunk in res: @@ -415,6 +427,15 @@ class LLMService: get_token_usage(chunk, token_usage) datasource_msg.append(AIMessage(full_text)) + self.current_logs[OperationEnum.CHOOSE_DATASOURCE] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.CHOOSE_DATASOURCE], + full_message=[ + {'type': msg.type, 'content': msg.content} + for msg in datasource_msg], + reasoning_content=full_thinking_text, + token_usage=token_usage) + json_str = extract_nested_json(full_text) _error: Exception | None = None @@ -460,17 +481,10 @@ class LLMService: _error = e if not ignore_auto_select: - self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id, - answer=orjson.dumps({'content': full_text, - 'reasoning_content': full_thinking_text}).decode(), - datasource=_datasource, - engine_type=_engine_type, - full_message=orjson.dumps(history_msg + - [{'type': msg.type, - 'content': msg.content} - for msg - in - datasource_msg]).decode()) + self.record = save_select_datasource_answer(session=self.session, record_id=self.record.id, + answer=orjson.dumps({'content': full_text}).decode(), + datasource=_datasource, + engine_type=_engine_type) self.init_messages() if _error: @@ -479,10 +493,15 @@ class LLMService: def generate_sql(self): # append current question self.sql_message.append(HumanMessage(self.chat_question.sql_user_question())) - self.record = save_full_sql_message(session=self.session, record_id=self.record.id, - full_message=orjson.dumps( - [{'type': msg.type, 'content': msg.content} for msg in - self.sql_message]).decode()) + + self.current_logs[OperationEnum.GENERATE_SQL] = start_log(session=self.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, 'content': msg.content} for msg + in self.sql_message]) full_thinking_text = '' full_sql_text = '' token_usage = {} @@ -503,24 +522,37 @@ class LLMService: get_token_usage(chunk, token_usage) self.sql_message.append(AIMessage(full_sql_text)) - self.record = save_full_sql_message_and_answer(session=self.session, record_id=self.record.id, - token_usage=token_usage, - answer=orjson.dumps({'content': full_sql_text, - 'reasoning_content': full_thinking_text}).decode(), - full_message=orjson.dumps( - [{'type': msg.type, 'content': msg.content} for msg in - self.sql_message]).decode()) + + self.current_logs[OperationEnum.GENERATE_SQL] = end_log(session=self.session, + log=self.current_logs[OperationEnum.GENERATE_SQL], + full_message=[{'type': msg.type, 'content': msg.content} + for msg in self.sql_message], + reasoning_content=full_thinking_text, + token_usage=token_usage) + self.record = save_sql_answer(session=self.session, record_id=self.record.id, + answer=orjson.dumps({'content': full_sql_text}).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())) + dynamic_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = [] + dynamic_sql_msg.append(SystemMessage(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=self.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, + 'content': msg.content} + for + msg in dynamic_sql_msg]) + full_thinking_text = '' full_dynamic_text = '' - res = self.llm.stream(msg) + res = self.llm.stream(dynamic_sql_msg) token_usage = {} for chunk in res: SQLBotLogUtil.info(chunk) @@ -533,6 +565,18 @@ class LLMService: full_dynamic_text += chunk.content get_token_usage(chunk, token_usage) + dynamic_sql_msg.append(AIMessage(full_dynamic_text)) + + self.current_logs[OperationEnum.GENERATE_DYNAMIC_SQL] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.GENERATE_DYNAMIC_SQL], + full_message=[ + {'type': msg.type, + '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 @@ -550,23 +594,23 @@ class LLMService: filter = json.dumps(filters, ensure_ascii=False) self.chat_question.sql = sql self.chat_question.filter = filter - msg: List[Union[BaseMessage, dict[str, Any]]] = [] - msg.append(SystemMessage(content=self.chat_question.filter_sys_question())) - msg.append(HumanMessage(content=self.chat_question.filter_user_question())) + permission_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = [] + permission_sql_msg.append(SystemMessage(content=self.chat_question.filter_sys_question())) + permission_sql_msg.append(HumanMessage(content=self.chat_question.filter_user_question())) - history_msg = [] - # if self.record.full_analysis_message and self.record.full_analysis_message.strip() != '': - # history_msg = orjson.loads(self.record.full_analysis_message) - - # self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='', - # full_message=orjson.dumps(history_msg + - # [{'type': msg.type, - # 'content': msg.content} for msg - # in - # msg]).decode()) + self.current_logs[OperationEnum.GENERATE_SQL_WITH_PERMISSIONS] = start_log(session=self.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, + 'content': msg.content} for + msg + in permission_sql_msg]) full_thinking_text = '' full_filter_text = '' - res = self.llm.stream(msg) + res = self.llm.stream(permission_sql_msg) token_usage = {} for chunk in res: SQLBotLogUtil.info(chunk) @@ -583,16 +627,18 @@ class LLMService: # yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk} get_token_usage(chunk, token_usage) - msg.append(AIMessage(full_filter_text)) - # self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, - # token_usage=token_usage, - # answer=orjson.dumps({'content': full_analysis_text, - # 'reasoning_content': full_thinking_text}).decode(), - # full_message=orjson.dumps(history_msg + - # [{'type': msg.type, - # 'content': msg.content} for msg - # in - # analysis_msg]).decode()) + permission_sql_msg.append(AIMessage(full_filter_text)) + + self.current_logs[OperationEnum.GENERATE_SQL_WITH_PERMISSIONS] = end_log(session=self.session, + log=self.current_logs[ + OperationEnum.GENERATE_SQL_WITH_PERMISSIONS], + full_message=[ + {'type': msg.type, + '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 @@ -616,10 +662,16 @@ class LLMService: def generate_chart(self): # append current question self.chart_message.append(HumanMessage(self.chat_question.chart_user_question())) - self.record = save_full_chart_message(session=self.session, record_id=self.record.id, - full_message=orjson.dumps( - [{'type': msg.type, 'content': msg.content} for msg in - self.chart_message]).decode()) + + self.current_logs[OperationEnum.GENERATE_CHART] = start_log(session=self.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, 'content': msg.content} for + msg + in self.chart_message]) full_thinking_text = '' full_chart_text = '' token_usage = {} @@ -640,13 +692,16 @@ class LLMService: get_token_usage(chunk, token_usage) self.chart_message.append(AIMessage(full_chart_text)) - self.record = save_full_chart_message_and_answer(session=self.session, record_id=self.record.id, - token_usage=token_usage, - answer=orjson.dumps({'content': full_chart_text, - 'reasoning_content': full_thinking_text}).decode(), - full_message=orjson.dumps( - [{'type': msg.type, 'content': msg.content} for msg in - self.chart_message]).decode()) + + self.record = save_chart_answer(session=self.session, record_id=self.record.id, + answer=orjson.dumps({'content': full_chart_text}).decode()) + self.current_logs[OperationEnum.GENERATE_CHART] = end_log(session=self.session, + log=self.current_logs[OperationEnum.GENERATE_CHART], + full_message=[ + {'type': msg.type, 'content': msg.content} + for msg in self.chart_message], + reasoning_content=full_thinking_text, + token_usage=token_usage) def check_sql(self, res: str) -> tuple[any]: json_str = extract_nested_json(res) @@ -829,8 +884,8 @@ class LLMService: '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' + 'engine_type': self.ds.type_name or self.ds.type, + 'type': 'datasource'}).decode() + '\n\n' self.chat_question.db_schema = self.out_ds_instance.get_db_schema( self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session, @@ -909,7 +964,8 @@ class LLMService: chart = self.check_save_chart(res=full_chart_text) SQLBotLogUtil.info(chart) if in_chat: - yield 'data:' + orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n' + yield 'data:' + orjson.dumps( + {'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n' else: data = [] _fields = {} diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 075d6590..2dd4c7b2 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -91,20 +91,20 @@ exclude = ["alembic"] [tool.ruff.lint] select = [ - "E", # pycodestyle errors - "W", # pycodestyle warnings - "F", # pyflakes - "I", # isort - "B", # flake8-bugbear - "C4", # flake8-comprehensions - "UP", # pyupgrade + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "B", # flake8-bugbear + "C4", # flake8-comprehensions + "UP", # pyupgrade "ARG001", # unused arguments in functions ] ignore = [ - "E501", # line too long, handled by black - "B008", # do not perform function calls in argument defaults - "W191", # indentation contains tabs - "B904", # Allow raising exceptions without from e, for HTTPException + "E501", # line too long, handled by black + "B008", # do not perform function calls in argument defaults + "W191", # indentation contains tabs + "B904", # Allow raising exceptions without from e, for HTTPException ] [tool.ruff.lint.pyupgrade]