pipu-music/services/recommender/pipu_rec/wave/store.py

"""Загрузка каталога, профиля и сессии волны из PostgreSQL / Redis."""

from __future__ import annotations

import json
import logging
import math
import threading
import time
from dataclasses import asdict
from typing import Any, cast

import numpy as np
import psycopg
import redis

from pipu_rec.wave.engine import CatalogIndex, Profile, Session

log = logging.getLogger(__name__)

INDEX_TTL_S = 300
SESSION_TTL_S = 12 * 3600
HALF_LIFE_DAYS = 90.0


def parse_vector(text: str | None) -> np.ndarray | None:
    """pgvector в текстовом виде «[0.1,0.2,...]»."""
    if not text:
        return None
    return np.array(text.strip("[]").split(","), dtype=np.float32)


def load_index(conn: psycopg.Connection[Any]) -> CatalogIndex:
    """Доступные треки: основной файл в подключённом хранилище и на месте."""
    rows = conn.execute(
        """
        SELECT t.id::text,
            COALESCE((SELECT ta.artist_id::text FROM track_artists ta
                WHERE ta.track_id = t.id AND ta.role = 'main' ORDER BY ta.position LIMIT 1),
                al.album_artist_id::text),
            e.audio::text, f.bpm, f.energy, f.danceability, f.valence, f.acousticness,
            COALESCE((SELECT sum(play_count) FROM user_track_stats s WHERE s.track_id = t.id), 0)
        FROM tracks t
        JOIN albums al ON al.id = t.album_id
        JOIN media_files mf ON mf.id = t.primary_file_id AND mf.missing_since IS NULL
        JOIN library_sources src ON src.id = mf.source_id AND src.status = 'mounted'
        LEFT JOIN track_embeddings e ON e.track_id = t.id
        LEFT JOIN track_features f ON f.track_id = t.id
        """
    ).fetchall()
    feats: dict[str, list[float | None]] = {
        k: [] for k in ("bpm", "energy", "danceability", "valence", "acousticness")
    }
    for r in rows:
        for k, v in zip(feats, r[3:8], strict=True):
            feats[k].append(None if v is None else float(v))
    return CatalogIndex.build(
        ids=[r[0] for r in rows],
        artist=[r[1] for r in rows],
        raw_emb=[parse_vector(r[2]) for r in rows],
        features=feats,
        plays=[int(r[8]) for r in rows],
    )


class IndexHolder:
    """Каталог в памяти; перечитывается раз в INDEX_TTL_S в фоне, не блокируя чтение."""

    def __init__(self, dsn: str) -> None:
        self.dsn = dsn
        self._index: CatalogIndex | None = None
        self._loaded = 0.0
        self._lock = threading.Lock()

    def get(self) -> CatalogIndex:
        if self._index is None:
            self.refresh()
        elif time.monotonic() - self._loaded > INDEX_TTL_S and self._lock.acquire(blocking=False):
            threading.Thread(target=self._refresh_locked, daemon=True).start()
        assert self._index is not None
        return self._index

    def refresh(self) -> None:
        with self._lock:
            self._load()

    def _refresh_locked(self) -> None:
        try:
            self._load()
        finally:
            self._lock.release()

    def _load(self) -> None:
        started = time.monotonic()
        with psycopg.connect(self.dsn) as conn:
            self._index = load_index(conn)
        self._loaded = time.monotonic()
        log.info(
            "catalog index loaded: tracks=%d analysed=%d in %.0f ms",
            len(self._index),
            int(self._index.has_emb.sum()),
            (self._loaded - started) * 1000,
        )


def load_profile(conn: psycopg.Connection[Any], user_id: str) -> Profile:
    """implicit_score из агрегатов с затуханием по давности; лайк +3, дизлайк — исключение."""
    now = time.time()
    scores: dict[str, float] = {}
    played: set[str] = set()
    last: dict[str, float] = {}
    for tid, plays, complete, skips, early, last_at in conn.execute(
        """
        SELECT track_id::text, play_count, complete_count, skip_count, early_skip_count, last_played_at
        FROM user_track_stats WHERE user_id = %s
        """,
        (user_id,),
    ):
        played.add(tid)
        raw = complete * 1.0 + max(plays - complete, 0) * 0.7 - early * 1.0 - max(skips - early, 0) * 0.2
        ts = last_at.timestamp() if last_at else now
        last[tid] = ts
        scores[tid] = raw * math.exp(-math.log(2) * (now - ts) / 86400 / HALF_LIFE_DAYS)
    disliked: set[str] = set()
    liked: set[str] = set()
    for tid, reaction in conn.execute(
        "SELECT track_id::text, reaction::text FROM track_reactions WHERE user_id = %s", (user_id,)
    ):
        if reaction == "like":
            liked.add(tid)
            scores[tid] = scores.get(tid, 0.0) + 3.0
        else:
            disliked.add(tid)
            scores.pop(tid, None)
    disliked_artists = {
        r[0]
        for r in conn.execute(
            "SELECT artist_id::text FROM artist_reactions WHERE user_id = %s AND reaction = 'dislike'",
            (user_id,),
        )
    }
    recent = {
        r[0]
        for r in conn.execute(
            """
            SELECT DISTINCT track_id::text FROM play_events
            WHERE user_id = %s AND type = 'play_start' AND occurred_at > now() - interval '3 hours'
            """,
            (user_id,),
        )
    }
    row = conn.execute("SELECT exploration FROM user_profiles_rec WHERE user_id = %s", (user_id,)).fetchone()
    exploration = float(row[0]) if row else 0.15
    if len(played) < 30:
        exploration = max(exploration, 0.4)  # холодный старт: больше исследования на первых треках
    day = 86400
    return Profile(
        scores=scores,
        disliked=disliked,
        disliked_artists=disliked_artists,
        recent=recent,
        played=played,
        liked_recent={t for t in liked if now - last.get(t, 0) < day},
        rediscover={t for t, s in scores.items() if s > 1.5 and now - last.get(t, now) > 60 * day},
        exploration=exploration,
    )


def seed_tracks(conn: psycopg.Connection[Any], seed_type: str, seed_id: str) -> list[str]:
    """Треки, от которых запущена волна: трек, альбом или исполнитель."""
    queries = {
        "track": "SELECT %s::uuid::text",
        "album": "SELECT id::text FROM tracks WHERE album_id = %s",
        "artist": "SELECT track_id::text FROM track_artists WHERE artist_id = %s",
    }
    q = queries.get(seed_type)
    if not q:
        return []
    return [r[0] for r in conn.execute(q, (seed_id,))]


class SessionStore:
    def __init__(self, url: str) -> None:
        self.rdb = redis.Redis.from_url(url, socket_connect_timeout=2, socket_timeout=2)

    def get(self, session_id: str) -> Session:
        raw = cast(bytes | None, self.rdb.get(f"wave:session:{session_id}"))
        if not raw:
            return Session()
        data = json.loads(raw)
        return Session(**data)

    def put(self, session_id: str, s: Session) -> None:
        self.rdb.set(f"wave:session:{session_id}", json.dumps(asdict(s)), ex=SESSION_TTL_S)