feat: predict data

Author: ulleo

backend/apps/chat/api/chat.py

@@ -169,8 +169,14 @@ def run_task():
     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):
+@router.post("/record/{chart_record_id}/{action_type}")
+async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chart_record_id: int, action_type: str):
+    if action_type != 'analysis' and action_type != 'predict':
+        raise HTTPException(
+            status_code=404,
+            detail="Not Found"
+        )
+
     record = session.query(ChatRecord).get(chart_record_id)
     if not record:
         raise HTTPException(
@@ -217,13 +223,29 @@ async def analysis(session: SessionDep, current_user: CurrentUser, chart_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'
+            if action_type == 'analysis':
+                # 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'
+
+            elif action_type == 'predict':
+                # generate predict
+                analysis_res = llm_service.generate_predict(session=session)
+                full_text = ''
+                for chunk in analysis_res:
+                    yield orjson.dumps({'content': chunk, 'type': 'predict-result'}).decode() + '\n\n'
+                    full_text += chunk
+                yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
+
+                _data = llm_service.check_save_predict_data(session=session, res=full_text)
+                yield orjson.dumps({'type': 'predict', 'content': _data}).decode() + '\n\n'
+
+                yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
 
-            yield orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
 
         except Exception as e:
             traceback.print_exc()

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.analysis, ChatRecord.predict, ChatRecord.finish,
-                  ChatRecord.error, ChatRecord.run_time)).filter(
+                  ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
+                  ChatRecord.predict_data, 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
@@ -178,12 +178,13 @@ def save_full_analysis_message_and_answer(session: SessionDep, record_id: int, a
 
 
 def save_full_predict_message_and_answer(session: SessionDep, record_id: int, answer: str,
-                                         full_message: str) -> ChatRecord:
+                                         full_message: str, data: 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
+    record.predict_data = data
 
     result = ChatRecord(**record.model_dump())
 
@@ -254,6 +255,23 @@ def save_chart(session: SessionDep, record_id: int, chart: str) -> ChatRecord:
     return result
 
 
+def save_predict_data(session: SessionDep, record_id: int, data: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.predict_data = data
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
 def save_error_message(session: SessionDep, record_id: int, message: str) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")

backend/apps/chat/models/chat_model.py

@@ -40,6 +40,7 @@ class ChatRecord(SQLModel, table=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))
+    predict_data: 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))

backend/apps/chat/task/llm.py

@@ -10,7 +10,7 @@
 from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_llm_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
+    finish_record, save_full_analysis_message_and_answer, save_full_predict_message_and_answer, save_predict_data
 from apps.chat.models.chat_model import ChatQuestion, ChatRecord
 from apps.datasource.models.datasource import CoreDatasource
 from apps.db.db import exec_sql
@@ -112,7 +112,7 @@ def get_record(self):
     def set_record(self, record: ChatRecord):
         self.record = record
 
-    def generate_analysis(self, session: SessionDep):
+    def get_fields_from_chart(self):
         chart_info = orjson.loads(self.record.chart)
         fields = []
         if chart_info.get('columns') and len(chart_info.get('columns')) > 0:
@@ -129,6 +129,10 @@ def generate_analysis(self, session: SessionDep):
                     if column.get('value') != column.get('name'):
                         column_str = column_str + '(' + column.get('name') + ')'
                     fields.append(column_str)
+        return fields
+
+    def generate_analysis(self, session: SessionDep):
+        fields = self.get_fields_from_chart()
 
         self.chat_question.fields = orjson.dumps(fields).decode()
         self.chat_question.data = orjson.dumps(orjson.loads(self.record.data).get('data')).decode()
@@ -169,6 +173,49 @@ def generate_analysis(self, session: SessionDep):
                                                                                        in
                                                                                        analysis_msg]).decode())
 
+    def generate_predict(self, session: SessionDep):
+        fields = self.get_fields_from_chart()
+
+        self.chat_question.fields = orjson.dumps(fields).decode()
+        self.chat_question.data = orjson.dumps(orjson.loads(self.record.data).get('data')).decode()
+        predict_msg: List[Union[BaseMessage, dict[str, Any]]] = []
+        predict_msg.append(SystemMessage(content=self.chat_question.predict_sys_question()))
+        predict_msg.append(HumanMessage(content=self.chat_question.predict_user_question()))
+
+        history_msg = []
+        if self.record.full_predict_message and self.record.full_predict_message.strip() != '':
+            history_msg = orjson.loads(self.record.full_predict_message)
+
+        self.record = save_full_predict_message_and_answer(session=session, record_id=self.record.id, answer='',
+                                                           data='',
+                                                           full_message=orjson.dumps(history_msg +
+                                                                                     [{'type': msg.type,
+                                                                                       'content': msg.content} for msg
+                                                                                      in
+                                                                                      predict_msg]).decode())
+
+        full_predict_text = ''
+        res = self.llm.stream(predict_msg)
+        for chunk in res:
+            print(chunk)
+            if isinstance(chunk, dict):
+                full_predict_text += chunk['content']
+                yield chunk['content']
+                continue
+            if isinstance(chunk, AIMessageChunk):
+                full_predict_text += chunk.content
+                yield chunk.content
+                continue
+
+        predict_msg.append(AIMessage(full_predict_text))
+        self.record = save_full_predict_message_and_answer(session=session, record_id=self.record.id,
+                                                           answer=full_predict_text, data='',
+                                                           full_message=orjson.dumps(history_msg +
+                                                                                     [{'type': msg.type,
+                                                                                       'content': msg.content} for msg
+                                                                                      in
+                                                                                      predict_msg]).decode())
+
     def generate_sql(self, session: SessionDep):
         # append current question
         self.sql_message.append(HumanMessage(self.chat_question.sql_user_question()))
