feat: chat question for mcp

Author: ulleo

backend/apps/chat/task/llm.py

@@ -2,6 +2,7 @@
 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,63 +566,81 @@ 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
         chart_res = llm_service.generate_chart()
         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

@@ -5,6 +5,7 @@
 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

@@ -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]]