341 lines
12 KiB
Python
341 lines
12 KiB
Python
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
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,
|
|
current_cube, limit_source, response_message,
|
|
request_content, 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],
|
|
"query": row[3],
|
|
"limitSource": row[4],
|
|
"message": row[5],
|
|
"requestContent": 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,
|
|
tenant_name: str,
|
|
query: dict[str, Any],
|
|
limit_source: str,
|
|
response_message: str,
|
|
request_content: 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, current_cube, limit_source,
|
|
response_message, request_content, 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,
|
|
current_cube = EXCLUDED.current_cube,
|
|
limit_source = EXCLUDED.limit_source,
|
|
response_message = EXCLUDED.response_message,
|
|
request_content = EXCLUDED.request_content,
|
|
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],
|
|
Jsonb(query),
|
|
normalized_source,
|
|
response_message,
|
|
request_content,
|
|
),
|
|
)
|
|
cur.execute(
|
|
"""
|
|
UPDATE ai_cube_saved_reports
|
|
SET requirement = %s,
|
|
tenant_id = %s,
|
|
tenant_name = %s,
|
|
cube_query = %s,
|
|
limit_source = %s,
|
|
response_message = %s,
|
|
updated_at = NOW()
|
|
WHERE dify_conversation_id = %s AND created_by = %s
|
|
""",
|
|
(
|
|
request_content,
|
|
tenant_id,
|
|
tenant_name[:200],
|
|
Jsonb(query),
|
|
normalized_source,
|
|
response_message,
|
|
conversation_id,
|
|
user_id,
|
|
),
|
|
)
|
|
|
|
|
|
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),
|
|
)
|
|
|
|
|
|
def _saved_report_row(row) -> dict[str, Any]:
|
|
return {
|
|
"id": str(row[0]),
|
|
"conversationId": str(row[1]),
|
|
"title": row[2],
|
|
"requirement": row[3],
|
|
"tenantId": str(row[4]),
|
|
"tenantName": row[5],
|
|
"query": row[6],
|
|
"limitSource": row[7],
|
|
"message": row[8],
|
|
"createdAt": row[9].isoformat() if row[9] else None,
|
|
"updatedAt": row[10].isoformat() if row[10] else None,
|
|
}
|
|
|
|
|
|
def list_saved_cube_reports(
|
|
*,
|
|
user_id: str,
|
|
tenant_id: str,
|
|
page: int,
|
|
page_size: int,
|
|
keyword: str = "",
|
|
) -> tuple[list[dict[str, Any]], int]:
|
|
conditions = ["created_by = %s", "tenant_id = %s"]
|
|
params: list[Any] = [user_id, tenant_id]
|
|
if keyword:
|
|
conditions.append("(title ILIKE %s OR requirement ILIKE %s)")
|
|
value = f"%{keyword}%"
|
|
params.extend([value, value])
|
|
where_clause = " AND ".join(conditions)
|
|
offset = (page - 1) * page_size
|
|
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
f"SELECT COUNT(*) FROM ai_cube_saved_reports WHERE {where_clause}",
|
|
tuple(params),
|
|
)
|
|
total = int(cur.fetchone()[0])
|
|
cur.execute(
|
|
f"""
|
|
SELECT id, dify_conversation_id, title, requirement,
|
|
tenant_id, tenant_name, cube_query, limit_source,
|
|
response_message, created_at, updated_at
|
|
FROM ai_cube_saved_reports
|
|
WHERE {where_clause}
|
|
ORDER BY updated_at DESC, id DESC
|
|
LIMIT %s OFFSET %s
|
|
""",
|
|
(*params, page_size, offset),
|
|
)
|
|
items = [_saved_report_row(row) for row in cur.fetchall()]
|
|
return items, total
|
|
|
|
|
|
def get_saved_cube_report(report_id: str, user_id: str) -> dict[str, Any] | None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT id, dify_conversation_id, title, requirement,
|
|
tenant_id, tenant_name, cube_query, limit_source,
|
|
response_message, created_at, updated_at
|
|
FROM ai_cube_saved_reports
|
|
WHERE id = %s AND created_by = %s
|
|
""",
|
|
(report_id, user_id),
|
|
)
|
|
row = cur.fetchone()
|
|
return _saved_report_row(row) if row else None
|
|
|
|
|
|
def get_saved_report_by_conversation(
|
|
conversation_id: str, user_id: str
|
|
) -> dict[str, Any] | None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT id, dify_conversation_id, title, requirement,
|
|
tenant_id, tenant_name, cube_query, limit_source,
|
|
response_message, created_at, updated_at
|
|
FROM ai_cube_saved_reports
|
|
WHERE dify_conversation_id = %s AND created_by = %s
|
|
""",
|
|
(conversation_id, user_id),
|
|
)
|
|
row = cur.fetchone()
|
|
return _saved_report_row(row) if row else None
|
|
|
|
|
|
def list_saved_report_links(
|
|
user_id: str, conversation_ids: list[str]
|
|
) -> dict[str, str]:
|
|
if not conversation_ids:
|
|
return {}
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT dify_conversation_id, id
|
|
FROM ai_cube_saved_reports
|
|
WHERE created_by = %s
|
|
AND dify_conversation_id = ANY(%s::varchar[])
|
|
""",
|
|
(user_id, conversation_ids),
|
|
)
|
|
return {str(row[0]): str(row[1]) for row in cur.fetchall()}
|
|
|
|
|
|
def save_cube_report_from_state(
|
|
*,
|
|
conversation_id: str,
|
|
user_id: str,
|
|
title: str,
|
|
) -> dict[str, Any] | None:
|
|
report_id = str(uuid4())
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO ai_cube_saved_reports (
|
|
id, dify_conversation_id, created_by, title, requirement,
|
|
tenant_id, tenant_name, cube_query, limit_source,
|
|
response_message, created_at, updated_at
|
|
)
|
|
SELECT %s, dify_conversation_id, created_by, %s,
|
|
request_content, tenant_id, tenant_name, current_cube,
|
|
limit_source, response_message, NOW(), NOW()
|
|
FROM ai_cube_report_states
|
|
WHERE dify_conversation_id = %s
|
|
AND created_by = %s
|
|
AND tenant_id IS NOT NULL
|
|
ON CONFLICT (created_by, dify_conversation_id) DO UPDATE SET
|
|
requirement = EXCLUDED.requirement,
|
|
tenant_id = EXCLUDED.tenant_id,
|
|
tenant_name = EXCLUDED.tenant_name,
|
|
cube_query = EXCLUDED.cube_query,
|
|
limit_source = EXCLUDED.limit_source,
|
|
response_message = EXCLUDED.response_message,
|
|
updated_at = NOW()
|
|
RETURNING id, dify_conversation_id, title, requirement,
|
|
tenant_id, tenant_name, cube_query, limit_source,
|
|
response_message, created_at, updated_at
|
|
""",
|
|
(
|
|
report_id,
|
|
title[:200] or "未命名报表",
|
|
conversation_id,
|
|
user_id,
|
|
),
|
|
)
|
|
row = cur.fetchone()
|
|
return _saved_report_row(row) if row else None
|
|
|
|
|
|
def rename_saved_cube_report(report_id: str, user_id: str, title: str) -> None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE ai_cube_saved_reports
|
|
SET title = %s, updated_at = NOW()
|
|
WHERE id = %s AND created_by = %s
|
|
""",
|
|
(title[:200], report_id, user_id),
|
|
)
|
|
|
|
|
|
def delete_saved_cube_report(report_id: str, user_id: str) -> None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
DELETE FROM ai_cube_saved_reports
|
|
WHERE id = %s AND created_by = %s
|
|
""",
|
|
(report_id, user_id),
|
|
)
|