diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index b23f4c41..f79df6f5 100644 --- a/backend/apps/chat/api/chat.py +++ b/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() diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index b7df18a9..9f9be02e 100644 --- a/backend/apps/chat/curd/chat.py +++ b/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") diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 06d08c03..e6c51fca 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/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)) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index a07ade68..1e3f440c 100644 --- a/backend/apps/chat/task/llm.py +++ b/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) diff --git a/backend/template.yaml b/backend/template.yaml index 3600eac5..2c811245 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -134,7 +134,10 @@ template: predict: system: | ### 说明： - 你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后1-2个周期的数据，用json格式返回。 + 你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后1-2个周期的数据，并将预测数据拼接在原始数据后，用json格式返回。 + ```json + + 无法预测或者不支持预测的数据请直接返回："抱歉，该数据无法进行预测。" user: | ### 请使用 i18n: {lang} 对应的语言输出你的结果 diff --git a/frontend/src/api/chat.ts b/frontend/src/api/chat.ts index 6d4ba976..58a03ef4 100644 --- a/frontend/src/api/chat.ts +++ b/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 diff --git a/frontend/src/views/chat/ChatAnswer.vue b/frontend/src/views/chat/ChatAnswer.vue index c0deb770..e32a71e5 100644 --- a/frontend/src/views/chat/ChatAnswer.vue +++ b/frontend/src/views/chat/ChatAnswer.vue @@ -1,9 +1,9 @@ @@ -250,23 +228,13 @@ function onTypeChange() { - + diff --git a/frontend/src/views/chat/component/DisplayChartBlock.vue b/frontend/src/views/chat/component/DisplayChartBlock.vue new file mode 100644 index 00000000..e16c6e9a --- /dev/null +++ b/frontend/src/views/chat/component/DisplayChartBlock.vue @@ -0,0 +1,88 @@ + + + + + diff --git a/frontend/src/views/chat/component/PredictChartBlock.vue b/frontend/src/views/chat/component/PredictChartBlock.vue new file mode 100644 index 00000000..220a3144 --- /dev/null +++ b/frontend/src/views/chat/component/PredictChartBlock.vue @@ -0,0 +1,57 @@ + + + + + diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index 20f68424..7b35a70d 100644 --- a/frontend/src/views/chat/index.vue +++ b/frontend/src/views/chat/index.vue @@ -32,23 +32,36 @@ :msg="message" >