AI问数

2026-07-16 09:28:57 +08:00
parent cb9a411c55
commit 7be72a50ad
13 changed files with 3383 additions and 0 deletions
@@ -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),
+            )