BBIT-Kai
115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
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
|
|
] |