feat: data analysis

Author: ulleo

backend/alembic/env.py

@@ -26,8 +26,8 @@
 
 # from apps.system.models.user import SQLModel  # noqa
 # from apps.settings.models.setting_models import SQLModel
-#from apps.chat.models.chat_model import SQLModel
-from apps.dashboard.models.dashboard_model import SQLModel
+from apps.chat.models.chat_model import SQLModel
+# from apps.dashboard.models.dashboard_model import SQLModel
 from common.core.config import settings # noqa
 
 target_metadata = SQLModel.metadata

…versions/a3af70d43e98_011_license_ddl.py …versions/a3af70d43e98_012_license_ddl.pybackend/alembic/versions/a3af70d43e98_011_license_ddl.py renamed to backend/alembic/versions/a3af70d43e98_012_license_ddl.py backend/alembic/versions/a3af70d43e98_011_license_ddl.py renamed to backend/alembic/versions/a3af70d43e98_012_license_ddl.py

@@ -11,7 +11,7 @@
 
 # revision identifiers, used by Alembic.
 revision = 'a3af70d43e98'
-down_revision = '8dc3b1bdbfef'
+down_revision = '941e2355a94d'
 branch_labels = None
 depends_on = None
 

backend/alembic/versions/bfa10ce83d73_013_modify_chat.py

