klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support command in mcp

Browse Source
This commit is contained in:
ulleo
2025-12-22 16:46:41 +08:00
committed by ulleo
parent 2522137678
commit b9b3376396
4 changed files with 40 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/api/chat.py
+13 -8
View File
@@ -198,11 +198,10 @@ async def question_answer_inner(session: SessionDep, current_user: CurrentUser,
if rec_first_chat:
raise Exception(f'Record id: {record_id} does not support this operation')
if command == QuickCommand.REGENERATE:
if rec_analysis_record_id:
raise Exception('Analysis record does not support this operation')
if rec_predict_record_id:
raise Exception('Predict data record does not support this operation')
if rec_analysis_record_id:
raise Exception('Analysis record does not support this operation')
if rec_predict_record_id:
raise Exception('Predict data record does not support this operation')
else: # get last record id
stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.regenerate_record_id).where(
@@ -233,10 +232,12 @@ async def question_answer_inner(session: SessionDep, current_user: CurrentUser,
finish_step, embedding)
elif command == QuickCommand.ANALYSIS:
return await analysis_or_predict(session, current_user, rec_id, 'analysis', current_assistant, in_chat, stream)
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, in_chat, stream)
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:
@@ -247,7 +248,11 @@ async def question_answer_inner(session: SessionDep, current_user: CurrentUser,
if stream:
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
if in_chat:
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
else:
yield f'❌ **ERROR:**\n'
yield f'> {str(_e)}\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
else:
backend/apps/chat/curd/chat.py
+1 -1
View File
@@ -168,7 +168,7 @@ def get_chat_chart_config(session: SessionDep, chat_record_id: int):
res = session.execute(stmt)
for row in res:
try:
return orjson.loads(row.data)
return orjson.loads(row.chart)
except Exception:
pass
return {}
backend/apps/chat/task/llm.py
+25 -14
View File
@@ -1172,15 +1172,17 @@ class LLMService:
else:
# generate picture
try:
if chart['type'] != 'table':
if chart.get('type') != 'table':
yield '### generated chart picture\n\n'
image_url = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(result))
image_url, error = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(result))
SQLBotLogUtil.info(image_url)
if stream:
yield f'![{chart["type"]}]({image_url})'
yield f'![{chart.get("type")}]({image_url})'
else:
json_result['image_url'] = image_url
if error is not None:
raise error
except Exception as e:
if stream:
raise e
@@ -1207,7 +1209,8 @@ class LLMService:
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
else:
if stream:
yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。'
yield f'❌ **ERROR:**\n'
yield f'> {error_msg}\n'
else:
json_result['success'] = False
json_result['message'] = error_msg
@@ -1327,19 +1330,21 @@ class LLMService:
# generate picture
try:
if chart['type'] != 'table':
if chart.get('type') != 'table':
yield '### generated chart picture\n\n'
_data = get_chat_chart_data(_session, self.record.id)
_data['data'] = _data['data'] + predict_data
_data['data'] = _data.get('data') + predict_data
image_url = request_picture(self.record.chat_id, self.record.id, chart,
format_json_data(_data))
image_url, error = 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})'
yield f'![{chart.get("type")}]({image_url})'
else:
json_result['image_url'] = image_url
if error is not None:
raise error
except Exception as e:
if stream:
raise e
@@ -1360,6 +1365,7 @@ class LLMService:
if not stream:
yield json_result
except Exception as e:
traceback.print_exc()
error_msg: str
if isinstance(e, SingleMessageError):
error_msg = str(e)
@@ -1371,7 +1377,8 @@ class LLMService:
yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
else:
if stream:
yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。'
yield f'❌ **ERROR:**\n'
yield f'> {error_msg}\n'
else:
json_result['success'] = False
json_result['message'] = error_msg
@@ -1451,16 +1458,20 @@ def request_picture(chat_id: int, record_id: int, chart: dict, data: dict):
request_obj = {
"path": os.path.join(settings.MCP_IMAGE_PATH, file_name),
"type": chart['type'],
"type": chart.get('type'),
"data": orjson.dumps(data.get('data') if data.get('data') else []).decode(),
"axis": orjson.dumps(axis).decode(),
}
requests.post(url=settings.MCP_IMAGE_HOST, json=request_obj)
_error = None
try:
requests.post(url=settings.MCP_IMAGE_HOST, json=request_obj, timeout=settings.SERVER_IMAGE_TIMEOUT)
except Exception as e:
_error = e
request_path = urllib.parse.urljoin(settings.SERVER_IMAGE_HOST, f"{file_name}.png")
return request_path
return request_path, _error
def get_token_usage(chunk: BaseMessageChunk, token_usage: dict = None):
backend/common/core/config.py
+1
View File
@@ -87,6 +87,7 @@ class Settings(BaseSettings):
EXCEL_PATH: str = '/opt/sqlbot/data/excel'
MCP_IMAGE_HOST: str = 'http://localhost:3000'
SERVER_IMAGE_HOST: str = 'http://YOUR_SERVE_IP:MCP_PORT/images/'
SERVER_IMAGE_TIMEOUT: int = 15
LOCAL_MODEL_PATH: str = '/opt/sqlbot/models'
DEFAULT_EMBEDDING_MODEL: str = 'shibing624/text2vec-base-chinese'