import asyncio import json from typing import Any from uuid import UUID from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile from fastapi.responses import Response, StreamingResponse import db.sqlserver as sqlserver from config.cubeReport import get_cube_report_settings from config.security import get_user_id_from_token from db.postgres.cube_report import ( delete_cube_report_state, get_cube_report_state, list_cube_report_states, rename_cube_report_state, upsert_cube_report_state, ) from models.BaseResponse import BaseResponse from models.CubeReportRequest import ( RenameCubeReportConversationRequest, SendCubeReportMessageRequest, ) from service.cube_report import ( delete_dify_conversation, get_dify_conversation, get_dify_conversation_variables, get_dify_messages, get_dify_parameters, get_latest_cube_response, list_dify_conversations, normalize_cube_load, normalize_cube_response, preview_dify_file, query_cube_export, query_cube_page, rename_dify_conversation, stop_dify_task, stream_dify_chat, upload_dify_file, ) from routers.dify_export_router import build_xlsx, sanitize_filename, upload_xlsx cubeReportRouter = APIRouter(prefix="/cube-report") def _sse(payload: dict[str, Any]) -> str: return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" def _scope(user_id: UUID) -> dict[str, Any]: scope = sqlserver.get_user_company_scope(str(user_id)) if not scope["companies"]: raise HTTPException(status_code=403, detail="当前用户没有可用的数据租户") return scope def _validate_scope( user_id: UUID, *, allow_global: bool, tenant_id: str | None, ) -> tuple[str | None, str]: scope = _scope(user_id) if allow_global: if not scope["canGlobal"]: raise HTTPException(status_code=403, detail="当前用户没有全局查询权限") return None, "不限租户" companies = {item["id"]: item["name"] for item in scope["companies"]} if not tenant_id or tenant_id not in companies: raise HTTPException(status_code=403, detail="无权访问所选租户") return tenant_id, companies[tenant_id] async def _save_cube_state( *, conversation_id: str, user_id: UUID, session: dict[str, Any], cube_response: dict[str, Any], ) -> None: await asyncio.to_thread( upsert_cube_report_state, conversation_id=conversation_id, user_id=str(user_id), title=str(session.get("title") or "新数据对话"), tenant_id=session.get("tenantId"), tenant_name=str(session.get("tenantName") or ""), allow_global=bool(session.get("allowGlobal")), query=cube_response["query"], limit_source=str(cube_response.get("limitSource") or "system"), response_message=str(cube_response.get("message") or ""), ) async def _resolve_cube_state( conversation_id: str, user_id: UUID, session: dict[str, Any], ) -> dict[str, Any] | None: state = await asyncio.to_thread( get_cube_report_state, conversation_id, str(user_id) ) if state: return state cube_response = await get_latest_cube_response(conversation_id, str(user_id)) if not cube_response or not cube_response["status"]: variables = await get_dify_conversation_variables( conversation_id, str(user_id) ) query = normalize_cube_load(variables.get("temp_request")) if not query: return None cube_response = { "query": query, "limitSource": str(variables.get("temp_limit_source") or "system"), "message": "", "status": True, } await _save_cube_state( conversation_id=conversation_id, user_id=user_id, session=session, cube_response=cube_response, ) return await asyncio.to_thread( get_cube_report_state, conversation_id, str(user_id) ) @cubeReportRouter.get("/scope") def get_scope(user_id: UUID = Depends(get_user_id_from_token)): return BaseResponse(data=_scope(user_id)) @cubeReportRouter.get("/app-parameters") async def get_app_parameters(user_id: UUID = Depends(get_user_id_from_token)): _scope(user_id) return BaseResponse(data=await get_dify_parameters()) @cubeReportRouter.get("/sessions") async def list_sessions(user_id: UUID = Depends(get_user_id_from_token)): conversations, states = await asyncio.gather( list_dify_conversations(str(user_id)), asyncio.to_thread(list_cube_report_states, str(user_id)), ) for item in conversations: item["hasData"] = item["id"] in states return BaseResponse(data=conversations) @cubeReportRouter.get("/sessions/{conversation_id}") async def get_session( conversation_id: str, user_id: UUID = Depends(get_user_id_from_token), ): session, messages = await asyncio.gather( get_dify_conversation(conversation_id, str(user_id)), get_dify_messages(conversation_id, str(user_id)), ) _validate_scope( user_id, allow_global=session["allowGlobal"], tenant_id=session["tenantId"], ) state = await _resolve_cube_state(conversation_id, user_id, session) session["hasData"] = bool(state and state.get("query")) return BaseResponse(data={"session": session, "messages": messages}) @cubeReportRouter.get("/sessions/{conversation_id}/data") async def get_session_data( conversation_id: str, page: int = Query(default=1, ge=1), pageSize: int = Query(default=20, ge=1, le=100), user_id: UUID = Depends(get_user_id_from_token), ): session = await get_dify_conversation(conversation_id, str(user_id)) tenant_id, _ = _validate_scope( user_id, allow_global=session["allowGlobal"], tenant_id=session["tenantId"], ) state = await _resolve_cube_state(conversation_id, user_id, session) if not state or not state.get("query"): return BaseResponse( data={ "columns": [], "rows": [], "page": page, "pageSize": pageSize, "total": 0, "sql": "", } ) data = await query_cube_page( base_load=state["query"], tenant_id=tenant_id, allow_global=session["allowGlobal"], page=page, page_size=pageSize, limit_source=state["limitSource"], ) data["title"] = session.get("title") or state.get("title") return BaseResponse(data=data) @cubeReportRouter.post("/sessions/{conversation_id}/export") async def export_session_data( conversation_id: str, user_id: UUID = Depends(get_user_id_from_token), ): session = await get_dify_conversation(conversation_id, str(user_id)) tenant_id, _ = _validate_scope( user_id, allow_global=session["allowGlobal"], tenant_id=session["tenantId"], ) state = await _resolve_cube_state(conversation_id, user_id, session) if not state or not state.get("query"): raise HTTPException(status_code=404, detail="当前会话还没有可导出的数据请求") export_data = await query_cube_export( base_load=state["query"], tenant_id=tenant_id, allow_global=session["allowGlobal"], limit_source=state["limitSource"], ) filename = sanitize_filename(session.get("title") or "Cube查询结果") contents, sheet_count = await asyncio.to_thread( build_xlsx, export_data["rows"], export_data["headers"] ) object_name, url = await asyncio.to_thread(upload_xlsx, contents, filename) return BaseResponse( data={ "url": url, "filename": filename, "bucketName": "dify-export", "objectName": object_name, "rowCount": len(export_data["rows"]), "sheetCount": sheet_count, } ) @cubeReportRouter.post("/messages/stream") async def send_message_stream( request: SendCubeReportMessageRequest, user_id: UUID = Depends(get_user_id_from_token), ): conversation_id = request.conversationId if conversation_id: session = await get_dify_conversation(conversation_id, str(user_id)) allow_global = session["allowGlobal"] tenant_id, tenant_name = _validate_scope( user_id, allow_global=allow_global, tenant_id=session["tenantId"], ) else: allow_global = request.allowGlobal tenant_id, tenant_name = _validate_scope( user_id, allow_global=allow_global, tenant_id=request.tenantId, ) session = { "title": "新数据对话", "tenantId": tenant_id, "tenantName": tenant_name, "allowGlobal": allow_global, } files = [item.model_dump() for item in request.files] inputs = { "allow_global": allow_global, "tenant_id": "__ALL__" if allow_global else tenant_id, "tenant_name": "不限租户" if allow_global else tenant_name, } async def event_stream(): active_conversation_id = conversation_id answer = "" task_id = None workflow_status_sent = False try: yield _sse({"type": "status", "text": "正在理解问题…"}) async for event in stream_dify_chat( user_id=str(user_id), content=request.content.strip(), conversation_id=conversation_id, inputs=inputs, files=files, ): event_type = event.get("event") active_conversation_id = ( event.get("conversation_id") or active_conversation_id ) current_task_id = event.get("task_id") if current_task_id and current_task_id != task_id: task_id = current_task_id yield _sse({"type": "task", "taskId": task_id}) if event_type in {"message", "agent_message"}: delta = str(event.get("answer") or "") answer += delta elif event_type == "message_replace": answer = str(event.get("answer") or "") elif ( event_type in {"workflow_started", "node_started"} and not workflow_status_sent ): workflow_status_sent = True yield _sse({"type": "status", "text": "正在生成数据请求…"}) elif event_type == "error": raise RuntimeError( str(event.get("message") or "Dify 工作流执行失败") ) cube_response = normalize_cube_response(answer) has_data = bool(cube_response and cube_response["status"]) if has_data and active_conversation_id and cube_response: await _save_cube_state( conversation_id=active_conversation_id, user_id=user_id, session=session, cube_response=cube_response, ) display_message = ( cube_response["message"] if cube_response else answer or "工作流没有返回可识别的数据请求" ) yield _sse({"type": "message_replace", "content": display_message}) yield _sse( { "type": "complete", "conversationId": active_conversation_id, "hasData": has_data, } ) except Exception as error: message = error.detail if isinstance(error, HTTPException) else str(error) yield _sse({"type": "error", "message": message}) settings = get_cube_report_settings() return StreamingResponse( event_stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", "X-Page-Size": str(settings.default_page_size), }, ) @cubeReportRouter.post("/tasks/{task_id}/stop") async def stop_task( task_id: str, user_id: UUID = Depends(get_user_id_from_token), ): return BaseResponse(data=await stop_dify_task(task_id, str(user_id))) @cubeReportRouter.post("/files/upload") async def upload_file( file: UploadFile = File(...), user_id: UUID = Depends(get_user_id_from_token), ): filename = (file.filename or "").lower() if not filename.endswith((".xls", ".xlsx")): raise HTTPException(status_code=415, detail="当前仅支持 XLS、XLSX 文件") return BaseResponse(data=await upload_dify_file(file, str(user_id))) @cubeReportRouter.post("/sessions/{conversation_id}/name") async def rename_session( conversation_id: str, request: RenameCubeReportConversationRequest, user_id: UUID = Depends(get_user_id_from_token), ): await get_dify_conversation(conversation_id, str(user_id)) data = await rename_dify_conversation( conversation_id, str(user_id), request.name.strip() ) await asyncio.to_thread( rename_cube_report_state, conversation_id, str(user_id), request.name.strip(), ) return BaseResponse(data=data) @cubeReportRouter.delete("/sessions/{conversation_id}") async def delete_session( conversation_id: str, user_id: UUID = Depends(get_user_id_from_token), ): await get_dify_conversation(conversation_id, str(user_id)) await delete_dify_conversation(conversation_id, str(user_id)) await asyncio.to_thread( delete_cube_report_state, conversation_id, str(user_id) ) return BaseResponse(data={"result": "success"}) @cubeReportRouter.get("/files/{file_id}/preview") async def preview_file( file_id: str, conversationId: str = Query(min_length=1), asAttachment: bool = Query(default=False), user_id: UUID = Depends(get_user_id_from_token), ): messages = await get_dify_messages(conversationId, str(user_id)) owned = any( file["id"] == file_id for message in messages for file in message.get("files", []) ) if not owned: raise HTTPException(status_code=404, detail="附件不存在或无权访问") content, headers = await preview_dify_file( file_id, as_attachment=asAttachment ) media_type = headers.pop("content-type", "application/octet-stream") return Response(content=content, media_type=media_type, headers=headers)