完善权限系统

bbit_ai/app/db/postgres/aimessage.py

@@ -0,0 +1,50 @@
+from langchain_postgres import PostgresChatMessageHistory
+
+from config.pgDb import pg_pool
+
+database_name = "ai_chat_history"
+
+
+def insert_message(session_id: str, isAI: bool, content: str):
+    with pg_pool.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):
+    simplified = []
+    with pg_pool.getConn() as conn:
+        history = PostgresChatMessageHistory(
+            database_name, session_id, sync_connection=conn
+        )
+        for msg in history.messages:
+            simplified.append({"type": msg.type, "content": msg.content})
+        return simplified
+
+
+def get_history_with_time(session_id: str, number: int):
+    simplified = []
+    with pg_pool.getConn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                f"SELECT message, created_at FROM ai_chat_history WHERE session_id = '{session_id}' ORDER BY created_at DESC LIMIT {number}"
+            )
+            rows = cur.fetchall()
+            simplified = []
+
+            for row in rows:
+                msg_dict = row[0]
+                simplified.append(
+                    {
+                        "type": msg_dict.get("type"),
+                        "created_at": row[1].isoformat(),
+                        "content": msg_dict.get("data", {}).get("content"),
+                    }
+                )
+
+        return simplified