src/memory/session.py

"""
Session store backed by Qdrant Cloud.

Each session is one Qdrant point (payload-only; vector is a dummy scalar).
Payload schema:
  {
    "session_id":  str,
    "created_at":  float,   # unix timestamp
    "last_active": float,   # unix timestamp — updated on every turn
    "messages":    [{"role": str, "content": str, "timestamp": float}, ...]
  }

This module is the single source of truth for session lifecycle.
It is also used by cleanup_sessions.py.
"""
from __future__ import annotations
import time
import uuid
from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance,
    FieldCondition,
    Filter,
    PayloadSchemaType,
    PointIdsList,
    PointStruct,
    Range,
    VectorParams,
)
from src.config import QDRANT_URL, QDRANT_API_KEY, QDRANT_LOG_COLLECTION


class SessionStore:
    # Qdrant requires a vector; we use a single dummy dimension.
    _VECTOR_DIM = 1

    def __init__(self):
        self._client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
        self._ensure_collection()

    def _ensure_collection(self):
        names = [c.name for c in self._client.get_collections().collections]
        if QDRANT_LOG_COLLECTION not in names:
            self._client.create_collection(
                collection_name=QDRANT_LOG_COLLECTION,
                vectors_config=VectorParams(
                    size=self._VECTOR_DIM,
                    distance=Distance.COSINE,
                ),
            )
        self._client.create_payload_index(
            collection_name=QDRANT_LOG_COLLECTION,
            field_name="last_active",
            field_schema=PayloadSchemaType.FLOAT,
        )

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def new_session(self) -> str:
        sid = str(uuid.uuid4())
        now = time.time()
        self._upsert(sid, {"session_id": sid, "created_at": now, "last_active": now, "messages": []})
        return sid

    def touch(self, session_id: str):
        """Update last_active without adding a message."""
        payload = self._get_payload(session_id)
        if payload:
            payload["last_active"] = time.time()
            self._upsert(session_id, payload)

    def log_turn(self, session_id: str, role: str, content: str):
        payload = self._get_payload(session_id)
        if not payload:
            return
        payload["messages"].append({"role": role, "content": content, "timestamp": time.time()})
        payload["last_active"] = time.time()
        self._upsert(session_id, payload)

    def delete_session(self, session_id: str):
        self._client.delete(
            collection_name=QDRANT_LOG_COLLECTION,
            points_selector=PointIdsList(points=[session_id]),
        )

    # ── Queries ───────────────────────────────────────────────────────────────

    def get_messages(self, session_id: str) -> list[dict]:
        """Return stored message list for a session, or [] if not found."""
        payload = self._get_payload(session_id)
        return payload.get("messages", []) if payload else []

    def stale_sessions(self, older_than_hours: int = 48) -> list[str]:
        """Return session IDs whose last_active is older than the threshold."""
        cutoff = time.time() - older_than_hours * 3600
        records, _ = self._client.scroll(
            collection_name=QDRANT_LOG_COLLECTION,
            scroll_filter=Filter(
                must=[FieldCondition(key="last_active", range=Range(lt=cutoff))]
            ),
            with_payload=["session_id"],
            limit=1000,
        )
        return [r.payload["session_id"] for r in records if r.payload]

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _upsert(self, session_id: str, payload: dict):
        self._client.upsert(
            collection_name=QDRANT_LOG_COLLECTION,
            points=[PointStruct(id=session_id, vector=[0.0], payload=payload)],
        )

    def _get_payload(self, session_id: str) -> dict | None:
        records = self._client.retrieve(
            collection_name=QDRANT_LOG_COLLECTION,
            ids=[session_id],
            with_payload=True,
        )
        return records[0].payload if records else None


# Singleton — shared across Chainlit coroutines in the same process
_store: SessionStore | None = None


def get_store() -> SessionStore:
    global _store
    if _store is None:
        _store = SessionStore()
    return _store