feat: split analysis & predict in chat

Author: ulleo

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 ###

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,13 +157,12 @@ 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(
             status_code=500,
             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")

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='')
 

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

backend/apps/chat/task/llm.py

@@ -17,8 +17,8 @@
 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 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
 
         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

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")

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[] => {