feat: export chart

Author: ulleo

backend/apps/chat/api/chat.py

@@ -1,12 +1,16 @@
+import datetime
 import traceback
 
+import numpy as np
+import pandas as pd
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 
 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
-from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion
+from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion, ExcelData
 from apps.chat.task.llm import LLMService
+from common.core.config import settings
 from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
 
 router = APIRouter(tags=["Data Q&A"], prefix="/chat")
@@ -183,3 +187,27 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
         )
 
     return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
+
+
+@router.post("/excel/export")
+async def export_excel(excel_data: ExcelData):
+    _fields_list = []
+    data = []
+    for _data in excel_data.data:
+        _row = []
+        for field in excel_data.axis:
+            _row.append(_data.get(field.value))
+        data.append(_row)
+    for field in excel_data.axis:
+        _fields_list.append(field.name)
+    df = pd.DataFrame(np.array(data), columns=_fields_list)
+
+    file_name = f"{excel_data.name}-{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.xlsx"
+
+    file_path = f'{(settings.EXCEL_PATH if settings.EXCEL_PATH[-1] == "/" else (settings.EXCEL_PATH + "/"))}{file_name}'
+
+    file_download_path = f'{(settings.SERVER_EXCEL_HOST if settings.SERVER_EXCEL_HOST[-1] == "/" else (settings.SERVER_EXCEL_HOST + "/"))}{file_name}'
+
+    df.to_excel(file_path, index=False)
+
+    return file_download_path

backend/apps/chat/models/chat_model.py

@@ -180,3 +180,15 @@ class McpQuestion(BaseModel):
     question: str = Body(description='用户提问')
     chat_id: int = Body(description='会话ID')
     token: str = Body(description='token')
+
+
+class AxisObj(BaseModel):
+    name: str = ''
+    value: str = ''
+    type: str | None = None
+
+
+class ExcelData(BaseModel):
+    axis: list[AxisObj] = []
+    data: list[dict] = []
+    name: str = 'Excel'

backend/apps/datasource/crud/datasource.py

@@ -4,10 +4,9 @@
 
 from fastapi import HTTPException
 from sqlalchemy import and_, text, func
-
 from sqlmodel import select
 
-
+from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
 from apps.datasource.utils.utils import aes_decrypt
 from apps.db.constant import DB
 from apps.db.db import get_engine, get_tables, get_fields, exec_sql
@@ -20,7 +19,7 @@
 from ..crud.table import delete_table_by_ds_id, update_table
 from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, TableObj, \
     DatasourceConf, TableAndFields
-from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
+
 
 def get_datasource_list(session: SessionDep, user: CurrentUser, oid: Optional[int] = None) -> List[CoreDatasource]:
     current_oid = user.oid if user.oid is not None else 1
@@ -249,11 +248,13 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table
     f_list = [f for f in data.fields if f.checked]
     if is_normal_user(current_user):
         # column is checked, and, column permission for data.fields
-        f_list = get_column_permission_fields(session=session, current_user=current_user, table=data.table, fields=f_list) or f_list
+        f_list = get_column_permission_fields(session=session, current_user=current_user, table=data.table,
+                                              fields=f_list) or f_list
 
         # row permission tree
         where_str = ''
-        filter_mapping = get_row_permission_filters(session=session, current_user=current_user, ds=ds, tables=None, single_table=data.table)
+        filter_mapping = get_row_permission_filters(session=session, current_user=current_user, ds=ds, tables=None,
+                                                    single_table=data.table)
         if filter_mapping:
             mapping_dict = filter_mapping[0]
             where_str = mapping_dict.get('filter')
@@ -325,13 +326,12 @@ def get_table_obj_by_ds(session: SessionDep, current_user: CurrentUser, ds: Core
         fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all()
 
         # do column permissions, filter fields
-        fields = get_column_permission_fields(session=session, current_user=current_user, table=table, fields=fields) or fields
+        fields = get_column_permission_fields(session=session, current_user=current_user, table=table,
+                                              fields=fields) or fields
         _list.append(TableAndFields(schema=schema, table=table, fields=fields))
     return _list
 
 
-    
-
 def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource) -> str:
     schema_str = ""
     table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
@@ -360,4 +360,5 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
                 field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})")
         schema_str += ",\n".join(field_list)
         schema_str += '\n]\n'
+    # todo 外键
     return schema_str

backend/common/core/config.py

@@ -79,8 +79,10 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
         )
 
     MCP_IMAGE_PATH: str = '/opt/sqlbot/images'
+    EXCEL_PATH: str
     MCP_IMAGE_HOST: str = 'http://localhost:3000'
     SERVER_IMAGE_HOST: str
+    SERVER_EXCEL_HOST: str
 
 settings = Settings()  # type: ignore
 print(settings)

backend/main.py

@@ -44,6 +44,9 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     lifespan=lifespan
 )
 
+app.mount("/excel", StaticFiles(directory=settings.EXCEL_PATH), name="excel")
+
+
 mcp_app = FastAPI()
 # mcp server, images path
 mcp_app.mount("/images", StaticFiles(directory=settings.MCP_IMAGE_PATH), name="images")

frontend/src/api/chat.ts

@@ -329,4 +329,5 @@ export const chatApi = {
     return request.fetchStream(`/chat/recommend_questions/${record_id}`, {}, controller)
   },
   checkLLMModel: () => request.get('/system/aimodel/default', { requestOptions: { silent: true } }),
+  export2Excel: (data: any) => request.post('/chat/excel/export', data),
 }