SPB Git forge

spb/uqo-chat

Public
14commits 1branches 0releases
1.4 MBsize
maindefault branch
17 days agolast push
Python 64.6% TypeScript 33.7% CSS 0.8%
8.7 KB · 210 lines python
Raw Blame History
1"""Conversation / message persistence helpers."""23from __future__ import annotations45from datetime import timedelta6from typing import Any78from sqlalchemy import delete, func, or_, select9from sqlalchemy.orm import selectinload1011from app.db import SessionLocal12from app.models import Conversation, Message, ToolCall, utcnow131415def conv_to_dict(c: Conversation) -> dict[str, Any]:16    return {"id": c.id, "title": c.title, "course": c.course_code, "pinned": c.pinned,17            "archived": c.archived, "created_at": c.created_at.isoformat(),18            "updated_at": c.updated_at.isoformat()}192021def message_to_dict(m: Message, include_tools: bool = True) -> dict[str, Any]:22    d: dict[str, Any] = {23        "id": m.id, "role": m.role, "content": m.content, "model": m.model,24        "created_at": m.created_at.isoformat(), "feedback": m.feedback,25        "attachments": m.attachments or [], "cost_usd": m.cost_usd,26        "tokens": {"in": m.tokens_in, "out": m.tokens_out}, "latency_ms": m.latency_ms,27    }28    if include_tools:29        d["tool_calls"] = [{30            "id": t.id, "name": t.name, "arguments": t.arguments, "summary": t.result_summary,31            "payload": t.result_payload, "status": t.status, "duration_ms": t.duration_ms,32        } for t in m.tool_calls]33    return d343536async def list_conversations(user_id: str, q: str | None = None, archived: bool = False,37                             limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:38    async with SessionLocal() as session:39        stmt = select(Conversation).where(Conversation.user_id == user_id,40                                          Conversation.archived == archived)41        if q:42            like = f"%{q}%"43            sub = select(Message.conversation_id).where(Message.content.ilike(like))44            stmt = stmt.where(or_(Conversation.title.ilike(like), Conversation.id.in_(sub)))45        stmt = stmt.order_by(Conversation.pinned.desc(), Conversation.updated_at.desc()) \46            .limit(limit).offset(offset)47        rows = (await session.execute(stmt)).scalars()48        return [conv_to_dict(c) for c in rows]495051async def create_conversation(user_id: str, course: str, title: str | None = None) -> dict[str, Any]:52    async with SessionLocal() as session:53        c = Conversation(user_id=user_id, course_code=course.upper(),54                         title=title or "Nouvelle conversation")55        session.add(c)56        await session.commit()57        await session.refresh(c)58        return conv_to_dict(c)596061async def get_conversation(conv_id: str, user_id: str | None = None) -> Conversation | None:62    async with SessionLocal() as session:63        c = await session.get(Conversation, conv_id)64        if c is None or (user_id is not None and c.user_id != user_id):65            return None66        return c676869async def get_conversation_full(conv_id: str, user_id: str) -> dict[str, Any] | None:70    async with SessionLocal() as session:71        c = await session.scalar(72            select(Conversation).where(Conversation.id == conv_id, Conversation.user_id == user_id)73            .options(selectinload(Conversation.messages).selectinload(Message.tool_calls)))74        if c is None:75            return None76        return {**conv_to_dict(c), "messages": [message_to_dict(m) for m in c.messages]}777879async def update_conversation(conv_id: str, user_id: str, **fields: Any) -> dict[str, Any] | None:80    async with SessionLocal() as session:81        c = await session.get(Conversation, conv_id)82        if c is None or c.user_id != user_id:83            return None84        for k, v in fields.items():85            if v is not None and k in {"title", "pinned", "archived", "course_code"}:86                setattr(c, k, v)87        await session.commit()88        await session.refresh(c)89        return conv_to_dict(c)909192async def delete_conversation(conv_id: str, user_id: str) -> bool:93    async with SessionLocal() as session:94        c = await session.get(Conversation, conv_id)95        if c is None or c.user_id != user_id:96            return False97        msg_ids = [m for m in (await session.execute(98            select(Message.id).where(Message.conversation_id == conv_id))).scalars()]99        if msg_ids:100            await session.execute(delete(ToolCall).where(ToolCall.message_id.in_(msg_ids)))101            await session.execute(delete(Message).where(Message.id.in_(msg_ids)))102        await session.delete(c)103        await session.commit()104        return True105106107async def history(conv_id: str, limit: int = 40) -> list[Message]:108    async with SessionLocal() as session:109        rows = (await session.execute(110            select(Message).where(Message.conversation_id == conv_id)111            .options(selectinload(Message.tool_calls))112            .order_by(Message.created_at.desc()).limit(limit))).scalars()113        return list(reversed(list(rows)))114115116async def add_message(conv_id: str, role: str, content: str, attachments: list[str] | None = None,117                      **fields: Any) -> Message:118    async with SessionLocal() as session:119        m = Message(conversation_id=conv_id, role=role, content=content,120                    attachments=attachments or [], **fields)121        session.add(m)122        c = await session.get(Conversation, conv_id)123        if c:124            c.updated_at = utcnow()125        await session.commit()126        await session.refresh(m)127        return m128129130async def save_assistant_message(conv_id: str, message_id: str, content: str, model: str,131                                 usage: dict[str, Any], latency_ms: int,132                                 tool_calls: list[dict[str, Any]]) -> None:133    async with SessionLocal() as session:134        m = Message(id=message_id, conversation_id=conv_id, role="assistant", content=content,135                    model=model, tokens_in=int(usage.get("input_tokens", 0)),136                    tokens_out=int(usage.get("output_tokens", 0)),137                    cost_usd=float(usage.get("cost_usd", 0.0)), latency_ms=latency_ms)138        session.add(m)139        for tc in tool_calls:140            session.add(ToolCall(id=tc["id"], message_id=message_id, name=tc["name"],141                                 arguments=tc.get("arguments", {}),142                                 result_summary=tc.get("summary", ""),143                                 result_payload=tc.get("payload", {}),144                                 status=tc.get("status", "ok"),145                                 duration_ms=int(tc.get("duration_ms", 0))))146        c = await session.get(Conversation, conv_id)147        if c:148            c.updated_at = utcnow()149        await session.commit()150151152async def set_title(conv_id: str, title: str) -> None:153    async with SessionLocal() as session:154        c = await session.get(Conversation, conv_id)155        if c:156            c.title = title[:120]157            await session.commit()158159160async def set_feedback(message_id: str, user_id: str, feedback: str | None) -> bool:161    async with SessionLocal() as session:162        m = await session.get(Message, message_id)163        if not m:164            return False165        c = await session.get(Conversation, m.conversation_id)166        if not c or c.user_id != user_id:167            return False168        m.feedback = feedback169        await session.commit()170        return True171172173async def delete_messages_after(conv_id: str, message_id: str) -> None:174    """Remove the assistant message (and later ones) for regeneration."""175    async with SessionLocal() as session:176        target = await session.get(Message, message_id)177        if not target:178            return179        ids = [m for m in (await session.execute(180            select(Message.id).where(Message.conversation_id == conv_id,181                                     Message.created_at >= target.created_at))).scalars()]182        if ids:183            await session.execute(delete(ToolCall).where(ToolCall.message_id.in_(ids)))184            await session.execute(delete(Message).where(Message.id.in_(ids)))185        await session.commit()186187188async def redact_old_messages(months: int) -> int:189    cutoff = utcnow() - timedelta(days=30 * months)190    async with SessionLocal() as session:191        rows = (await session.execute(192            select(Message).where(Message.created_at < cutoff,193                                  Message.content_redacted_at.is_(None)))).scalars()194        n = 0195        for m in rows:196            m.content = ""197            m.content_redacted_at = utcnow()198            n += 1199        await session.commit()200        return n201202203async def count_messages_today(user_id: str) -> int:204    since = utcnow() - timedelta(days=1)205    async with SessionLocal() as session:206        return int(await session.scalar(207            select(func.count(Message.id)).join(Conversation)208            .where(Conversation.user_id == user_id, Message.role == "user",209                   Message.created_at >= since)) or 0)210