AI问数
This commit is contained in:
BBIT-Kai
@@ -0,0 +1,430 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user