主干Ai实验室后端项目

2025-09-05 09:37:47 +08:00
parent aa25f914ab
commit 4a0e79b35a
25 changed files with 628 additions and 0 deletions
@@ -0,0 +1,115 @@
+import psycopg
+from langchain_postgres import PostgresChatMessageHistory
+from config.pgDb import database_name,getConn
+from typing import List, Dict
+# ————————————————————————————————————————————————————AI角色———————————————————————————————
+
+def get_ai_personality(ai_id: str):
+    with getConn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                "SELECT ai_personality FROM ai_chat_profiles WHERE id = %s",
+                (ai_id,)
+            )
+            row = cur.fetchone()
+            if row:
+                return row[0]
+            else:
+                return "你是一个乐于助人的AI助手，请保持中文简洁回答用户。"
+            
+def get_all_ai(user_id: str) -> List[Dict]:
+    with getConn() as conn:
+        with conn.cursor() as cur:
+            # 查询用户角色
+            cur.execute(
+                "SELECT roles FROM users WHERE id = %s",
+                (user_id,)
+            )
+            role_row = cur.fetchone()
+            if not role_row:
+                return []  # 用户不存在
+            user_roles = role_row[0]
+
+            # 查询 AI 角色 JSON 字段包含用户角色
+            cur.execute(
+                """
+                SELECT id, name, welcome_words
+                FROM ai_chat_profiles
+                WHERE availabel_roles::jsonb ?| %s
+                """,
+                (user_roles,)  # user_roles 是 list，比如 ["a", "b", "c"]
+            )
+            rows = cur.fetchall()
+            return [
+                {
+                    "id": row[0],
+                    "name": row[1],
+                    "welcomeWords": row[2],
+                }
+                for row in rows
+            ]
+
+# ————————————————————————————————————————————————————消息———————————————————————————————
+def insert_message(session_id: str, isAI: bool, content: str):
+    with getConn() as conn:
+        history = PostgresChatMessageHistory(
+            database_name,
+            session_id,
+            sync_connection=conn
+        )
+        if isAI:
+            history.add_ai_message(content)
+        else:
+            history.add_user_message(content)
+
+def get_history(session_id: str):
+    with getConn() as conn:
+        history = PostgresChatMessageHistory(
+            database_name,
+            session_id,
+            sync_connection=conn
+        )
+        simplified = []
+        for msg in history.messages:
+            simplified.append({
+                "type": msg.type,
+                "content": msg.content
+            })
+
+        return simplified
+
+# ————————————————————————————————————————————————————会话———————————————————————————————
+def insert_session(user_id: str,ai_id:str, session_id: str,session_title: str):
+    with getConn() as coon:
+        with coon.cursor() as cur:
+            cur.execute(
+                "INSERT INTO ai_chat_sessions (id ,user_id, ai_id, title, created_at, updated_at) VALUES (%s, %s, %s, %s, NOW(), NOW())",
+                (session_id, user_id, ai_id, session_title )
+            )
+            coon.commit()
+
+def update_session_updated_at(session_id: str):
+    with getConn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                "UPDATE ai_chat_sessions SET updated_at = NOW() WHERE id = %s",
+                (session_id,)
+            )
+            conn.commit()
+
+def get_sessions(user_id: str):
+    with getConn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                "SELECT id, title,  updated_at FROM ai_chat_sessions WHERE user_id = %s ORDER BY updated_at DESC",
+                (user_id,)
+            )
+            sessions = cur.fetchall()
+            return [
+                {
+                    "id": row[0],
+                    "title": row[1],
+                    "updated_at": row[2]
+                }
+                for row in sessions
+            ]