AI问数
This commit is contained in:
BBIT-Kai
@@ -11,6 +11,7 @@ from models.BaseResponse import BaseResponse
|
||||
from routers.AnnualMeeting import amRouter
|
||||
from routers.Bot import botRouter
|
||||
from routers.Chat import chatRouter
|
||||
from routers.CubeReport import cubeReportRouter
|
||||
from routers.Datasource import reportDataRouter
|
||||
from routers.Iot import iot_router
|
||||
from routers.Knowledge import knowledgeRouter
|
||||
@@ -57,6 +58,7 @@ async def ai_lab():
|
||||
)
|
||||
routers = [
|
||||
chatRouter,
|
||||
cubeReportRouter,
|
||||
reportRouter,
|
||||
knowledgeRouter,
|
||||
reportDataRouter,
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CubeReportSettings:
|
||||
dify_api_base: str
|
||||
dify_api_key: str
|
||||
cube_api_base: str
|
||||
cube_api_token: str
|
||||
request_timeout_seconds: float
|
||||
default_page_size: int
|
||||
max_page_size: int
|
||||
|
||||
|
||||
def get_cube_report_settings() -> CubeReportSettings:
|
||||
return CubeReportSettings(
|
||||
dify_api_base=os.getenv(
|
||||
"DIFY_DATABASE_ASSISTANT_API_BASE", "https://chat.bbitcn.net/v1"
|
||||
).rstrip("/"),
|
||||
dify_api_key=os.getenv(
|
||||
"DIFY_DATABASE_ASSISTANT_API_KEY", "app-uibWo8ZEpqHCsWXREPTCBDH6"
|
||||
),
|
||||
cube_api_base=os.getenv(
|
||||
"CUBE_API_BASE_URL", "http://10.10.12.101:4001/cubejs-api/v1"
|
||||
).rstrip("/"),
|
||||
cube_api_token=os.getenv("CUBE_API_TOKEN", ""),
|
||||
request_timeout_seconds=float(
|
||||
os.getenv("CUBE_REPORT_REQUEST_TIMEOUT_SECONDS", "180")
|
||||
),
|
||||
default_page_size=int(os.getenv("CUBE_REPORT_DEFAULT_PAGE_SIZE", "20")),
|
||||
max_page_size=int(os.getenv("CUBE_REPORT_MAX_PAGE_SIZE", "100")),
|
||||
)
|
||||
@@ -0,0 +1,134 @@
|
||||
from typing import Any
|
||||
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from config.pgDb import pg_pool
|
||||
|
||||
|
||||
def list_cube_report_states(user_id: str) -> dict[str, dict[str, Any]]:
|
||||
with pg_pool.getConn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT dify_conversation_id, current_cube, limit_source,
|
||||
response_message, updated_at
|
||||
FROM ai_cube_report_states
|
||||
WHERE created_by = %s
|
||||
""",
|
||||
(user_id,),
|
||||
)
|
||||
return {
|
||||
str(row[0]): {
|
||||
"query": row[1],
|
||||
"limitSource": row[2],
|
||||
"message": row[3],
|
||||
"updatedAt": row[4].isoformat() if row[4] else None,
|
||||
}
|
||||
for row in cur.fetchall()
|
||||
}
|
||||
|
||||
|
||||
def get_cube_report_state(
|
||||
conversation_id: str, user_id: str
|
||||
) -> dict[str, Any] | None:
|
||||
with pg_pool.getConn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT title, tenant_id, tenant_name, allow_global,
|
||||
current_cube, limit_source, response_message,
|
||||
created_at, updated_at
|
||||
FROM ai_cube_report_states
|
||||
WHERE dify_conversation_id = %s AND created_by = %s
|
||||
""",
|
||||
(conversation_id, user_id),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
return {
|
||||
"conversationId": conversation_id,
|
||||
"title": row[0],
|
||||
"tenantId": str(row[1]) if row[1] is not None else None,
|
||||
"tenantName": row[2],
|
||||
"allowGlobal": row[3],
|
||||
"query": row[4],
|
||||
"limitSource": row[5],
|
||||
"message": row[6],
|
||||
"createdAt": row[7].isoformat() if row[7] else None,
|
||||
"updatedAt": row[8].isoformat() if row[8] else None,
|
||||
}
|
||||
|
||||
|
||||
def upsert_cube_report_state(
|
||||
*,
|
||||
conversation_id: str,
|
||||
user_id: str,
|
||||
title: str,
|
||||
tenant_id: str | None,
|
||||
tenant_name: str,
|
||||
allow_global: bool,
|
||||
query: dict[str, Any],
|
||||
limit_source: str,
|
||||
response_message: str,
|
||||
) -> None:
|
||||
normalized_source = "user" if limit_source == "user" else "system"
|
||||
with pg_pool.getConn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO ai_cube_report_states (
|
||||
dify_conversation_id, created_by, title, tenant_id,
|
||||
tenant_name, allow_global, current_cube, limit_source,
|
||||
response_message, created_at, updated_at
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
|
||||
ON CONFLICT (dify_conversation_id) DO UPDATE SET
|
||||
title = EXCLUDED.title,
|
||||
tenant_id = EXCLUDED.tenant_id,
|
||||
tenant_name = EXCLUDED.tenant_name,
|
||||
allow_global = EXCLUDED.allow_global,
|
||||
current_cube = EXCLUDED.current_cube,
|
||||
limit_source = EXCLUDED.limit_source,
|
||||
response_message = EXCLUDED.response_message,
|
||||
updated_at = NOW()
|
||||
WHERE ai_cube_report_states.created_by = EXCLUDED.created_by
|
||||
""",
|
||||
(
|
||||
conversation_id,
|
||||
user_id,
|
||||
title[:200] or "新数据对话",
|
||||
tenant_id,
|
||||
tenant_name[:200],
|
||||
allow_global,
|
||||
Jsonb(query),
|
||||
normalized_source,
|
||||
response_message,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def rename_cube_report_state(
|
||||
conversation_id: str, user_id: str, title: str
|
||||
) -> None:
|
||||
with pg_pool.getConn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE ai_cube_report_states
|
||||
SET title = %s, updated_at = NOW()
|
||||
WHERE dify_conversation_id = %s AND created_by = %s
|
||||
""",
|
||||
(title[:200], conversation_id, user_id),
|
||||
)
|
||||
|
||||
|
||||
def delete_cube_report_state(conversation_id: str, user_id: str) -> None:
|
||||
with pg_pool.getConn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM ai_cube_report_states
|
||||
WHERE dify_conversation_id = %s AND created_by = %s
|
||||
""",
|
||||
(conversation_id, user_id),
|
||||
)
|
||||
@@ -14,6 +14,21 @@ def executeSQL(sql: str):
|
||||
return [dict(row._mapping) for row in result]
|
||||
|
||||
|
||||
def execute_parameterized_sql(sql: str, params: dict | None = None):
|
||||
"""执行参数化只读 SQL，返回列名和字典行。"""
|
||||
with mssql_pool.getConn() as conn:
|
||||
result = conn.execute(text(sql), params or {})
|
||||
columns = list(result.keys())
|
||||
rows = [dict(row._mapping) for row in result]
|
||||
return columns, rows
|
||||
|
||||
|
||||
def execute_parameterized_scalar(sql: str, params: dict | None = None):
|
||||
with mssql_pool.getConn() as conn:
|
||||
result = conn.execute(text(sql), params or {})
|
||||
return result.scalar()
|
||||
|
||||
|
||||
def get_company_list(user_id: str):
|
||||
# 1️⃣ 从 PostgreSQL 获取 tenant_id
|
||||
with pg_pool.getConn() as pg_conn:
|
||||
@@ -35,3 +50,33 @@ def get_company_list(user_id: str):
|
||||
with mssql_pool.getConn() as mssql_conn:
|
||||
result = mssql_conn.execute(query, params)
|
||||
return [{"id": str(row[0]), "name": row[1]} for row in result.fetchall()]
|
||||
|
||||
|
||||
def get_user_company_scope(user_id: str):
|
||||
"""返回数据助手使用的租户范围；用户不存在时绝不能按全局用户处理。"""
|
||||
with pg_pool.getConn() as pg_conn:
|
||||
with pg_conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT bbit_tenant_id FROM sys_users WHERE id = %s", (user_id,)
|
||||
)
|
||||
row = cur.fetchone()
|
||||
|
||||
if not row:
|
||||
return {"canGlobal": False, "companies": []}
|
||||
|
||||
tenant_id = row[0]
|
||||
if tenant_id:
|
||||
query = text("SELECT Id, Name FROM dbo.POC_TENANTS WHERE Id = :tenant_id")
|
||||
params = {"tenant_id": tenant_id}
|
||||
can_global = False
|
||||
else:
|
||||
query = text("SELECT Id, Name FROM dbo.POC_TENANTS ORDER BY Name")
|
||||
params = {}
|
||||
can_global = True
|
||||
|
||||
with mssql_pool.getConn() as mssql_conn:
|
||||
result = mssql_conn.execute(query, params)
|
||||
companies = [
|
||||
{"id": str(item[0]), "name": item[1]} for item in result.fetchall()
|
||||
]
|
||||
return {"canGlobal": can_global, "companies": companies}
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
|
||||
class CubeReportFileInput(BaseModel):
|
||||
id: str = Field(min_length=1)
|
||||
name: str = Field(min_length=1, max_length=255)
|
||||
type: Literal["custom", "document", "image"] = "custom"
|
||||
|
||||
|
||||
class SendCubeReportMessageRequest(BaseModel):
|
||||
content: str = Field(default="", max_length=8000)
|
||||
conversationId: str | None = None
|
||||
allowGlobal: bool = False
|
||||
tenantId: str | None = Field(default=None, max_length=64)
|
||||
tenantName: str | None = Field(default=None, max_length=200)
|
||||
files: list[CubeReportFileInput] = Field(default_factory=list, max_length=5)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_content(self):
|
||||
if not self.content.strip() and not self.files:
|
||||
raise ValueError("消息内容和附件不能同时为空")
|
||||
return self
|
||||
|
||||
|
||||
class RenameCubeReportConversationRequest(BaseModel):
|
||||
name: str = Field(min_length=1, max_length=200)
|
||||
@@ -1,4 +1,5 @@
|
||||
fastapi==0.116.1
|
||||
httpx>=0.28.0
|
||||
langchain==0.3.27
|
||||
langchain_community==0.3.29
|
||||
langchain_milvus==0.2.1
|
||||
|
||||
@@ -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)
|
||||
@@ -0,0 +1,821 @@
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from collections.abc import AsyncIterator
|
||||
from copy import deepcopy
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException, UploadFile
|
||||
|
||||
from config.cubeReport import get_cube_report_settings
|
||||
from db.sqlserver import execute_parameterized_scalar, execute_parameterized_sql
|
||||
|
||||
|
||||
def _require_dify_key() -> str:
|
||||
key = get_cube_report_settings().dify_api_key
|
||||
if not key:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="尚未配置 DIFY_DATABASE_ASSISTANT_API_KEY",
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
def _headers(token: str, content_type: str | None = "application/json") -> dict[str, str]:
|
||||
headers: dict[str, str] = {}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
if content_type:
|
||||
headers["Content-Type"] = content_type
|
||||
return headers
|
||||
|
||||
|
||||
def _timeout() -> httpx.Timeout:
|
||||
seconds = get_cube_report_settings().request_timeout_seconds
|
||||
return httpx.Timeout(seconds, connect=15.0)
|
||||
|
||||
|
||||
def _raise_dify_error(response: httpx.Response) -> None:
|
||||
if not response.is_error:
|
||||
return
|
||||
try:
|
||||
payload = response.json()
|
||||
detail = payload.get("message") or payload.get("detail") or response.text
|
||||
except (ValueError, TypeError):
|
||||
detail = response.text
|
||||
raise HTTPException(
|
||||
status_code=response.status_code,
|
||||
detail=f"Dify 请求失败：{str(detail)[:500]}",
|
||||
)
|
||||
|
||||
|
||||
def _iso_timestamp(value: Any) -> str | None:
|
||||
try:
|
||||
return datetime.fromtimestamp(float(value), tz=timezone.utc).isoformat()
|
||||
except (TypeError, ValueError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def _as_bool(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return str(value).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _decode_json_objects(text: str) -> list[Any]:
|
||||
"""从混合文本中提取完整 JSON，兼容工作流误输出连续多个对象。"""
|
||||
decoder = json.JSONDecoder()
|
||||
values: list[Any] = []
|
||||
index = 0
|
||||
while index < len(text):
|
||||
object_start = text.find("{", index)
|
||||
array_start = text.find("[", index)
|
||||
starts = [item for item in (object_start, array_start) if item >= 0]
|
||||
if not starts:
|
||||
break
|
||||
start = min(starts)
|
||||
try:
|
||||
value, end = decoder.raw_decode(text, start)
|
||||
except json.JSONDecodeError:
|
||||
index = start + 1
|
||||
continue
|
||||
values.append(value)
|
||||
index = end
|
||||
return values
|
||||
|
||||
|
||||
def normalize_cube_response(value: Any) -> dict[str, Any] | None:
|
||||
"""解析工作流输出；多个有效 JSON 连续出现时，以最后一个为准。"""
|
||||
current = value
|
||||
if isinstance(current, str):
|
||||
text = current.strip()
|
||||
if text.startswith("```"):
|
||||
text = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
|
||||
text = re.sub(r"\s*```$", "", text)
|
||||
try:
|
||||
current = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
candidates = [
|
||||
item
|
||||
for item in _decode_json_objects(text)
|
||||
if isinstance(item, dict) and isinstance(item.get("query"), dict)
|
||||
]
|
||||
if not candidates:
|
||||
return None
|
||||
current = candidates[-1]
|
||||
if not isinstance(current, dict) or not isinstance(current.get("query"), dict):
|
||||
return None
|
||||
return {
|
||||
"query": current["query"],
|
||||
"message": str(current.get("message") or "数据请求已生成"),
|
||||
"status": _as_bool(current.get("status", True)),
|
||||
"limitSource": str(current.get("limit_source") or "system"),
|
||||
}
|
||||
|
||||
|
||||
async def get_dify_parameters() -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/parameters",
|
||||
headers=_headers(key),
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
payload = response.json()
|
||||
questions = payload.get("suggested_questions", [])
|
||||
questions = questions if isinstance(questions, list) else []
|
||||
return {
|
||||
"openingStatement": str(
|
||||
payload.get("opening_statement") or payload.get("introduction") or ""
|
||||
),
|
||||
"suggestedQuestions": [str(item) for item in questions if item],
|
||||
"suggestedQuestionsAfterAnswer": payload.get(
|
||||
"suggested_questions_after_answer", {"enabled": False}
|
||||
),
|
||||
"fileUpload": payload.get("file_upload", {}),
|
||||
"systemParameters": payload.get("system_parameters", {}),
|
||||
"userInputForm": payload.get("user_input_form", []),
|
||||
}
|
||||
|
||||
|
||||
def normalize_conversation(item: dict[str, Any]) -> dict[str, Any]:
|
||||
inputs = item.get("inputs") if isinstance(item.get("inputs"), dict) else {}
|
||||
allow_global = _as_bool(inputs.get("allow_global"))
|
||||
return {
|
||||
"id": str(item.get("id") or ""),
|
||||
"title": str(item.get("name") or "新数据对话"),
|
||||
"tenantId": None if allow_global else str(inputs.get("tenant_id") or ""),
|
||||
"tenantName": "不限租户"
|
||||
if allow_global
|
||||
else str(inputs.get("tenant_name") or "未指定租户"),
|
||||
"allowGlobal": allow_global,
|
||||
"createdAt": _iso_timestamp(item.get("created_at")),
|
||||
"updatedAt": _iso_timestamp(item.get("updated_at")),
|
||||
"inputs": inputs,
|
||||
}
|
||||
|
||||
|
||||
async def list_dify_conversations(
|
||||
user_id: str, *, max_items: int = 500
|
||||
) -> list[dict[str, Any]]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
conversations: list[dict[str, Any]] = []
|
||||
last_id: str | None = None
|
||||
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
while len(conversations) < max_items:
|
||||
params: dict[str, Any] = {
|
||||
"user": user_id,
|
||||
"limit": min(100, max_items - len(conversations)),
|
||||
"sort_by": "-updated_at",
|
||||
}
|
||||
if last_id:
|
||||
params["last_id"] = last_id
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/conversations",
|
||||
headers=_headers(key),
|
||||
params=params,
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
payload = response.json()
|
||||
batch = payload.get("data", [])
|
||||
if not isinstance(batch, list) or not batch:
|
||||
break
|
||||
conversations.extend(item for item in batch if isinstance(item, dict))
|
||||
if not payload.get("has_more"):
|
||||
break
|
||||
last_id = str(batch[-1].get("id") or "")
|
||||
if not last_id:
|
||||
break
|
||||
return [normalize_conversation(item) for item in conversations]
|
||||
|
||||
|
||||
async def get_dify_conversation(
|
||||
conversation_id: str, user_id: str
|
||||
) -> dict[str, Any]:
|
||||
conversations = await list_dify_conversations(user_id)
|
||||
for item in conversations:
|
||||
if item["id"] == conversation_id:
|
||||
return item
|
||||
raise HTTPException(status_code=404, detail="Dify 会话不存在或无权访问")
|
||||
|
||||
|
||||
def _message_files(item: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
files = item.get("message_files")
|
||||
if not isinstance(files, list):
|
||||
return []
|
||||
return [
|
||||
{
|
||||
"id": str(file.get("id") or ""),
|
||||
"name": str(file.get("name") or file.get("id") or "附件"),
|
||||
"type": str(file.get("type") or "custom"),
|
||||
"belongsTo": str(file.get("belongs_to") or "user"),
|
||||
}
|
||||
for file in files
|
||||
if isinstance(file, dict) and file.get("id")
|
||||
]
|
||||
|
||||
|
||||
async def get_dify_message_records(
|
||||
conversation_id: str, user_id: str
|
||||
) -> list[dict[str, Any]]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
data: list[dict[str, Any]] = []
|
||||
first_id: str | None = None
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
for _ in range(10):
|
||||
params: dict[str, Any] = {
|
||||
"conversation_id": conversation_id,
|
||||
"user": user_id,
|
||||
"limit": 100,
|
||||
}
|
||||
if first_id:
|
||||
params["first_id"] = first_id
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/messages",
|
||||
headers=_headers(key),
|
||||
params=params,
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
payload = response.json()
|
||||
batch = payload.get("data", [])
|
||||
if not isinstance(batch, list) or not batch:
|
||||
break
|
||||
data.extend(item for item in batch if isinstance(item, dict))
|
||||
if not payload.get("has_more"):
|
||||
break
|
||||
first_id = str(batch[0].get("id") or "")
|
||||
if not first_id:
|
||||
break
|
||||
|
||||
unique_data = {str(item.get("id")): item for item in data if item.get("id")}
|
||||
data = sorted(unique_data.values(), key=lambda item: item.get("created_at", 0))
|
||||
|
||||
return data
|
||||
|
||||
|
||||
async def get_dify_messages(
|
||||
conversation_id: str, user_id: str
|
||||
) -> list[dict[str, Any]]:
|
||||
data = await get_dify_message_records(conversation_id, user_id)
|
||||
messages: list[dict[str, Any]] = []
|
||||
for item in data:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
message_id = str(item.get("id") or "")
|
||||
created_at = _iso_timestamp(item.get("created_at"))
|
||||
query = str(item.get("query") or "")
|
||||
if query or _message_files(item):
|
||||
messages.append(
|
||||
{
|
||||
"id": f"{message_id}-user",
|
||||
"role": "user",
|
||||
"content": query,
|
||||
"status": "completed",
|
||||
"difyMessageId": message_id,
|
||||
"createdAt": created_at,
|
||||
"files": [
|
||||
file
|
||||
for file in _message_files(item)
|
||||
if file["belongsTo"] == "user"
|
||||
],
|
||||
}
|
||||
)
|
||||
raw_answer = str(item.get("answer") or "")
|
||||
cube_response = normalize_cube_response(raw_answer)
|
||||
answer = cube_response["message"] if cube_response else raw_answer
|
||||
if answer:
|
||||
messages.append(
|
||||
{
|
||||
"id": f"{message_id}-assistant",
|
||||
"role": "assistant",
|
||||
"content": answer,
|
||||
"status": "completed",
|
||||
"difyMessageId": message_id,
|
||||
"createdAt": created_at,
|
||||
"files": [
|
||||
file
|
||||
for file in _message_files(item)
|
||||
if file["belongsTo"] == "assistant"
|
||||
],
|
||||
}
|
||||
)
|
||||
return messages
|
||||
|
||||
|
||||
async def get_latest_cube_response(
|
||||
conversation_id: str, user_id: str
|
||||
) -> dict[str, Any] | None:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/messages",
|
||||
headers=_headers(key),
|
||||
params={
|
||||
"conversation_id": conversation_id,
|
||||
"user": user_id,
|
||||
"limit": 20,
|
||||
},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
records = response.json().get("data", [])
|
||||
records = records if isinstance(records, list) else []
|
||||
records.sort(
|
||||
key=lambda item: item.get("created_at", 0) if isinstance(item, dict) else 0,
|
||||
reverse=True,
|
||||
)
|
||||
for item in records:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
cube_response = normalize_cube_response(item.get("answer"))
|
||||
if cube_response and cube_response["status"]:
|
||||
return cube_response
|
||||
return None
|
||||
|
||||
|
||||
async def get_dify_conversation_variables(
|
||||
conversation_id: str, user_id: str
|
||||
) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/conversations/{conversation_id}/variables",
|
||||
headers=_headers(key),
|
||||
params={"user": user_id, "limit": 100},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
payload = response.json()
|
||||
return {
|
||||
item["name"]: item.get("value")
|
||||
for item in payload.get("data", [])
|
||||
if isinstance(item, dict) and item.get("name")
|
||||
}
|
||||
|
||||
|
||||
async def stream_dify_chat(
|
||||
*,
|
||||
user_id: str,
|
||||
content: str,
|
||||
conversation_id: str | None,
|
||||
inputs: dict[str, Any],
|
||||
files: list[dict[str, Any]],
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
payload = {
|
||||
"query": content or "请分析上传的文件",
|
||||
"inputs": inputs,
|
||||
"files": [
|
||||
{
|
||||
"type": item.get("type", "custom"),
|
||||
"transfer_method": "local_file",
|
||||
"upload_file_id": item["id"],
|
||||
}
|
||||
for item in files
|
||||
],
|
||||
"response_mode": "streaming",
|
||||
"conversation_id": conversation_id or "",
|
||||
"user": user_id,
|
||||
"auto_generate_name": True,
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
f"{settings.dify_api_base}/chat-messages",
|
||||
headers=_headers(key),
|
||||
json=payload,
|
||||
) as response:
|
||||
if response.is_error:
|
||||
body = (await response.aread()).decode("utf-8", errors="replace")
|
||||
raise HTTPException(
|
||||
status_code=response.status_code,
|
||||
detail=f"Dify 请求失败：{body[:500]}",
|
||||
)
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
try:
|
||||
event = json.loads(line[6:])
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(event, dict):
|
||||
yield event
|
||||
|
||||
|
||||
async def stop_dify_task(task_id: str, user_id: str) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.post(
|
||||
f"{settings.dify_api_base}/chat-messages/{task_id}/stop",
|
||||
headers=_headers(key),
|
||||
json={"user": user_id},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
return response.json()
|
||||
|
||||
|
||||
async def upload_dify_file(file: UploadFile, user_id: str) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
content = await file.read()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.post(
|
||||
f"{settings.dify_api_base}/files/upload",
|
||||
headers=_headers(key, None),
|
||||
data={"user": user_id},
|
||||
files={
|
||||
"file": (
|
||||
file.filename or "upload.xlsx",
|
||||
content,
|
||||
file.content_type or "application/octet-stream",
|
||||
)
|
||||
},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
return response.json()
|
||||
|
||||
|
||||
async def rename_dify_conversation(
|
||||
conversation_id: str, user_id: str, name: str
|
||||
) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.post(
|
||||
f"{settings.dify_api_base}/conversations/{conversation_id}/name",
|
||||
headers=_headers(key),
|
||||
json={"name": name, "auto_generate": False, "user": user_id},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
return normalize_conversation(response.json())
|
||||
|
||||
|
||||
async def delete_dify_conversation(conversation_id: str, user_id: str) -> None:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.request(
|
||||
"DELETE",
|
||||
f"{settings.dify_api_base}/conversations/{conversation_id}",
|
||||
headers=_headers(key),
|
||||
json={"user": user_id},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
|
||||
|
||||
async def preview_dify_file(
|
||||
file_id: str, *, as_attachment: bool = False
|
||||
) -> tuple[bytes, dict[str, str]]:
|
||||
settings = get_cube_report_settings()
|
||||
key = _require_dify_key()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.get(
|
||||
f"{settings.dify_api_base}/files/{file_id}/preview",
|
||||
headers=_headers(key),
|
||||
params={"as_attachment": str(as_attachment).lower()},
|
||||
)
|
||||
_raise_dify_error(response)
|
||||
headers = {
|
||||
name: value
|
||||
for name, value in response.headers.items()
|
||||
if name.lower() in {"content-disposition", "content-length", "cache-control"}
|
||||
}
|
||||
headers["content-type"] = response.headers.get(
|
||||
"content-type", "application/octet-stream"
|
||||
)
|
||||
return response.content, headers
|
||||
|
||||
|
||||
def normalize_cube_load(value: Any) -> dict[str, Any] | None:
|
||||
response = normalize_cube_response(value)
|
||||
if response:
|
||||
return response["query"]
|
||||
current = value
|
||||
for _ in range(3):
|
||||
if isinstance(current, str):
|
||||
text = current.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
current = json.loads(text)
|
||||
except json.JSONDecodeError as error:
|
||||
raise ValueError(f"Dify 返回的 Cube Load 不是合法 JSON：{error.msg}") from error
|
||||
continue
|
||||
break
|
||||
if isinstance(current, dict) and isinstance(current.get("query"), dict):
|
||||
current = current["query"]
|
||||
return current if isinstance(current, dict) else None
|
||||
|
||||
|
||||
def _walk_filter_nodes(filters: Any):
|
||||
if not isinstance(filters, list):
|
||||
return
|
||||
for item in filters:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if "member" in item:
|
||||
yield item
|
||||
for key in ("and", "or"):
|
||||
children = item.get(key)
|
||||
if isinstance(children, list):
|
||||
yield from _walk_filter_nodes(children)
|
||||
|
||||
|
||||
def ensure_tenant_filter(query: dict[str, Any], tenant_id: str) -> None:
|
||||
tenant_member = None
|
||||
for item in _walk_filter_nodes(query.get("filters", [])):
|
||||
member = str(item.get("member", ""))
|
||||
values = [str(value) for value in item.get("values", [])]
|
||||
operator = str(item.get("operator") or "").lower()
|
||||
if (
|
||||
re.search(r"tenant|租户", member, re.IGNORECASE)
|
||||
and operator == "equals"
|
||||
and tenant_id in values
|
||||
):
|
||||
tenant_member = member
|
||||
break
|
||||
if not tenant_member:
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail="Cube Query 缺少当前租户的 equals 过滤条件，已阻止数据查询",
|
||||
)
|
||||
|
||||
# 顶层 filters 数组按 AND 组合。即使 Dify 将原租户条件放进 OR，
|
||||
# 这里也会额外增加不可绕过的租户约束。
|
||||
filters = query.get("filters")
|
||||
if not isinstance(filters, list):
|
||||
filters = []
|
||||
query["filters"] = filters
|
||||
filters.append(
|
||||
{
|
||||
"member": tenant_member,
|
||||
"operator": "equals",
|
||||
"values": [tenant_id],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _mask_value(key: str, title: str, value: Any) -> Any:
|
||||
if value is None:
|
||||
return value
|
||||
text = str(value)
|
||||
field = f"{key} {title}"
|
||||
if re.search(r"idcard|身份证", field, re.IGNORECASE) and len(text) >= 10:
|
||||
return f"{text[:6]}********{text[-4:]}"
|
||||
if re.search(r"phone|mobile|手机号", field, re.IGNORECASE) and len(text) == 11:
|
||||
return f"{text[:3]}****{text[-4:]}"
|
||||
return value
|
||||
|
||||
|
||||
def _positive_int(value: Any) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
return max(int(value), 0)
|
||||
return None
|
||||
|
||||
|
||||
def _ensure_page_order(query: dict[str, Any]) -> None:
|
||||
if query.get("order"):
|
||||
return
|
||||
dimensions = query.get("dimensions")
|
||||
if isinstance(dimensions, list) and dimensions:
|
||||
query["order"] = {str(dimensions[0]): "asc"}
|
||||
return
|
||||
time_dimensions = query.get("timeDimensions")
|
||||
if isinstance(time_dimensions, list):
|
||||
for item in time_dimensions:
|
||||
if isinstance(item, dict) and item.get("dimension"):
|
||||
query["order"] = {str(item["dimension"]): "asc"}
|
||||
return
|
||||
|
||||
|
||||
def _validate_read_only_sql(sql: str) -> None:
|
||||
normalized = re.sub(r"\s+", " ", sql).strip().lower()
|
||||
if not (normalized.startswith("select ") or normalized.startswith("with ")):
|
||||
raise HTTPException(status_code=502, detail="Cube 返回了非查询 SQL")
|
||||
if ";" in normalized.rstrip(";") or re.search(
|
||||
r"\b(insert|update|delete|drop|alter|create|truncate|merge|exec|execute)\b",
|
||||
normalized,
|
||||
):
|
||||
raise HTTPException(status_code=502, detail="Cube 返回的 SQL 未通过只读校验")
|
||||
|
||||
|
||||
async def _generate_cube_sql(query: dict[str, Any]) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
async with httpx.AsyncClient(timeout=_timeout()) as client:
|
||||
response = await client.post(
|
||||
f"{settings.cube_api_base}/sql",
|
||||
headers=_headers(settings.cube_api_token),
|
||||
json={"query": query},
|
||||
)
|
||||
if response.is_error:
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail=f"Cube SQL 生成失败：{response.text[:500]}",
|
||||
)
|
||||
payload = response.json()
|
||||
sql_node = payload.get("sql") if isinstance(payload, dict) else None
|
||||
sql_tuple = sql_node.get("sql") if isinstance(sql_node, dict) else None
|
||||
if (
|
||||
not isinstance(sql_tuple, list)
|
||||
or not sql_tuple
|
||||
or not isinstance(sql_tuple[0], str)
|
||||
):
|
||||
raise HTTPException(status_code=502, detail="Cube 没有返回可执行 SQL")
|
||||
raw_params = sql_tuple[1] if len(sql_tuple) > 1 else []
|
||||
if raw_params is None:
|
||||
raw_params = []
|
||||
elif not isinstance(raw_params, list):
|
||||
raw_params = [raw_params]
|
||||
aliases = sql_node.get("aliasNameToMember", {})
|
||||
return {
|
||||
"sql": sql_tuple[0].strip().rstrip(";"),
|
||||
"params": raw_params,
|
||||
"aliases": aliases if isinstance(aliases, dict) else {},
|
||||
}
|
||||
|
||||
|
||||
def _prepare_cube_sql(raw_sql: str, values: list[Any]) -> tuple[str, dict[str, Any]]:
|
||||
bindings: dict[str, Any] = {}
|
||||
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
index = int(match.group(1))
|
||||
if index < 1 or index > len(values):
|
||||
raise HTTPException(status_code=502, detail="Cube SQL 参数数量不匹配")
|
||||
name = f"cube_param_{index}"
|
||||
bindings[name] = values[index - 1]
|
||||
return f":{name}"
|
||||
|
||||
sql = re.sub(r"@_(\d+)", replace, raw_sql)
|
||||
_validate_read_only_sql(sql)
|
||||
return sql, bindings
|
||||
|
||||
|
||||
def _sql_literal(value: Any) -> str:
|
||||
if value is None:
|
||||
return "NULL"
|
||||
if isinstance(value, bool):
|
||||
return "1" if value else "0"
|
||||
if isinstance(value, (int, float)):
|
||||
return str(value)
|
||||
text = str(value).replace("'", "''")
|
||||
return f"N'{text}'"
|
||||
|
||||
|
||||
def _display_sql(raw_sql: str, values: list[Any]) -> str:
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
index = int(match.group(1))
|
||||
return (
|
||||
_sql_literal(values[index - 1])
|
||||
if 0 < index <= len(values)
|
||||
else match.group(0)
|
||||
)
|
||||
|
||||
return re.sub(r"@_(\d+)", replace, raw_sql)
|
||||
|
||||
|
||||
def _column_title(column: str, aliases: dict[str, Any]) -> str:
|
||||
member = str(aliases.get(column) or column)
|
||||
return member
|
||||
|
||||
|
||||
async def query_cube_page(
|
||||
*,
|
||||
base_load: dict[str, Any],
|
||||
tenant_id: str | None,
|
||||
allow_global: bool,
|
||||
page: int,
|
||||
page_size: int,
|
||||
limit_source: str | None,
|
||||
) -> dict[str, Any]:
|
||||
settings = get_cube_report_settings()
|
||||
page_size = min(max(page_size, 1), settings.max_page_size)
|
||||
page_offset = (page - 1) * page_size
|
||||
canonical_query = deepcopy(base_load)
|
||||
if not allow_global:
|
||||
if not tenant_id:
|
||||
raise HTTPException(status_code=422, detail="当前会话缺少租户信息")
|
||||
ensure_tenant_filter(canonical_query, tenant_id)
|
||||
|
||||
is_user_limit = str(limit_source).lower() == "user"
|
||||
semantic_limit = (
|
||||
_positive_int(canonical_query.get("limit")) if is_user_limit else None
|
||||
)
|
||||
semantic_offset = (
|
||||
_positive_int(canonical_query.get("offset")) or 0 if is_user_limit else 0
|
||||
)
|
||||
if semantic_limit is not None and page_offset >= semantic_limit:
|
||||
return {
|
||||
"columns": [],
|
||||
"rows": [],
|
||||
"page": page,
|
||||
"pageSize": page_size,
|
||||
"total": semantic_limit,
|
||||
"sql": "",
|
||||
"limitSource": "user",
|
||||
}
|
||||
|
||||
page_query = deepcopy(canonical_query)
|
||||
page_query["offset"] = semantic_offset + page_offset
|
||||
page_query["limit"] = (
|
||||
min(page_size, semantic_limit - page_offset)
|
||||
if semantic_limit is not None
|
||||
else page_size
|
||||
)
|
||||
page_query.pop("total", None)
|
||||
_ensure_page_order(page_query)
|
||||
|
||||
count_query = deepcopy(canonical_query)
|
||||
for key in ("limit", "offset", "order", "total"):
|
||||
count_query.pop(key, None)
|
||||
|
||||
page_sql, count_sql = await asyncio.gather(
|
||||
_generate_cube_sql(page_query),
|
||||
_generate_cube_sql(count_query),
|
||||
)
|
||||
executable_sql, bindings = _prepare_cube_sql(page_sql["sql"], page_sql["params"])
|
||||
count_executable, count_bindings = _prepare_cube_sql(
|
||||
count_sql["sql"], count_sql["params"]
|
||||
)
|
||||
wrapped_count_sql = (
|
||||
"SELECT COUNT_BIG(1) AS __cube_total FROM (\n"
|
||||
f"{count_executable}\n"
|
||||
") AS __cube_count"
|
||||
)
|
||||
(keys, rows), raw_total = await asyncio.gather(
|
||||
asyncio.to_thread(execute_parameterized_sql, executable_sql, bindings),
|
||||
asyncio.to_thread(
|
||||
execute_parameterized_scalar, wrapped_count_sql, count_bindings
|
||||
),
|
||||
)
|
||||
columns = [
|
||||
{
|
||||
"key": key,
|
||||
"title": _column_title(key, page_sql["aliases"]),
|
||||
"type": "text",
|
||||
}
|
||||
for key in keys
|
||||
]
|
||||
normalized_rows = []
|
||||
for row_index, row in enumerate(rows):
|
||||
normalized = {"__rowKey": f"{page}-{row_index}"}
|
||||
for column in columns:
|
||||
key = column["key"]
|
||||
normalized[key] = _mask_value(key, column["title"], row.get(key))
|
||||
normalized_rows.append(normalized)
|
||||
|
||||
try:
|
||||
total = max(int(raw_total or 0) - semantic_offset, 0)
|
||||
except (TypeError, ValueError):
|
||||
total = len(normalized_rows)
|
||||
if semantic_limit is not None:
|
||||
total = min(total, semantic_limit)
|
||||
return {
|
||||
"columns": columns,
|
||||
"rows": normalized_rows,
|
||||
"page": page,
|
||||
"pageSize": page_size,
|
||||
"total": total,
|
||||
"sql": _display_sql(page_sql["sql"], page_sql["params"]),
|
||||
"limitSource": "user" if is_user_limit else "system",
|
||||
}
|
||||
|
||||
|
||||
async def query_cube_export(
|
||||
*,
|
||||
base_load: dict[str, Any],
|
||||
tenant_id: str | None,
|
||||
allow_global: bool,
|
||||
limit_source: str | None,
|
||||
) -> dict[str, Any]:
|
||||
query = deepcopy(base_load)
|
||||
if not allow_global:
|
||||
if not tenant_id:
|
||||
raise HTTPException(status_code=422, detail="当前会话缺少租户信息")
|
||||
ensure_tenant_filter(query, tenant_id)
|
||||
if str(limit_source).lower() != "user":
|
||||
query.pop("limit", None)
|
||||
query.pop("offset", None)
|
||||
query.pop("total", None)
|
||||
|
||||
sql_info = await _generate_cube_sql(query)
|
||||
executable_sql, bindings = _prepare_cube_sql(sql_info["sql"], sql_info["params"])
|
||||
keys, rows = await asyncio.to_thread(
|
||||
execute_parameterized_sql, executable_sql, bindings
|
||||
)
|
||||
headers = {key: _column_title(key, sql_info["aliases"]) for key in keys}
|
||||
return {
|
||||
"headers": headers,
|
||||
"rows": rows,
|
||||
"sql": _display_sql(sql_info["sql"], sql_info["params"]),
|
||||
}
|
||||
Reference in New Issue
Block a user