"""Загрузка каталога, профиля и сессии волны из 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), e.cf::text 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], raw_cf=[parse_vector(r[9]) 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() recent_likes: list[str] = [] for tid, reaction in conn.execute( "SELECT track_id::text, reaction::text FROM track_reactions WHERE user_id = %s " "ORDER BY created_at DESC", (user_id,), ): if reaction == "like": liked.add(tid) recent_likes.append(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, taste_cf::text, bandit, onboarding FROM user_profiles_rec WHERE user_id = %s", (user_id,), ).fetchone() if row and row[3] and len(scores) < ONBOARDING_UNTIL: apply_onboarding(conn, scores, row[3]) exploration = float(row[0]) if row else 0.15 taste_cf = parse_vector(row[1]) if row else None bandit = {str(k): [float(x) for x in v] for k, v in (row[2] or {}).items()} if row else {} 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, taste_cf=taste_cf.tolist() if taste_cf is not None else None, bandit=bandit, recent_likes=recent_likes, ) # Пока своих оценок меньше, выбор онбординга задаёт начальный вкус (docs/recommendations.md#7). ONBOARDING_UNTIL = 30 ONBOARDING_LIMIT = 300 def apply_onboarding(conn: psycopg.Connection[Any], scores: dict[str, float], choice: dict[str, Any]) -> None: """Псевдооценки: треки выбранных исполнителей +1, выбранных жанров +0.5 (до ONBOARDING_LIMIT треков). Собственные оценки не перезаписываются; в «слышанное» эти треки не попадают.""" artists = [str(a) for a in choice.get("artist_ids") or []] genres = [str(g).lower() for g in choice.get("genres") or []] pseudo: dict[str, float] = {} if artists: for (tid,) in conn.execute( "SELECT DISTINCT track_id::text FROM track_artists WHERE artist_id::text = ANY(%s) LIMIT %s", (artists, ONBOARDING_LIMIT), ): pseudo[tid] = pseudo.get(tid, 0.0) + 1.0 if genres: for (tid,) in conn.execute( """ SELECT DISTINCT tg.track_id::text FROM track_genres tg JOIN genres g ON g.id = tg.genre_id WHERE lower(g.name) = ANY(%s) LIMIT %s """, (genres, ONBOARDING_LIMIT), ): pseudo[tid] = pseudo.get(tid, 0.0) + 0.5 for tid, v in pseudo.items(): scores.setdefault(tid, v) 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)