"""Repository functions over the storage models. Pure async + SQLAlchemy 2.0 style.""" from __future__ import annotations import logging import secrets from typing import Any from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.config import get_settings from app.storage.db import SessionLocal from app.storage.models import Message, Session, SkillPref, ToolRun logger = logging.getLogger(__name__) def _new_id() -> str: return secrets.token_urlsafe(16) # ---------- sessions ---------- async def _prune_old_sessions(db: AsyncSession, user_id: str) -> int: """If the user has more than `max_sessions_per_user` sessions, delete the oldest by `created_at` until the count is at most that limit. Returns the number deleted. Messages cascade. `max_sessions_per_user == 0` disables the cap. """ cap = get_settings().max_sessions_per_user if cap <= 0: return 0 result = await db.execute( select(Session) .where(Session.user_id == user_id) .order_by(Session.created_at.desc()) ) rows = result.scalars().all() if len(rows) <= cap: return 0 excess = [r.id for r in rows[cap:]] from sqlalchemy import delete await db.execute(delete(Session).where(Session.id.in_(excess))) return len(excess) async def create_session_async(title: str, user_id: str) -> str: sid = _new_id() async with SessionLocal() as db: db.add(Session(id=sid, title=title, user_id=user_id)) await db.commit() # Enforce retention: keep at most max_sessions_per_user pruned = await _prune_old_sessions(db, user_id) if pruned: await db.commit() logger.info("Pruned %d oldest session(s) for user=%s", pruned, user_id) return sid def create_session(title: str, user_id: str = "default") -> str: import asyncio return asyncio.run(create_session_async(title, user_id)) async def list_sessions_async(user_id: str, limit: int) -> list[dict[str, Any]]: async with SessionLocal() as db: result = await db.execute( select(Session).where(Session.user_id == user_id).order_by(Session.updated_at.desc()).limit(limit) ) rows = result.scalars().all() return [ { "id": s.id, "title": s.title, "created_at": s.created_at.isoformat(), "updated_at": s.updated_at.isoformat(), } for s in rows ] def list_sessions(user_id: str = "default", limit: int = 50) -> list[dict[str, Any]]: import asyncio return asyncio.run(list_sessions_async(user_id, limit)) async def get_session_async(session_id: str) -> dict[str, Any] | None: async with SessionLocal() as db: s = await db.get(Session, session_id) if not s: return None return { "id": s.id, "title": s.title, "user_id": s.user_id, "created_at": s.created_at.isoformat(), "updated_at": s.updated_at.isoformat(), } def get_session(session_id: str) -> dict[str, Any] | None: import asyncio return asyncio.run(get_session_async(session_id)) async def update_session_title_async(session_id: str, title: str) -> None: async with SessionLocal() as db: s = await db.get(Session, session_id) if s: s.title = title await db.commit() def update_session_title(session_id: str, title: str) -> None: import asyncio asyncio.run(update_session_title_async(session_id, title)) async def delete_session_async(session_id: str) -> None: async with SessionLocal() as db: s = await db.get(Session, session_id) if s: await db.delete(s) await db.commit() def delete_session(session_id: str) -> None: import asyncio asyncio.run(delete_session_async(session_id)) async def delete_all_sessions_async(user_id: str = "default") -> int: """Delete every session belonging to `user_id`. Returns the count.""" from sqlalchemy import delete async with SessionLocal() as db: result = await db.execute(delete(Session).where(Session.user_id == user_id)) await db.commit() return result.rowcount or 0 def delete_all_sessions(user_id: str = "default") -> int: import asyncio return asyncio.run(delete_all_sessions_async(user_id)) # ---------- messages ---------- async def save_message_async( session_id: str, role: str, content: str, tool_calls: list[dict[str, Any]] | None = None, tool_call_id: str | None = None, thinking: dict[str, Any] | None = None, ) -> str: import json mid = _new_id() async with SessionLocal() as db: db.add( Message( id=mid, session_id=session_id, role=role, content=content, tool_calls_json=json.dumps(tool_calls) if tool_calls else None, tool_call_id=tool_call_id, thinking_json=json.dumps(thinking) if thinking else None, ) ) # bump session updated_at s = await db.get(Session, session_id) if s: from datetime import datetime, timezone s.updated_at = datetime.now(timezone.utc) await db.commit() return mid def save_message( session_id: str, role: str, content: str, tool_calls: list[dict[str, Any]] | None = None, tool_call_id: str | None = None, thinking: dict[str, Any] | None = None, ) -> str: import asyncio return asyncio.run( save_message_async( session_id, role, content, tool_calls, tool_call_id, thinking ) ) async def list_messages_async(session_id: str) -> list[dict[str, Any]]: import json async with SessionLocal() as db: result = await db.execute( select(Message).where(Message.session_id == session_id).order_by(Message.created_at) ) rows = result.scalars().all() out: list[dict[str, Any]] = [] for m in rows: entry: dict[str, Any] = { "id": m.id, "role": m.role, "content": m.content, "created_at": m.created_at.isoformat(), } if m.tool_calls_json: entry["tool_calls"] = json.loads(m.tool_calls_json) if m.tool_call_id: entry["tool_call_id"] = m.tool_call_id if m.thinking_json: entry["thinking"] = json.loads(m.thinking_json) out.append(entry) return out def list_messages(session_id: str) -> list[dict[str, Any]]: import asyncio return asyncio.run(list_messages_async(session_id))