feat: predict data

backend/apps/chat/api/chat.py

@@ -169,8 +169,14 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     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 langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AI
 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 @@ class LLMService:
     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 @@ class LLMService:
                     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 @@ class LLMService:
                                                                                        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 @@ class LLMService:
 
         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,23 +228,13 @@ 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>
@@ -274,7 +242,7 @@ function onTypeChange() {
         </div>
       </template>
     </el-container>
-    <slot name="footer"></slot>
+    <slot :data="{ id: message.record?.id, chartType: chartType, chartObject: chartObject }"></slot>
   </el-container>
 </template>
 

frontend/src/views/chat/component/DisplayChartBlock.vue

@@ -0,0 +1,88 @@
+<script setup lang="ts">
+import ChartComponent from '@/views/chat/component/ChartComponent.vue'
+import type { ChatMessage } from '@/api/chat.ts'
+import { computed, nextTick, ref } from 'vue'
+import type { ChartTypes } from '@/views/chat/component/BaseChart.ts'
+
+const props = defineProps<{
+  id?: number | string
+  chartType: ChartTypes
+  message: ChatMessage
+  data: Array<{ [key: string]: any }>
+}>()
+
+const chartObject = computed<{
+  type: ChartTypes
+  title: string
+  axis: {
+    x: { name: string; value: string }
+    y: { name: string; value: string }
+    series: { name: string; value: string }
+  }
+  columns: Array<{ name: string; value: string }>
+}>(() => {
+  if (props.message?.record?.chart) {
+    return JSON.parse(props.message.record.chart)
+  }
+  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 chartRef = ref()
+
+function onTypeChange() {
+  nextTick(() => {
+    chartRef.value?.destroyChart()
+    chartRef.value?.renderChart()
+  })
+}
+
+defineExpose({
+  onTypeChange,
+})
+</script>
+
+<template>
+  <div v-if="message?.record">
+    <div v-if="message.record.chart" class="chart-base-container">
+      <div>
+        <ChartComponent
+          v-if="message.record.id"
+          :id="id ?? 'default_chat_id'"
+          ref="chartRef"
+          :type="chartType"
+          :columns="chartObject?.columns"
+          :x="xAxis"
+          :y="yAxis"
+          :series="series"
+          :data="data"
+        />
+      </div>
+    </div>
+  </div>
+</template>
+
+<style scoped lang="less">
+.chart-base-container {
+  padding: 20px;
+  background: rgba(224, 224, 226, 0.29);
+}
+</style>

frontend/src/views/chat/component/PredictChartBlock.vue

@@ -0,0 +1,57 @@
+<script setup lang="ts">
+import type { ChatMessage } from '@/api/chat.ts'
+import { computed, watch, ref } from 'vue'
+import type { ChartTypes } from '@/views/chat/component/BaseChart.ts'
+import DisplayChartBlock from '@/views/chat/component/DisplayChartBlock.vue'
+
+const props = defineProps<{
+  id?: number | string
+  chartType: ChartTypes
+  message: ChatMessage
+  data: string
+}>()
+
+const _data = computed(() => {
+  if (
+    props.data &&
+    props.data.length > 0 &&
+    props.data.trim().startsWith('[') &&
+    props.data.trim().endsWith(']')
+  ) {
+    try {
+      return JSON.parse(props.data)
+    } catch (e) {
+      console.error(e)
+    }
+  }
+
+  return []
+})
+
+const blockRef = ref()
+
+watch(
+  () => props.chartType,
+  () => {
+    blockRef.value?.onTypeChange()
+  }
+)
+</script>
+
+<template>
+  <DisplayChartBlock
+    v-if="_data.length > 0"
+    :id="id"
+    ref="blockRef"
+    :chart-type="chartType"
+    :message="message"
+    :data="_data"
+  />
+</template>
+
+<style scoped lang="less">
+.chart-base-container {
+  padding: 20px;
+  background: rgba(224, 224, 226, 0.29);
+}
+</style>

frontend/src/views/chat/index.vue

@@ -32,23 +32,36 @@
               :msg="message"
             >
               <template v-if="message.role === 'assistant'">
-                <ChatAnswer :message="message">
-                  <template #footer>
-                    <div style="padding: 0 22px; display: flex; justify-content: flex-end">
-                      <el-button text type="primary" @click="clickAnalysis(message.record?.id)">
-                        {{ t('chat.data_analysis') }}
-                      </el-button>
-                      <el-button text type="primary" @click="clickPredict(message.record?.id)">
-                        {{ t('chat.data_predict') }}
-                      </el-button>
-                    </div>
-                    <div class="analysis-container">
-                      <MdComponent
-                        v-if="message.record?.analysis || isAnalysisTyping"
-                        :message="message.record?.analysis"
+                <ChatAnswer v-slot="{ data }" :message="message">
+                  <div style="padding: 0 22px; display: flex; justify-content: flex-end">
+                    <el-button text type="primary" @click="clickAnalysis(message.record?.id)">
+                      {{ t('chat.data_analysis') }}
+                    </el-button>
+                    <el-button text type="primary" @click="clickPredict(message.record?.id)">
+                      {{ t('chat.data_predict') }}
+                    </el-button>
+                  </div>
+                  <div class="analysis-container">
+                    <MdComponent
+                      v-if="message.record?.analysis || isAnalysisTyping"
+                      :message="message.record?.analysis"
+                    />
+                    <el-divider
+                      v-if="
+                        (message.record?.analysis || isAnalysisTyping) &&
+                        (message.record?.predict || isPredictTyping)
+                      "
+                    />
+                    <template v-if="message.record?.predict || isPredictTyping">
+                      <MdComponent :message="message.record?.predict" />
+                      <PredictChartBlock
+                        :id="message.record?.id + '-predict'"
+                        :data="message.record?.predict_data ?? '[]'"
+                        :message="message"
+                        :chart-type="data.chartType"
                       />
-                    </div>
-                  </template>
+                    </template>
+                  </div>
                 </ChatAnswer>
               </template>
               <template v-if="message.role === 'assistant'" #footer>
@@ -67,7 +80,7 @@
         <div class="input-wrapper">
           <el-input
             v-model="inputMessage"
-            :disabled="isTyping"
+            :disabled="isTyping || isPredictTyping || isAnalysisTyping"
             class="input-area"
             type="textarea"
             :rows="1"
@@ -80,7 +93,7 @@
             link
             type="primary"
             class="input-icon"
-            :disabled="isTyping"
+            :disabled="isTyping || isPredictTyping || isAnalysisTyping"
             @click="sendMessage"
           >
             <el-icon size="20">
@@ -101,6 +114,7 @@ import ChatList from './ChatList.vue'
 import ChatRow from './ChatRow.vue'
 import ChatAnswer from './ChatAnswer.vue'
 import MdComponent from './component/MdComponent.vue'
+import PredictChartBlock from './component/PredictChartBlock.vue'
 import { useI18n } from 'vue-i18n'
 import { find } from 'lodash-es'
 
@@ -389,6 +403,7 @@ async function clickAnalysis(id?: number) {
   if (currentRecord == undefined) {
     return
   }
+  isAnalysisTyping.value = true
   currentChat.value.records[_index].analysis = ''
 
   try {
@@ -484,7 +499,9 @@ async function clickPredict(id?: number) {
   if (currentRecord == undefined) {
     return
   }
+  isPredictTyping.value = true
   currentChat.value.records[_index].predict = ''
+  currentChat.value.records[_index].predict_data = ''
 
   try {
     const response = await chatApi.predict(id)
@@ -548,6 +565,9 @@ async function clickPredict(id?: number) {
             predict_answer += data.content
             currentChat.value.records[_index].predict = predict_answer
             break
+          case 'predict':
+            currentChat.value.records[_index].predict_data = data.content
+            break
           case 'predict_finish':
             isPredictTyping.value = false
             break