diff --git a/backend/alembic/versions/024_modify_chat_record.py b/backend/alembic/versions/024_modify_chat_record.py new file mode 100644 index 00000000..2b7e635b --- /dev/null +++ b/backend/alembic/versions/024_modify_chat_record.py @@ -0,0 +1,33 @@ +"""024_modify_chat_record + +Revision ID: 806bc67ff45f +Revises: f535d09946f6 +Create Date: 2025-07-11 18:09:52.417628 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '806bc67ff45f' +down_revision = 'f535d09946f6' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('chat_record', sa.Column('analysis_record_id', sa.BigInteger(), nullable=True)) + op.add_column('chat_record', sa.Column('predict_record_id', sa.BigInteger(), nullable=True)) + op.drop_column('chat_record', 'run_time') + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('chat_record', sa.Column('run_time', sa.DOUBLE_PRECISION(precision=53), autoincrement=False, nullable=False)) + op.drop_column('chat_record', 'predict_record_id') + op.drop_column('chat_record', 'analysis_record_id') + # ### end Alembic commands ### diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 4044d898..82c18003 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -129,7 +129,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que detail=str(e) ) - return StreamingResponse(run_task(llm_service, session), media_type="text/event-stream") + return StreamingResponse(run_task(llm_service), media_type="text/event-stream") @router.post("/record/{chat_record_id}/{action_type}") @@ -157,7 +157,6 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch try: llm_service = LLMService(session, current_user, request_question) - llm_service.set_record(record) except Exception as e: traceback.print_exc() raise HTTPException( @@ -165,5 +164,5 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch detail=str(e) ) - return StreamingResponse(run_analysis_or_predict_task(llm_service, action_type), + return StreamingResponse(run_analysis_or_predict_task(llm_service, action_type, record), media_type="text/event-stream") diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 0c399bda..f653658c 100644 --- a/backend/apps/chat/curd/chat.py +++ b/backend/apps/chat/curd/chat.py @@ -82,9 +82,9 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr load_only(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time, ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, ChatRecord.data, ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict, - ChatRecord.datasource_select_answer, + ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id, ChatRecord.recommended_question, ChatRecord.first_chat, - ChatRecord.predict_data, ChatRecord.finish, ChatRecord.error, ChatRecord.run_time)).filter( + ChatRecord.predict_data, ChatRecord.finish, ChatRecord.error)).filter( and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all() result = list(map(format_record, record_list)) @@ -130,9 +130,11 @@ def format_record(record: ChatRecord): return _dict -def list_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> List[ChatRecord]: +def list_base_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> List[ChatRecord]: record_list = session.query(ChatRecord).filter( - and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all() + and_(Chat.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).all() return record_list @@ -225,6 +227,34 @@ def save_question(session: SessionDep, current_user: CurrentUser, question: Chat return result +def save_analysis_predict_record(session: SessionDep, base_record: ChatRecord, action_type: str) -> ChatRecord: + record = ChatRecord() + record.question = base_record.question + record.chat_id = base_record.chat_id + record.datasource = base_record.datasource + record.engine_type = base_record.engine_type + record.ai_modal_id = base_record.ai_modal_id + record.create_time = datetime.datetime.now() + record.create_by = base_record.id + record.chart = base_record.chart + record.data = base_record.data + + if action_type == 'analysis': + record.analysis_record_id = base_record.id + elif action_type == 'predict': + record.predict_record_id = base_record.id + + result = ChatRecord(**record.model_dump()) + + session.add(record) + session.flush() + session.refresh(record) + result.id = record.id + session.commit() + + 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='') diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 9832a432..4aa02266 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -62,8 +62,8 @@ class ChatRecord(SQLModel, table=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)) - run_time: float = Field(default=0) - + analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) + predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) class CreateChat(BaseModel): id: int = None diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 674d644c..d495e0dc 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -17,8 +17,8 @@ from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_confi 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, list_records, save_full_recommend_question_message_and_answer, \ - get_old_questions + save_full_select_datasource_message_and_answer, save_full_recommend_question_message_and_answer, \ + get_old_questions, save_analysis_predict_record, list_base_records from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat from apps.datasource.crud.datasource import get_table_schema from apps.datasource.models.datasource import CoreDatasource @@ -63,9 +63,9 @@ class LLMService: history_records: List[ChatRecord] = list( map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False, - list_records(session=self.session, - current_user=current_user, - chart_id=chat_question.chat_id)))) + list_base_records(session=self.session, + current_user=current_user, + chart_id=chat_question.chat_id)))) # get schema if ds: chat_question.db_schema = get_table_schema(session=self.session, ds=ds) @@ -606,7 +606,7 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str: raise RuntimeError(error_msg) -def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True): +def run_task(llm_service: LLMService, in_chat: bool = True): try: # return id if in_chat: @@ -626,7 +626,7 @@ def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True) yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name, 'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n' - llm_service.chat_question.db_schema = get_table_schema(session=session, ds=llm_service.ds) + llm_service.chat_question.db_schema = get_table_schema(session=llm_service.session, ds=llm_service.ds) # generate sql sql_res = llm_service.generate_sql() @@ -720,8 +720,10 @@ def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True) yield f'> ❌ **ERROR**\n\n> \n\n> {str(e)}。' -def run_analysis_or_predict_task(llm_service: LLMService, action_type: str): +def run_analysis_or_predict_task(llm_service: LLMService, action_type: str, base_record: ChatRecord): try: + llm_service.set_record(save_analysis_predict_record(llm_service.session, base_record, action_type)) + if action_type == 'analysis': # generate analysis analysis_res = llm_service.generate_analysis() @@ -752,10 +754,10 @@ def run_analysis_or_predict_task(llm_service: LLMService, action_type: str): yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n' - + llm_service.finish() except Exception as e: traceback.print_exc() - # llm_service.save_error(session=session, message=str(e)) + llm_service.save_error(message=str(e)) yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n' finally: # end diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 03792190..3530f493 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -66,4 +66,4 @@ async def mcp_question(session: SessionDep, chat: ChatMcp): llm_service = LLMService(session, user, chat) llm_service.init_record() - return StreamingResponse(run_task(llm_service, session, False), media_type="text/event-stream") + return StreamingResponse(run_task(llm_service, False), media_type="text/event-stream") diff --git a/frontend/src/api/chat.ts b/frontend/src/api/chat.ts index 2e3c6857..2b85bd20 100644 --- a/frontend/src/api/chat.ts +++ b/frontend/src/api/chat.ts @@ -47,6 +47,8 @@ export class ChatRecord { run_time: number = 0 first_chat: boolean = false recommended_question?: string + analysis_record_id?: number + predict_record_id?: number constructor() constructor( @@ -69,7 +71,9 @@ export class ChatRecord { error: string | undefined, run_time: number, first_chat: boolean, - recommended_question: string | undefined + recommended_question: string | undefined, + analysis_record_id: number | undefined, + predict_record_id: number | undefined ) constructor( id?: number, @@ -91,7 +95,9 @@ export class ChatRecord { error?: string, run_time?: number, first_chat?: boolean, - recommended_question?: string + recommended_question?: string, + analysis_record_id?: number, + predict_record_id?: number ) { this.id = id this.chat_id = chat_id @@ -113,6 +119,8 @@ export class ChatRecord { this.run_time = run_time ?? 0 this.first_chat = !!first_chat this.recommended_question = recommended_question + this.analysis_record_id = analysis_record_id + this.predict_record_id = predict_record_id } } @@ -235,7 +243,9 @@ const toChatRecord = (data?: any): ChatRecord | undefined => { data.error, data.run_time, data.first_chat, - data.recommended_question + data.recommended_question, + data.analysis_record_id, + data.predict_record_id ) } const toChatRecordList = (list: any = []): ChatRecord[] => { diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index 4ce182ed..95efba46 100644 --- a/frontend/src/views/chat/index.vue +++ b/frontend/src/views/chat/index.vue @@ -46,7 +46,13 @@