gogo-telefon/gogo/chat/router.py

"""Chat widget backend (§7): WebSocket endpoint + embeddable script serving.

Same agent, same tools, text instead of voice. No client accounts; sessions are
ephemeral and resumable for 24h via a localStorage token. Rate-limited per IP.
"""

from __future__ import annotations

import logging
import time
import uuid
from collections import defaultdict, deque
from datetime import UTC, datetime, timedelta
from pathlib import Path

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi.responses import Response
from sqlalchemy import select

from gogo.agent.loop import AgentConversation
from gogo.agent.prompt import compose_greeting
from gogo.db import get_sessionmaker
from gogo.models import ChatMessage, ChatSession, Tenant

log = logging.getLogger("gogo.chat")

router = APIRouter()

MAX_MESSAGE_CHARS = 500
MAX_MESSAGES_PER_SESSION = 40
MIN_SECONDS_BETWEEN_MESSAGES = 1.0
SESSIONS_PER_IP_PER_HOUR = 10
RESUME_WINDOW = timedelta(hours=24)

# ip -> deque of session-creation timestamps (in-memory; fine for MVP scale §15)
_session_creations: dict[str, deque[float]] = defaultdict(deque)


def _ip_allowed(ip: str) -> bool:
    now = time.monotonic()
    q = _session_creations[ip]
    while q and now - q[0] > 3600:
        q.popleft()
    if len(q) >= SESSIONS_PER_IP_PER_HOUR:
        return False
    q.append(now)
    return True


WIDGET_JS = (Path(__file__).parent.parent / "static" / "widget.js").read_text


@router.get("/widget.js")
async def widget_js():
    return Response(WIDGET_JS(), media_type="application/javascript")


@router.websocket("/widget/ws")
async def widget_ws(ws: WebSocket):
    key = ws.query_params.get("key", "")
    resume = ws.query_params.get("resume", "")
    ip = ws.client.host if ws.client else ""

    async with get_sessionmaker()() as session:
        tenant = (
            await session.execute(
                select(Tenant).where(Tenant.widget_public_key == key, Tenant.status == "active")
            )
        ).scalar_one_or_none()
        if tenant is None:
            await ws.close(code=4404)
            return

        chat, history = await _resume_or_create(session, tenant, resume, ip)
        if chat is None:
            await ws.close(code=4429)  # rate limited
            return
        await ws.accept()

        convo = AgentConversation(
            session, tenant, channel="chat", chat_session_id=chat.id
        )
        greeting = compose_greeting(tenant, "chat")
        if history:
            # rebuild the model-visible history from stored messages
            for m in history:
                convo.messages.append(
                    {"role": "user" if m.role == "user" else "assistant", "content": m.content}
                )
            await ws.send_json(
                {
                    "type": "resume",
                    "session": chat.resume_token,
                    "messages": [{"role": m.role, "text": m.content} for m in history],
                }
            )
        else:
            convo.messages.append({"role": "assistant", "content": greeting})
            session.add(ChatMessage(session_id=chat.id, role="assistant", content=greeting))
            await session.commit()
            await ws.send_json(
                {"type": "greeting", "session": chat.resume_token, "text": greeting}
            )

        message_count = len(history)
        last_message_at = 0.0
        try:
            while True:
                data = await ws.receive_json()
                if data.get("type") != "message":
                    continue
                if data.get("website"):  # honeypot field — bots fill it, humans never see it
                    continue
                text = str(data.get("text", "")).strip()[:MAX_MESSAGE_CHARS]
                if not text:
                    continue
                now = time.monotonic()
                if now - last_message_at < MIN_SECONDS_BETWEEN_MESSAGES:
                    await ws.send_json({"type": "error", "code": "too_fast"})
                    continue
                last_message_at = now
                message_count += 1
                if message_count > MAX_MESSAGES_PER_SESSION:
                    await ws.send_json(
                        {
                            "type": "message",
                            "text": (
                                "Dostigli ste ograničenje poruka. Molimo pozovite salon "
                                "telefonom. Hvala!"
                            ),
                        }
                    )
                    break

                session.add(ChatMessage(session_id=chat.id, role="user", content=text))
                await ws.send_json({"type": "typing"})
                try:
                    reply = await convo.user_turn(text)
                except Exception:  # noqa: BLE001
                    log.exception("chat agent error (tenant %s)", tenant.slug)
                    reply = (
                        "Izvinite, došlo je do tehničke greške. Molimo pokušajte "
                        "ponovo ili pozovite salon telefonom."
                    )
                session.add(ChatMessage(session_id=chat.id, role="assistant", content=reply))
                chat.outcome = convo.outcome
                await session.commit()
                await ws.send_json({"type": "message", "text": reply})
        except WebSocketDisconnect:
            # No awaits here: the task may be cancelled during teardown and an
            # in-flight DB call on a closing connection can never complete.
            # Outcome is already persisted after every message.
            pass


async def _resume_or_create(
    session, tenant: Tenant, resume_token: str, ip: str
) -> tuple[ChatSession | None, list[ChatMessage]]:
    if resume_token:
        chat = (
            await session.execute(
                select(ChatSession).where(
                    ChatSession.tenant_id == tenant.id,
                    ChatSession.resume_token == resume_token,
                )
            )
        ).scalar_one_or_none()
        if chat is not None:
            started = chat.started_at
            if started.tzinfo is None:
                started = started.replace(tzinfo=UTC)
            if datetime.now(UTC) - started < RESUME_WINDOW:
                history = (
                    (
                        await session.execute(
                            select(ChatMessage)
                            .where(ChatMessage.session_id == chat.id)
                            .order_by(ChatMessage.at)
                        )
                    )
                    .scalars()
                    .all()
                )
                return chat, list(history)
    if not _ip_allowed(ip):
        return None, []
    chat = ChatSession(
        tenant_id=tenant.id,
        resume_token=uuid.uuid4().hex,
        client_ip=ip,
        outcome="abandoned",  # upgraded as the conversation progresses
    )
    session.add(chat)
    from gogo.metering import record_chat_session

    await record_chat_session(session, tenant.id)
    await session.commit()
    return chat, []