45 lines
1.8 KiB
Python
45 lines
1.8 KiB
Python
|
from datetime import UTC, datetime
|
||
|
|
|
||
|
|
from sqlalchemy import select, update
|
||
|
|
from sqlalchemy.engine import CursorResult
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from app.core.errors import ConflictAgentError, ResourceNotFoundError
|
||
|
|
from app.model.session import ConversationSession
|
||
|
|
|
||
|
|
|
||
|
|
class SessionRepository:
|
||
|
|
def __init__(self, session: AsyncSession) -> None:
|
||
|
|
self.session = session
|
||
|
|
|
||
|
|
async def owned(
|
||
|
|
self, session_id: str, user_id: int, *, lock: bool = False
|
||
|
|
) -> ConversationSession:
|
||
|
|
query = select(ConversationSession).where(ConversationSession.session_id == session_id,
|
||
|
|
ConversationSession.user_id == user_id)
|
||
|
|
if lock:
|
||
|
|
query = query.with_for_update()
|
||
|
|
row = await self.session.scalar(query)
|
||
|
|
if row is None:
|
||
|
|
raise ResourceNotFoundError("会话不存在")
|
||
|
|
return row
|
||
|
|
|
||
|
|
async def require_active(self, session_id: str, user_id: int, agent_type: str) -> None:
|
||
|
|
row = await self.owned(session_id, user_id, lock=True)
|
||
|
|
if row.status != "active":
|
||
|
|
raise ConflictAgentError("会话已结束")
|
||
|
|
if row.agent_type != agent_type:
|
||
|
|
raise ConflictAgentError("Agent 与会话不一致")
|
||
|
|
row.message_count += 1
|
||
|
|
row.last_active_at = datetime.now(UTC).replace(tzinfo=None)
|
||
|
|
|
||
|
|
async def advance_clarification(self, session_id: str, expected: int) -> bool:
|
||
|
|
if expected >= 10:
|
||
|
|
return False
|
||
|
|
result = await self.session.execute(update(ConversationSession).where(
|
||
|
|
ConversationSession.session_id == session_id,
|
||
|
|
ConversationSession.clarification_round == expected,
|
||
|
|
ConversationSession.status == "active",
|
||
|
|
).values(clarification_round=expected + 1))
|
||
|
|
return isinstance(result, CursorResult) and result.rowcount == 1
|