63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
"""ChatSessionRepository — /chat 세션 상태 DB 영속화 (P8-A).
|
|
|
|
인메모리 대신 learning.chat_sessions 에 진행 상태를 저장 → 서버 재시작/멀티워커 안전.
|
|
company_id 스코프. ChatSession(dataclass) ↔ row 직렬화.
|
|
"""
|
|
|
|
import uuid
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from common.database.db_session_manager import DB_SESSION_MNG
|
|
from common.database.model.models import ChatSessionRow
|
|
from common.enums import DBType, DBWRType, ErrorType
|
|
from common.utils.gtime import GTime
|
|
from negotiation.chat.service.chat_engine import ChatSession
|
|
|
|
|
|
class ChatSessionRepository:
|
|
def __init__(self, company_id: str):
|
|
self.company_id = company_id
|
|
|
|
async def get(self, session_id: Optional[str]) -> Optional[ChatSession]:
|
|
if not session_id:
|
|
return None
|
|
|
|
def _q(s: AsyncSession):
|
|
stmt = select(ChatSessionRow).where(
|
|
ChatSessionRow.session_id == session_id,
|
|
ChatSessionRow.company_id == self.company_id,
|
|
).limit(1)
|
|
return DB_SESSION_MNG.execute(s, stmt)
|
|
|
|
err, rows = await DB_SESSION_MNG.execute_lambda(DBType.MAIN.value, DBWRType.DB_READ.value, _q)
|
|
if err != ErrorType.SUCCESS or not rows:
|
|
return None
|
|
r = rows[0]
|
|
return ChatSession(
|
|
session_id=str(r.session_id), tenant_id=r.tenant_id, company_id=r.company_id,
|
|
rq_type=r.rq_type, step=r.step, context=dict(r.context or {}),
|
|
used_action_ids=set(r.used_action_ids or []), action_space_size=r.action_space_size,
|
|
ended=r.ended,
|
|
)
|
|
|
|
async def save(self, session: ChatSession) -> ErrorType:
|
|
values = dict(
|
|
session_id=session.session_id, company_id=session.company_id, tenant_id=session.tenant_id,
|
|
rq_type=session.rq_type, step=session.step, context=session.context,
|
|
used_action_ids=sorted(session.used_action_ids), action_space_size=session.action_space_size,
|
|
ended=session.ended, updated_at=GTime.UTC(),
|
|
)
|
|
|
|
def _do(s: AsyncSession):
|
|
stmt = pg_insert(ChatSessionRow.__table__).values(**values).on_conflict_do_update(
|
|
index_elements=[ChatSessionRow.session_id],
|
|
set_={k: values[k] for k in ("step", "context", "used_action_ids", "ended", "updated_at")},
|
|
)
|
|
return DB_SESSION_MNG.add(s, stmt)
|
|
|
|
return await DB_SESSION_MNG.execute_lambda_run([DBType.MAIN.value], [_do])
|