Первый коммит
This commit is contained in:
93
app/modules/chat/repository.py
Normal file
93
app/modules/chat/repository.py
Normal file
@@ -0,0 +1,93 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user