qzl
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
from uuid import UUID
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from models.agent_chat_session import AgentChatSession
|
|
|
|
|
|
class AgentRepository:
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self._session = session
|
|
|
|
async def get_session_owner(self, *, session_id: str) -> str:
|
|
try:
|
|
session_uuid = UUID(session_id)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=422, detail="Invalid session_id") from exc
|
|
|
|
stmt = select(AgentChatSession.user_id).where(
|
|
AgentChatSession.id == session_uuid
|
|
)
|
|
owner_id = (await self._session.execute(stmt)).scalar_one_or_none()
|
|
if owner_id is None:
|
|
raise HTTPException(status_code=404, detail="Session not found")
|
|
return str(owner_id)
|
|
|
|
async def create_session_for_user(self, *, user_id: str) -> str:
|
|
try:
|
|
user_uuid = UUID(user_id)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=422, detail="Invalid user_id") from exc
|
|
|
|
session = AgentChatSession(
|
|
user_id=user_uuid,
|
|
)
|
|
self._session.add(session)
|
|
await self._session.flush()
|
|
await self._session.refresh(session)
|
|
return str(session.id)
|
|
|
|
async def commit(self) -> None:
|
|
await self._session.commit()
|
|
|
|
async def rollback(self) -> None:
|
|
await self._session.rollback()
|
|
|
|
async def delete_session(self, *, session_id: str) -> None:
|
|
try:
|
|
session_uuid = UUID(session_id)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=422, detail="Invalid session_id") from exc
|
|
session = await self._session.get(AgentChatSession, session_uuid)
|
|
if session is not None:
|
|
await self._session.delete(session)
|
|
await self._session.flush()
|