154 lines
5.3 KiB
Python
154 lines
5.3 KiB
Python
"""Persistence queries for user-owned assistant conversations."""
|
|
|
|
from sqlalchemy import and_, func, or_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import selectinload
|
|
|
|
from app.models.entities import Conversation, ConversationMessage, IntegrationConversation
|
|
|
|
|
|
class ConversationRepository:
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self.session = session
|
|
|
|
async def is_external(self, conversation_id: int) -> bool:
|
|
return (
|
|
await self.session.scalar(
|
|
select(IntegrationConversation.id).where(IntegrationConversation.conversation_id == conversation_id)
|
|
)
|
|
is not None
|
|
)
|
|
|
|
async def list_for_user(
|
|
self,
|
|
user_id: int,
|
|
*,
|
|
limit: int = 50,
|
|
search: str | None = None,
|
|
) -> list[Conversation]:
|
|
filters = [Conversation.user_id == user_id]
|
|
if search:
|
|
filters.append(Conversation.title.ilike(f"%{search}%"))
|
|
statement = (
|
|
select(Conversation)
|
|
.where(*filters)
|
|
.options(selectinload(Conversation.messages))
|
|
.order_by(Conversation.updated_at.desc(), Conversation.id.desc())
|
|
.limit(limit)
|
|
)
|
|
return list(await self.session.scalars(statement))
|
|
|
|
async def get_page_for_user(
|
|
self,
|
|
conversation_id: int,
|
|
user_id: int,
|
|
*,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> tuple[Conversation, list[ConversationMessage], int] | None:
|
|
conversation = await self.session.scalar(
|
|
select(Conversation).where(
|
|
Conversation.id == conversation_id,
|
|
Conversation.user_id == user_id,
|
|
)
|
|
)
|
|
if conversation is None:
|
|
return None
|
|
return await self.get_message_page(conversation, limit=limit, offset=offset)
|
|
|
|
async def get_message_page(
|
|
self,
|
|
conversation: Conversation,
|
|
*,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> tuple[Conversation, list[ConversationMessage], int]:
|
|
"""Page messages after the caller has verified ownership of the conversation."""
|
|
|
|
conversation_id = conversation.id
|
|
total = await self.session.scalar(
|
|
select(func.count())
|
|
.select_from(ConversationMessage)
|
|
.where(
|
|
ConversationMessage.conversation_id == conversation_id,
|
|
)
|
|
)
|
|
messages = list(
|
|
await self.session.scalars(
|
|
select(ConversationMessage)
|
|
.where(ConversationMessage.conversation_id == conversation_id)
|
|
.order_by(ConversationMessage.created_at.desc(), ConversationMessage.id.desc())
|
|
.offset(offset)
|
|
.limit(limit)
|
|
)
|
|
)
|
|
messages.reverse()
|
|
return conversation, messages, int(total or 0)
|
|
|
|
async def get_for_user(self, conversation_id: int, user_id: int) -> Conversation | None:
|
|
statement = (
|
|
select(Conversation)
|
|
.where(Conversation.id == conversation_id, Conversation.user_id == user_id)
|
|
.options(selectinload(Conversation.messages))
|
|
)
|
|
return await self.session.scalar(statement)
|
|
|
|
async def delete_for_user(self, conversation_id: int, user_id: int) -> bool:
|
|
"""Delete a conversation only when it belongs to the requesting user."""
|
|
|
|
conversation = await self.get_for_user(conversation_id, user_id)
|
|
if conversation is None:
|
|
return False
|
|
await self.session.delete(conversation)
|
|
await self.session.flush()
|
|
return True
|
|
|
|
async def get_withdraw_target_for_user(
|
|
self,
|
|
conversation_id: int,
|
|
message_id: int,
|
|
user_id: int,
|
|
) -> tuple[Conversation, ConversationMessage | None] | None:
|
|
"""Check ownership and load the target in one query without loading history."""
|
|
|
|
row = (
|
|
await self.session.execute(
|
|
select(Conversation, ConversationMessage)
|
|
.outerjoin(
|
|
ConversationMessage,
|
|
and_(
|
|
ConversationMessage.conversation_id == Conversation.id,
|
|
ConversationMessage.id == message_id,
|
|
),
|
|
)
|
|
.where(
|
|
Conversation.id == conversation_id,
|
|
Conversation.user_id == user_id,
|
|
)
|
|
)
|
|
).one_or_none()
|
|
return (row[0], row[1]) if row is not None else None
|
|
|
|
async def get_next_message(
|
|
self,
|
|
conversation_id: int,
|
|
message: ConversationMessage,
|
|
) -> ConversationMessage | None:
|
|
"""Load only the message immediately following the target message."""
|
|
|
|
return await self.session.scalar(
|
|
select(ConversationMessage)
|
|
.where(
|
|
ConversationMessage.conversation_id == conversation_id,
|
|
or_(
|
|
ConversationMessage.created_at > message.created_at,
|
|
and_(
|
|
ConversationMessage.created_at == message.created_at,
|
|
ConversationMessage.id > message.id,
|
|
),
|
|
),
|
|
)
|
|
.order_by(ConversationMessage.created_at, ConversationMessage.id)
|
|
.limit(1)
|
|
)
|