94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
import json
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.modules.shared.db import get_engine
|
|
|
|
|
|
class ChatRepository:
|
|
def ensure_tables(self) -> None:
|
|
with get_engine().connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS dialog_sessions (
|
|
dialog_session_id VARCHAR(64) PRIMARY KEY,
|
|
rag_session_id VARCHAR(64) NOT NULL,
|
|
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS chat_messages (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
dialog_session_id VARCHAR(64) NOT NULL,
|
|
task_id VARCHAR(64),
|
|
role VARCHAR(16) NOT NULL,
|
|
content TEXT NOT NULL,
|
|
payload JSONB,
|
|
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS task_id VARCHAR(64)"))
|
|
conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS payload JSONB"))
|
|
conn.commit()
|
|
|
|
def create_dialog(self, dialog_session_id: str, rag_session_id: str) -> None:
|
|
with get_engine().connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO dialog_sessions (dialog_session_id, rag_session_id)
|
|
VALUES (:did, :sid)
|
|
"""
|
|
),
|
|
{"did": dialog_session_id, "sid": rag_session_id},
|
|
)
|
|
conn.commit()
|
|
|
|
def get_dialog(self, dialog_session_id: str) -> dict | None:
|
|
with get_engine().connect() as conn:
|
|
row = conn.execute(
|
|
text(
|
|
"""
|
|
SELECT dialog_session_id, rag_session_id
|
|
FROM dialog_sessions
|
|
WHERE dialog_session_id = :did
|
|
"""
|
|
),
|
|
{"did": dialog_session_id},
|
|
).mappings().fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def add_message(
|
|
self,
|
|
dialog_session_id: str,
|
|
role: str,
|
|
content: str,
|
|
task_id: str | None = None,
|
|
payload: dict | None = None,
|
|
) -> None:
|
|
payload_json = json.dumps(payload, ensure_ascii=False) if payload is not None else None
|
|
with get_engine().connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO chat_messages (dialog_session_id, task_id, role, content, payload)
|
|
VALUES (:did, :task_id, :role, :content, CAST(:payload AS JSONB))
|
|
"""
|
|
),
|
|
{
|
|
"did": dialog_session_id,
|
|
"task_id": task_id,
|
|
"role": role,
|
|
"content": content,
|
|
"payload": payload_json,
|
|
},
|
|
)
|
|
conn.commit()
|