@@ -0,0 +1,35 @@
+"""010_modify_chat
+
+Revision ID: bfa10ce83d73
+Revises: 8dc3b1bdbfef
+Create Date: 2025-06-18 14:16:39.230619
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = 'bfa10ce83d73'
+down_revision = 'a3af70d43e98'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('chat_record', sa.Column('analysis', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('predict', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('full_analysis_message', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('full_predict_message', sa.Text(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('chat_record', 'full_predict_message')
+    op.drop_column('chat_record', 'full_analysis_message')
+    op.drop_column('chat_record', 'predict')
+    op.drop_column('chat_record', 'analysis')
+    # ### end Alembic commands ###

backend/apps/chat/api/chat.py

@@ -113,19 +113,20 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
                                                      chart_id=request_question.chat_id)
     # get schema
     request_question.db_schema = get_table_schema(session=session, ds=ds)
-    llm_service = LLMService(request_question, history_records, CoreDatasource(**ds.model_dump()), aimodel)
+    db_user = get_user_info(session=session, user_id=current_user.id)
+    request_question.lang = db_user.language
 
-    llm_service.init_record(session=session, current_user=current_user)
+    llm_service = LLMService(request_question, aimodel, history_records, CoreDatasource(**ds.model_dump()))
 
-    db_user = get_user_info(session=session, user_id=current_user.id)
+    llm_service.init_record(session=session, current_user=current_user)
 
     def run_task():
         try:
             # return id
             yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
 
             # generate sql
-            sql_res = llm_service.generate_sql(session=session, lang=db_user.language)
+            sql_res = llm_service.generate_sql(session=session)
             full_sql_text = ''
             for chunk in sql_res:
                 full_sql_text += chunk
@@ -144,7 +145,7 @@ def run_task():
             yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n'
 
             # generate chart
-            chart_res = llm_service.generate_chart(session=session, lang=db_user.language)
+            chart_res = llm_service.generate_chart(session=session)
             full_chart_text = ''
             for chunk in chart_res:
                 full_chart_text += chunk
@@ -166,3 +167,67 @@ def run_task():
             yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
 
     return StreamingResponse(run_task(), media_type="text/event-stream")
+
+
+@router.post("/record/{chart_record_id}/analysis")
+async def analysis(session: SessionDep, current_user: CurrentUser, chart_record_id: int):
+    record = session.query(ChatRecord).get(chart_record_id)
+    if not record:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Chat record with id {chart_record_id} not found"
+        )
+
+    if not record.chart:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Chat record with id {chart_record_id} has not generated chart, do not support to analyze it"
+        )
+
+    chat = session.query(Chat).filter(Chat.id == record.chat_id).first()
+    if not chat:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Chat with id {record.chart_id} not found"
+        )
+
+    if chat.create_by != current_user.id:
+        raise HTTPException(
+            status_code=401,
+            detail=f"You cannot use the chat with id {record.chart_id}"
+        )
+
+    # Get available AI model
+    aimodel = session.exec(select(AiModelDetail).where(
+        AiModelDetail.status == True,
+        AiModelDetail.api_key.is_not(None)
+    )).first()
+    if not aimodel:
+        raise HTTPException(
+            status_code=500,
+            detail="No available AI model configuration found"
+        )
+
+    request_question = ChatQuestion(chat_id=chat.id, question='')
+    db_user = get_user_info(session=session, user_id=current_user.id)
+    request_question.lang = db_user.language
+
+    llm_service = LLMService(request_question, aimodel)
+    llm_service.set_record(record)
+
+    def run_task():
+        try:
+            # generate analysis
+            analysis_res = llm_service.generate_analysis(session=session)
+            for chunk in analysis_res:
+                yield orjson.dumps({'content': chunk, 'type': 'analysis-result'}).decode() + '\n\n'
+            yield orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
+
+            yield orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
+
+        except Exception as e:
+            traceback.print_exc()
+            # llm_service.save_error(session=session, message=str(e))
+            yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+
+    return StreamingResponse(run_task(), media_type="text/event-stream")

backend/apps/chat/curd/chat.py

@@ -59,8 +59,8 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
     record_list = session.query(ChatRecord).options(
         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.finish, ChatRecord.error,
-                  ChatRecord.run_time)).filter(
+                  ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict, ChatRecord.finish,
+                  ChatRecord.error, ChatRecord.run_time)).filter(
         and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
 
     chat_info.records = record_list
@@ -158,6 +158,44 @@ def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer
     return result
 
 
+def save_full_analysis_message_and_answer(session: SessionDep, record_id: int, answer: str,
+                                          full_message: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.full_analysis_message = full_message
+    record.analysis = answer
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
+def save_full_predict_message_and_answer(session: SessionDep, record_id: int, answer: str,
+                                         full_message: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.full_predict_message = full_message
+    record.predict = answer
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
 def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")

backend/apps/chat/models/chat_model.py

@@ -5,7 +5,9 @@
 from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity, Boolean
 from sqlmodel import SQLModel, Field
 
+from apps.template.generate_analysis.generator import get_analysis_template
 from apps.template.generate_chart.generator import get_chart_template
+from apps.template.generate_predict.generator import get_predict_template
 from apps.template.generate_sql.generator import get_sql_template
 
 
@@ -36,8 +38,12 @@ class ChatRecord(SQLModel, table=True):
     data: str = Field(sa_column=Column(Text, nullable=True))
     chart_answer: str = Field(sa_column=Column(Text, nullable=True))
     chart: str = Field(sa_column=Column(Text, nullable=True))
+    analysis: str = Field(sa_column=Column(Text, nullable=True))
+    predict: str = Field(sa_column=Column(Text, nullable=True))
     full_sql_message: str = Field(sa_column=Column(Text, nullable=True))
     full_chart_message: str = Field(sa_column=Column(Text, nullable=True))
+    full_analysis_message: str = Field(sa_column=Column(Text, nullable=True))
+    full_predict_message: str = Field(sa_column=Column(Text, 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)
@@ -72,6 +78,8 @@ class AiModelQuestion(BaseModel):
     db_schema: str = ""
     sql: str = ""
     rule: str = ""
+    fields: str = ""
+    data: str = ""
     lang: str = "zh-CN"
 
     def sql_sys_question(self):
@@ -87,6 +95,18 @@ def chart_sys_question(self):
     def chart_user_question(self):
         return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule, lang=self.lang)
 
+    def analysis_sys_question(self):
+        return get_analysis_template()['system']
+
+    def analysis_user_question(self):
+        return get_analysis_template()['user'].format(fields=self.fields, data=self.data, lang=self.lang)
+
+    def predict_sys_question(self):
+        return get_predict_template()['system']
+
+    def predict_user_question(self):
+        return get_predict_template()['user'].format(fields=self.fields, data=self.data, lang=self.lang)
+
 
 class ChatQuestion(AiModelQuestion):
     question: str