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: chat question for mcp

Browse Source
This commit is contained in:
ulleo
2025-07-02 16:11:58 +08:00
parent cc9af61915
commit 8afb148e33
3 changed files with 47 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+44 -25
View File
@@ -2,6 +2,7 @@ import logging
import warnings
from typing import Any, List, Union, Dict
import numpy as np
import orjson
import pandas as pd
from langchain_community.utilities import SQLDatabase
@@ -543,16 +544,16 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
raise RuntimeError(error_msg)
def run_task(llm_service: LLMService, session: SessionDep, stream: bool = True):
def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True):
try:
# return id
if stream:
if in_chat:
yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
# select datasource if datasource is none
if not llm_service.ds:
ds_res = llm_service.select_datasource()
if stream:
if in_chat:
for chunk in ds_res:
yield orjson.dumps({'content': chunk, 'type': 'datasource-result'}).decode() + '\n\n'
yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
@@ -565,22 +566,24 @@ def run_task(llm_service: LLMService, session: SessionDep, stream: bool = True):
full_sql_text = ''
for chunk in sql_res:
full_sql_text += chunk
if stream:
if in_chat:
yield orjson.dumps({'content': chunk, 'type': 'sql-result'}).decode() + '\n\n'
if stream:
if in_chat:
yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
# filter sql
print(full_sql_text)
sql = llm_service.check_save_sql(res=full_sql_text)
print(sql)
if stream:
if in_chat:
yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n'
else:
yield f'```sql\n{sql}\n```\n\n'
# execute sql
result = llm_service.execute_sql(sql=sql)
llm_service.save_sql_data(data_obj=result)
if stream:
if in_chat:
yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n'
# generate chart
@@ -588,40 +591,56 @@ def run_task(llm_service: LLMService, session: SessionDep, stream: bool = True):
full_chart_text = ''
for chunk in chart_res:
full_chart_text += chunk
if stream:
if in_chat:
yield orjson.dumps({'content': chunk, 'type': 'chart-result'}).decode() + '\n\n'
if stream:
if in_chat:
yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
# filter chart
print(full_chart_text)
chart = llm_service.check_save_chart(res=full_chart_text)
print(chart)
if stream:
if in_chat:
yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
else:
data = []
_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')
_fields_list = []
_fields_skip = False
for _data in result.get('data'):
_row = []
for field in result.get('fields'):
_row.append(_data.get(field))
if not _fields_skip:
_fields_list.append(field if not _fields.get(field) else _fields.get(field))
data.append(_row)
_fields_skip = True
df = pd.DataFrame(np.array(data), columns=_fields_list)
markdown_table = df.to_markdown(index=False)
yield markdown_table + '\n\n'
record = llm_service.finish()
if stream:
if in_chat:
yield orjson.dumps({'type': 'finish'}).decode() + '\n\n'
else:
md_str = f'```sql\n{sql}\n```\n\n'
# todo generate picture
if chart['type'] == 'table':
data = {}
for _data in result['data']:
for field in result['fields']:
if not data[field]:
data[field] = []
data[field].append(_data[field])
df = pd.DataFrame(data, columns=result['fields'])
markdown_table = df.to_markdown(index=False)
md_str += markdown_table
else:
md_str += ''
if chart['type'] != 'table':
yield '# todo generate chart picture'
except Exception as e:
llm_service.save_error(message=str(e))
if stream:
if in_chat:
yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
else:
raise e
backend/apps/mcp/mcp.py
+2 -19
View File
@@ -5,6 +5,7 @@ from datetime import timedelta
from typing import Annotated
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.security import OAuth2PasswordRequestForm
from apps.chat.api.chat import create_chat
@@ -59,22 +60,4 @@ async def mcp_question(session: SessionDep, chat: ChatMcp):
llm_service = LLMService(session, user, chat)
llm_service.init_record()
run_task(llm_service, session)
# return await stream_sql(session, user, chat)
return {"content": """这是一段写死的测试内容：
步骤1: 确定需要查询的字段。
我们需要统计上海的订单总数，因此需要从"城市"字段中筛选出值为"上海"的记录，并使用COUNT函数计算这些记录的数量。
步骤2: 确定筛选条件。
问题要求统计上海的订单总数，所以我们需要在SQL语句中添加WHERE "城市" = '上海'来筛选出符合条件的记录。
步骤3: 避免关键字冲突。
因为这个Excel/CSV数据库是 PostgreSQL 类型，所以在schema、表名、字段名和别名外层加双引号。
最终答案:
```json
{"success":true,"sql":"SELECT COUNT(*) AS \"TotalOrders\" FROM \"public\".\"Sheet1_c27345b66e\" WHERE \"城市\" = '上海';"}
```
<img src="https://sqlbot.fit2cloud.cn/images/111.png">"""}
return StreamingResponse(run_task(llm_service, session, False), media_type="text/event-stream")
backend/pyproject.toml
+1
View File
@@ -35,6 +35,7 @@ dependencies = [
"oracledb (>=3.1.1,<4.0.0)",
"pyyaml (>=6.0.2,<7.0.0)",
"fastapi-mcp (>=0.3.4,<0.4.0)",
"tabulate>=0.9.0",
"sqlbot-xpack==0.0.3.1",
]
[[tool.uv.index]]