""" Per-user chat database models. All conversation and message data are stored in the main PostgreSQL database and separated strictly using the user_id column. Schema: conversations – chat sessions (like ChatGPT sidebar items) messages – individual messages within a conversation """ import os import uuid from datetime import datetime, timezone from typing import List, Optional from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from sqlalchemy import String, DateTime, Text, ForeignKey, select, update from loguru import logger from app.db.database import engine, AsyncSessionLocal class ChatBase(DeclarativeBase): pass class Conversation(ChatBase): __tablename__ = "conversations" id: Mapped[str] = mapped_column(String(36), primary_key=True) user_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) title: Mapped[str] = mapped_column(String(512), default="New Chat") created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc) ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc), ) class Message(ChatBase): __tablename__ = "messages" id: Mapped[str] = mapped_column(String(36), primary_key=True) conversation_id: Mapped[str] = mapped_column( String(36), ForeignKey("conversations.id"), nullable=False, index=True ) role: Mapped[str] = mapped_column(String(20), nullable=False) # "user" | "assistant" content: Mapped[str] = mapped_column(Text, nullable=False) sources_json: Mapped[str | None] = mapped_column(Text) # JSON-serialised SourceChunk list meta_json: Mapped[str | None] = mapped_column(Text) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc) ) # ─── Connection Sharing ──────────────────────────────────────────────────────── def _get_user_engine(user_id: str): return engine async def ensure_user_db(user_id: str) -> None: # No-op since tables are created on the main connection during startup pass def _session_maker(user_id: str) -> async_sessionmaker: return AsyncSessionLocal # ─── CRUD helpers ───────────────────────────────────────────────────────────── async def create_conversation(user_id: str, title: str = "New Chat") -> Conversation: async with _session_maker(user_id)() as db: conv = Conversation(id=str(uuid.uuid4()), user_id=user_id, title=title) db.add(conv) await db.commit() await db.refresh(conv) return conv async def list_conversations(user_id: str) -> List[Conversation]: async with _session_maker(user_id)() as db: result = await db.execute( select(Conversation) .where(Conversation.user_id == user_id) .order_by(Conversation.updated_at.desc()) ) return result.scalars().all() async def get_conversation(user_id: str, conv_id: str) -> Optional[Conversation]: async with _session_maker(user_id)() as db: result = await db.execute( select(Conversation).where( Conversation.id == conv_id, Conversation.user_id == user_id ) ) return result.scalar_one_or_none() async def rename_conversation(user_id: str, conv_id: str, new_title: str) -> None: async with _session_maker(user_id)() as db: await db.execute( update(Conversation) .where(Conversation.id == conv_id, Conversation.user_id == user_id) .values(title=new_title, updated_at=datetime.now(timezone.utc)) ) await db.commit() async def delete_conversation(user_id: str, conv_id: str) -> None: async with _session_maker(user_id)() as db: result = await db.execute( select(Conversation).where( Conversation.id == conv_id, Conversation.user_id == user_id ) ) conv = result.scalar_one_or_none() if conv: # cascade delete messages first msgs = await db.execute(select(Message).where(Message.conversation_id == conv_id)) for msg in msgs.scalars().all(): await db.delete(msg) await db.delete(conv) await db.commit() async def add_message( user_id: str, conv_id: str, role: str, content: str, sources_json: str | None = None, meta_json: str | None = None, ) -> Message: async with _session_maker(user_id)() as db: # Verify ownership conv_check = await db.execute( select(Conversation).where( Conversation.id == conv_id, Conversation.user_id == user_id ) ) if not conv_check.scalar_one_or_none(): raise ValueError("Conversation not found or access denied") msg = Message( id=str(uuid.uuid4()), conversation_id=conv_id, role=role, content=content, sources_json=sources_json, meta_json=meta_json, ) db.add(msg) # Touch conversation updated_at so it floats to top await db.execute( update(Conversation) .where(Conversation.id == conv_id) .values(updated_at=datetime.now(timezone.utc)) ) await db.commit() await db.refresh(msg) return msg async def get_messages(user_id: str, conv_id: str) -> List[Message]: async with _session_maker(user_id)() as db: # Verify ownership conv_check = await db.execute( select(Conversation).where( Conversation.id == conv_id, Conversation.user_id == user_id ) ) if not conv_check.scalar_one_or_none(): return [] result = await db.execute( select(Message) .where(Message.conversation_id == conv_id) .order_by(Message.created_at.asc()) ) return result.scalars().all() async def auto_title_conversation(user_id: str, conv_id: str, first_question: str) -> None: """Set a smart title from the first user message (truncated).""" title = first_question[:60] + ("…" if len(first_question) > 60 else "") await rename_conversation(user_id, conv_id, title)