@@ -274,6 +321,17 @@ def check_save_chart(self, session: SessionDep, res: str) -> Dict[str, Any]:
 
         return chart
 
+    def check_save_predict_data(self, session: SessionDep, res: str) -> Dict[str, Any]:
+
+        json_str = extract_nested_json(res)
+
+        if not json_str:
+            json_str = ''
+
+        save_predict_data(session=session, record_id=self.record.id, data=json_str)
+
+        return json_str
+
     def save_error(self, session: SessionDep, message: str):
         return save_error_message(session=session, record_id=self.record.id, message=message)
 

backend/template.yaml

@@ -134,7 +134,10 @@ template:
   predict:
     system: |
       ### 说明：
-      你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后1-2个周期的数据，用json格式返回。
+      你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后1-2个周期的数据，并将预测数据拼接在原始数据后，用json格式返回。
+      ```json
+      
+      无法预测或者不支持预测的数据请直接返回："抱歉，该数据无法进行预测。"
 
     user: |
       ### 请使用 i18n: {lang} 对应的语言输出你的结果

frontend/src/api/chat.ts

@@ -38,6 +38,7 @@ export class ChatRecord {
   chart?: string
   analysis?: string
   predict?: string
+  predict_data?: string
   finish?: boolean = false
   error?: string
   run_time: number = 0
@@ -56,6 +57,7 @@ export class ChatRecord {
     chart: string | undefined,
     analysis: string | undefined,
     predict: string | undefined,
+    predict_data: string | undefined,
     finish: boolean,
     error: string | undefined,
     run_time: number
@@ -73,6 +75,7 @@ export class ChatRecord {
     chart?: string,
     analysis?: string,
     predict?: string,
+    predict_data?: string,
     finish?: boolean,
     error?: string,
     run_time?: number
@@ -89,6 +92,7 @@ export class ChatRecord {
     this.chart = chart
     this.analysis = analysis
     this.predict = predict
+    this.predict_data = predict_data
     this.finish = finish
     this.error = error
     this.run_time = run_time ?? 0
@@ -207,6 +211,7 @@ function toChatRecord(data?: any): ChatRecord | undefined {
     data.chart,
     data.analysis,
     data.predict,
+    data.predict_data,
     data.finish,
     data.error,
     data.run_time

frontend/src/views/chat/ChatAnswer.vue

@@ -1,9 +1,9 @@
 <script setup lang="ts">
 import type { ChatMessage } from '@/api/chat.ts'
-import { computed, nextTick, ref } from 'vue'
+import { computed, ref } from 'vue'
 import { Loading } from '@element-plus/icons-vue'
-import ChartComponent from './component/ChartComponent.vue'
 import MdComponent from './component/MdComponent.vue'
+import DisplayChartBlock from './component/DisplayChartBlock.vue'
 import type { ChartTypes } from '@/views/chat/component/BaseChart.ts'
 import { ArrowDown } from '@element-plus/icons-vue'
 import ICON_BAR from '@/assets/svg/chart/bar.svg'
@@ -70,25 +70,6 @@ const chartObject = computed<{
   return {}
 })
 
-const xAxis = computed(() => {
-  if (chartObject.value?.axis?.x) {
-    return [chartObject.value.axis.x]
-  }
-  return []
-})
-const yAxis = computed(() => {
-  if (chartObject.value?.axis?.y) {
-    return [chartObject.value.axis.y]
-  }
-  return []
-})
-const series = computed(() => {
-  if (chartObject.value?.axis?.series) {
-    return [chartObject.value.axis.series]
-  }
-  return []
-})
-
 const currentChartType = ref<ChartTypes | undefined>(undefined)
 
 const chartType = computed<ChartTypes>({
@@ -163,10 +144,7 @@ const currentChartTypeIcon = computed(() => {
 const chartRef = ref()
 
 function onTypeChange() {
-  nextTick(() => {
-    chartRef.value?.destroyChart()
-    chartRef.value?.renderChart()
-  })
+  chartRef.value?.onTypeChange()
 }
 </script>
 
@@ -250,31 +228,21 @@ function onTypeChange() {
             </div>
           </template>
           <template v-else-if="settings.type === 'chart'">
-            <div>
-              <div v-if="message.record.chart" class="chart-base-container">
-                <div>
-                  <ChartComponent
-                    v-if="message.record.id"
-                    :id="message.record.id"
-                    ref="chartRef"
-                    :type="chartType"
-                    :columns="chartObject?.columns"
-                    :x="xAxis"
-                    :y="yAxis"
-                    :series="series"
-                    :data="dataObject.data"
-                  />
-                </div>
-              </div>
-            </div>
+            <DisplayChartBlock
+              :id="message.record.id"
+              ref="chartRef"
+              :chart-type="chartType"
+              :message="message"
+              :data="dataObject.data"
+            />
             <div v-if="message.record.error" style="color: red">
               {{ message.record.error }}
             </div>
           </template>
         </div>
       </template>
     </el-container>
-    <slot name="footer"></slot>
+    <slot :data="{ id: message.record?.id, chartType: chartType, chartObject: chartObject }"></slot>
   </el-container>
 </template>