feat: support command in mcp

This commit is contained in:
ulleo
2025-12-18 19:01:11 +08:00
committed by ulleo
parent 96ab4818fe
commit 37bf819b29
5 changed files with 238 additions and 141 deletions
+82 -25
View File
@@ -8,12 +8,13 @@ import pandas as pd
from fastapi import APIRouter, HTTPException, Path
from fastapi.responses import StreamingResponse
from sqlalchemy import and_, select
from starlette.responses import JSONResponse
from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \
delete_chat, get_chat_chart_data, get_chat_predict_data, get_chat_with_records_with_data, get_chat_record_by_id, \
format_json_data, format_json_list_data, get_chart_config, list_recent_questions
from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion, AxisObj, QuickCommand, \
ChatInfo, Chat
ChatInfo, Chat, ChatFinishStep
from apps.chat.task.llm import LLMService
from apps.swagger.i18n import PLACEHOLDER_PREFIX
from apps.system.schemas.permission import SqlbotPermission, require_permissions
@@ -166,11 +167,18 @@ def find_base_question(record_id: int, session: SessionDep):
@require_permissions(permission=SqlbotPermission(type='chat', keyExpression="request_question.chat_id"))
async def question_answer(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
current_assistant: CurrentAssistant):
return await question_answer_inner(session, current_user, request_question, current_assistant, embedding=True)
async def question_answer_inner(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
current_assistant: Optional[CurrentAssistant] = None, in_chat: bool = True,
stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False):
try:
command, text_before_command, record_id, warning_info = parse_quick_command(request_question.question)
if command:
# todo 暂不支持分析和预测,需要改造前端
if command == QuickCommand.ANALYSIS or command == QuickCommand.PREDICT_DATA:
# todo 对话界面下,暂不支持分析和预测,需要改造前端
if in_chat and (command == QuickCommand.ANALYSIS or command == QuickCommand.PREDICT_DATA):
raise Exception(f'Command: {command.value} temporary not supported')
if record_id is not None:
@@ -221,53 +229,83 @@ async def question_answer(session: SessionDep, current_user: CurrentUser, reques
if command == QuickCommand.REGENERATE:
request_question.question = text_before_command
request_question.regenerate_record_id = rec_id
return await stream_sql(session, current_user, request_question, current_assistant)
return await stream_sql(session, current_user, request_question, current_assistant, in_chat, stream,
finish_step, embedding)
elif command == QuickCommand.ANALYSIS:
return await analysis_or_predict(session, current_user, rec_id, 'analysis', current_assistant)
return await analysis_or_predict(session, current_user, rec_id, 'analysis', current_assistant, in_chat, stream)
elif command == QuickCommand.PREDICT_DATA:
return await analysis_or_predict(session, current_user, rec_id, 'predict', current_assistant)
return await analysis_or_predict(session, current_user, rec_id, 'predict', current_assistant, in_chat, stream)
else:
raise Exception(f'Unknown command: {command.value}')
else:
return await stream_sql(session, current_user, request_question, current_assistant)
return await stream_sql(session, current_user, request_question, current_assistant, in_chat, stream,
finish_step, embedding)
except Exception as e:
traceback.print_exc()
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
if stream:
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
return JSONResponse(
content={'message': str(e)},
status_code=500,
)
async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
current_assistant: CurrentAssistant):
current_assistant: Optional[CurrentAssistant] = None, in_chat: bool = True, stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False):
try:
llm_service = await LLMService.create(session, current_user, request_question, current_assistant,
embedding=True)
embedding=embedding)
llm_service.init_record(session=session)
llm_service.run_task_async()
llm_service.run_task_async(in_chat=in_chat, stream=stream, finish_step=finish_step)
except Exception as e:
traceback.print_exc()
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
if stream:
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
return JSONResponse(
content={'message': str(e)},
status_code=500,
)
if stream:
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
else:
res = llm_service.await_result()
raw_data = {}
for chunk in res:
if chunk:
raw_data = chunk
status_code = 200
if not raw_data.get('success'):
status_code = 500
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
return JSONResponse(
content=raw_data,
status_code=status_code,
)
@router.post("/record/{chat_record_id}/{action_type}", summary=f"{PLACEHOLDER_PREFIX}analysis_or_predict")
async def analysis_or_predict_question(session: SessionDep, current_user: CurrentUser,
current_assistant: CurrentAssistant, chat_record_id: int,
action_type: str = Path(..., description=f"{PLACEHOLDER_PREFIX}analysis_or_predict_action_type")):
action_type: str = Path(...,
description=f"{PLACEHOLDER_PREFIX}analysis_or_predict_action_type")):
return await analysis_or_predict(session, current_user, chat_record_id, action_type, current_assistant)
async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str,
current_assistant: CurrentAssistant):
current_assistant: CurrentAssistant, in_chat: bool = True, stream: bool = True):
try:
if action_type != 'analysis' and action_type != 'predict':
raise Exception(f"Type {action_type} Not Found")
@@ -294,16 +332,35 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
llm_service = await LLMService.create(session, current_user, request_question, current_assistant)
llm_service.run_analysis_or_predict_task_async(session, action_type, record)
llm_service.run_analysis_or_predict_task_async(session, action_type, record, in_chat, stream)
except Exception as e:
traceback.print_exc()
if stream:
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
return JSONResponse(
content={'message': str(e)},
status_code=500,
)
if stream:
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
else:
res = llm_service.await_result()
raw_data = {}
for chunk in res:
if chunk:
raw_data = chunk
status_code = 200
if not raw_data.get('success'):
status_code = 500
return StreamingResponse(_err(e), media_type="text/event-stream")
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
return JSONResponse(
content=raw_data,
status_code=status_code,
)
@router.get("/record/{chat_record_id}/excel/export", summary=f"{PLACEHOLDER_PREFIX}export_chart_data")
+11
View File
@@ -163,6 +163,17 @@ def format_json_list_data(origin_data: list[dict]):
return data
def get_chat_chart_config(session: SessionDep, chat_record_id: int):
stmt = select(ChatRecord.chart).where(and_(ChatRecord.id == chat_record_id))
res = session.execute(stmt)
for row in res:
try:
return orjson.loads(row.data)
except Exception:
pass
return {}
def get_chat_chart_data(session: SessionDep, chat_record_id: int):
stmt = select(ChatRecord.data).where(and_(ChatRecord.id == chat_record_id))
res = session.execute(stmt)
+112 -45
View File
@@ -30,7 +30,8 @@ from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
save_select_datasource_answer, save_recommend_question_answer, \
get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \
get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate, get_chat_predict_data, \
get_chat_chart_config
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \
ChatFinishStep, AxisObj
from apps.data_training.curd.data_training import get_training_template
@@ -969,6 +970,10 @@ class LLMService:
'regenerate_record_id': self.get_record().regenerate_record_id}).decode() + '\n\n'
yield 'data:' + orjson.dumps(
{'type': 'question', 'question': self.get_record().question}).decode() + '\n\n'
else:
if stream:
yield '> ID: ' + str(self.get_record().id) + '\n'
yield '> ' + self.get_record().question + '\n\n'
if not stream:
json_result['record_id'] = self.get_record().id
@@ -1150,26 +1155,8 @@ class LLMService:
{'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
else:
if stream:
_fields = {}
if chart.get('columns'):
for _column in chart.get('columns'):
if _column:
_fields[_column.get('value')] = _column.get('name')
if chart.get('axis'):
if chart.get('axis').get('x'):
_fields[chart.get('axis').get('x').get('value')] = chart.get('axis').get('x').get('name')
if chart.get('axis').get('y'):
_fields[chart.get('axis').get('y').get('value')] = chart.get('axis').get('y').get('name')
if chart.get('axis').get('series'):
_fields[chart.get('axis').get('series').get('value')] = chart.get('axis').get('series').get(
'name')
_column_list = []
for field in result.get('fields'):
_column_list.append(
AxisObj(name=field if not _fields.get(field) else _fields.get(field), value=field))
md_data, _fields_list = DataFormat.convert_object_array_for_pandas(_column_list, result.get('data'))
md_data, _fields_list = DataFormat.convert_data_fields_for_pandas(chart, result.get('fields'),
result.get('data'))
# data, _fields_list, col_formats = self.format_pd_data(_column_list, result.get('data'))
if not md_data or not _fields_list:
@@ -1183,7 +1170,7 @@ class LLMService:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n'
else:
# todo generate picture
# generate picture
try:
if chart['type'] != 'table':
yield '### generated chart picture\n\n'
@@ -1255,51 +1242,123 @@ class LLMService:
finally:
session_maker.remove()
def run_analysis_or_predict_task_async(self, session: Session, action_type: str, base_record: ChatRecord):
def run_analysis_or_predict_task_async(self, session: Session, action_type: str, base_record: ChatRecord,
in_chat: bool = True, stream: bool = True):
self.set_record(save_analysis_predict_record(session, base_record, action_type))
self.future = executor.submit(self.run_analysis_or_predict_task_cache, action_type)
self.future = executor.submit(self.run_analysis_or_predict_task_cache, action_type, in_chat, stream)
def run_analysis_or_predict_task_cache(self, action_type: str):
for chunk in self.run_analysis_or_predict_task(action_type):
def run_analysis_or_predict_task_cache(self, action_type: str, in_chat: bool = True, stream: bool = True):
for chunk in self.run_analysis_or_predict_task(action_type, in_chat, stream):
self.chunk_list.append(chunk)
def run_analysis_or_predict_task(self, action_type: str):
def run_analysis_or_predict_task(self, action_type: str, in_chat: bool = True, stream: bool = True):
json_result: Dict[str, Any] = {'success': True}
_session = None
try:
_session = session_maker()
yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
else:
if stream:
yield '> ID: ' + str(self.get_record().id) + '\n'
yield '> ' + self.get_record().question + '\n\n'
if not stream:
json_result['record_id'] = self.get_record().id
if action_type == 'analysis':
# generate analysis
analysis_res = self.generate_analysis(_session)
full_text = ''
for chunk in analysis_res:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'analysis-result'}).decode() + '\n\n'
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
yield 'data:' + orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
full_text += chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'analysis-result'}).decode() + '\n\n'
else:
if stream:
yield chunk.get('content')
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
yield 'data:' + orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
else:
if stream:
yield '\n\n'
if not stream:
json_result['content'] = full_text
elif action_type == 'predict':
# generate predict
analysis_res = self.generate_predict(_session)
full_text = ''
for chunk in analysis_res:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'predict-result'}).decode() + '\n\n'
full_text += chunk.get('content')
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps(
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
'type': 'predict-result'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
_data = self.check_save_predict_data(session=_session, res=full_text)
if _data:
yield 'data:' + orjson.dumps({'type': 'predict-success'}).decode() + '\n\n'
has_data = self.check_save_predict_data(session=_session, res=full_text)
if has_data:
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict-success'}).decode() + '\n\n'
else:
chart = get_chat_chart_config(_session, self.record.id)
origin_data = get_chat_chart_data(_session, self.record.id)
predict_data = get_chat_predict_data(_session, self.record.id)
if stream:
md_data, _fields_list = DataFormat.convert_data_fields_for_pandas(chart,
origin_data.get('fields'),
predict_data)
if not md_data or not _fields_list:
yield 'Predict data result is empty.\n\n'
else:
df = pd.DataFrame(md_data, columns=_fields_list)
df_safe = DataFormat.safe_convert_to_string(df)
markdown_table = df_safe.to_markdown(index=False)
yield markdown_table + '\n\n'
else:
json_result['origin_data'] = origin_data
json_result['predict_data'] = predict_data
# generate picture
try:
if chart['type'] != 'table':
yield '### generated chart picture\n\n'
_data = get_chat_chart_data(_session, self.record.id)
_data['data'] = _data['data'] + predict_data
image_url = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(_data))
SQLBotLogUtil.info(image_url)
if stream:
yield f'![{chart["type"]}]({image_url})'
else:
json_result['image_url'] = image_url
except Exception as e:
if stream:
raise e
else:
yield 'data:' + orjson.dumps({'type': 'predict-failed'}).decode() + '\n\n'
yield 'data:' + orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict-failed'}).decode() + '\n\n'
else:
if stream:
yield full_text + '\n\n'
if not stream:
json_result['success'] = False
json_result['message'] = full_text
if in_chat:
yield 'data:' + orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
self.finish(_session)
if not stream:
yield json_result
except Exception as e:
error_msg: str
if isinstance(e, SingleMessageError):
@@ -1308,7 +1367,15 @@ class LLMService:
error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
if _session:
self.save_error(session=_session, message=error_msg)
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
else:
if stream:
yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。'
else:
json_result['success'] = False
json_result['message'] = error_msg
yield json_result
finally:
# end
session_maker.remove()
+6 -71
View File
@@ -1,22 +1,18 @@
# Author: Junjun
# Date: 2025/7/1
import json
import traceback
from datetime import timedelta
import jwt
from fastapi import HTTPException, status, APIRouter
from fastapi.responses import StreamingResponse
# from fastapi.security import OAuth2PasswordBearer
from jwt.exceptions import InvalidTokenError
from pydantic import ValidationError
from sqlmodel import select
from starlette.responses import JSONResponse
from apps.chat.api.chat import create_chat
from apps.chat.api.chat import create_chat, question_answer_inner
from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion, McpAssistant, ChatQuestion, \
ChatFinishStep
from apps.chat.task.llm import LLMService
from apps.system.crud.user import authenticate
from apps.system.crud.user import get_db_user
from apps.system.models.system_model import UserWsModel
@@ -110,39 +106,8 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question)
try:
llm_service = await LLMService.create(session, session_user, mcp_chat)
llm_service.init_record(session=session)
llm_service.run_task_async(False, chat.stream)
except Exception as e:
traceback.print_exc()
if chat.stream:
def _err(_e: Exception):
yield str(_e) + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
return JSONResponse(
content={'message': str(e)},
status_code=500,
)
if chat.stream:
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
else:
res = llm_service.await_result()
raw_data = {}
for chunk in res:
if chunk:
raw_data = chunk
status_code = 200
if not raw_data.get('success'):
status_code = 500
return JSONResponse(
content=raw_data,
status_code=status_code,
)
return await question_answer_inner(session=session, current_user=session_user, request_question=mcp_chat,
in_chat=False, stream=chat.stream)
@router.post("/mcp_assistant", operation_id="mcp_assistant")
@@ -166,36 +131,6 @@ async def mcp_assistant(session: SessionDep, chat: McpAssistant):
# assistant question
mcp_chat = ChatQuestion(chat_id=c.id, question=chat.question)
# ask
try:
llm_service = await LLMService.create(session, session_user, mcp_chat, mcp_assistant_header)
llm_service.init_record(session=session)
llm_service.run_task_async(False, chat.stream, ChatFinishStep.QUERY_DATA)
except Exception as e:
traceback.print_exc()
if chat.stream:
def _err(_e: Exception):
yield str(_e) + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
return JSONResponse(
content={'message': str(e)},
status_code=500,
)
if chat.stream:
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
else:
res = llm_service.await_result()
raw_data = {}
for chunk in res:
if chunk:
raw_data = chunk
status_code = 200
if not raw_data.get('success'):
status_code = 500
return JSONResponse(
content=raw_data,
status_code=status_code,
)
return await question_answer_inner(session=session, current_user=session_user, request_question=mcp_chat,
current_assistant=mcp_assistant_header,
in_chat=False, stream=chat.stream, finish_step=ChatFinishStep.QUERY_DATA)
+27
View File
@@ -1,5 +1,8 @@
import pandas as pd
from apps.chat.models.chat_model import AxisObj
class DataFormat:
@staticmethod
def safe_convert_to_string(df):
@@ -77,6 +80,30 @@ class DataFormat:
md_data.append(_row)
return md_data, _fields_list
@staticmethod
def convert_data_fields_for_pandas(chart: dict, fields: list, data: list):
_fields = {}
if chart.get('columns'):
for _column in chart.get('columns'):
if _column:
_fields[_column.get('value')] = _column.get('name')
if chart.get('axis'):
if chart.get('axis').get('x'):
_fields[chart.get('axis').get('x').get('value')] = chart.get('axis').get('x').get('name')
if chart.get('axis').get('y'):
_fields[chart.get('axis').get('y').get('value')] = chart.get('axis').get('y').get('name')
if chart.get('axis').get('series'):
_fields[chart.get('axis').get('series').get('value')] = chart.get('axis').get('series').get(
'name')
_column_list = []
for field in fields:
_column_list.append(
AxisObj(name=field if not _fields.get(field) else _fields.get(field), value=field))
md_data, _fields_list = DataFormat.convert_object_array_for_pandas(_column_list, data)
return md_data, _fields_list
@staticmethod
def format_pd_data(column_list: list, data_list: list, col_formats: dict = None):
# 预处理数据并记录每列的格式类型