klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: improve chat log table structure

Browse Source
This commit is contained in:
ulleo
2025-08-21 14:52:36 +08:00
parent 4264d8e4ad
commit 667536e16c
6 changed files with 642 additions and 321 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/alembic/versions/037_create_chat_log.py
+131
View File
@@ -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 ###
backend/alembic/versions/038_remove_chat_record_cloumns.py
+51
View File
@@ -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 ###
backend/apps/chat/curd/chat.py
+163 -141
View File
@@ -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:
backend/apps/chat/models/chat_model.py
+73 -12
View File
@@ -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 = ""
backend/apps/chat/task/llm.py
+213 -157
View File
@@ -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 = {}
backend/pyproject.toml
+11 -11
View File
@@ -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]