BBIT-Kai
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
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
|