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

"""Сервис «Моей волны»: собирает пачку из каталога, профиля, сессии и контекста."""

from __future__ import annotations

import logging
from typing import Literal

import numpy as np
import psycopg
from pydantic import BaseModel, Field

from pipu_rec.train.ranker import RankerHolder
from pipu_rec.wave.contexts import Contexts, Rule
from pipu_rec.wave.engine import Context, Pick, build_batch, model_version, taste_centroids, update_session
from pipu_rec.wave.store import IndexHolder, SessionStore, load_profile, seed_tracks

MAX_SERVED_MEMORY = 300

log = logging.getLogger(__name__)

Outcome = Literal["complete", "skip", "skip_early", "like", "dislike"]


class Feedback(BaseModel):
    track_id: str
    outcome: Outcome


class Seed(BaseModel):
    type: Literal["track", "album", "artist"]
    id: str


class BatchRequest(BaseModel):
    user_id: str
    session_id: str
    activity: str | None = None
    mood: str | None = None
    seed: Seed | None = None
    exclude_ids: list[str] = Field(default_factory=list, max_length=1000)
    friend_track_ids: list[str] = Field(default_factory=list, max_length=30)
    feedback: list[Feedback] = Field(default_factory=list, max_length=200)
    n: int = Field(default=5, ge=1, le=20)


class BatchTrack(BaseModel):
    id: str
    source: str
    score: float


class BatchResponse(BaseModel):
    tracks: list[BatchTrack]
    model_ver: str


class WaveService:
    def __init__(self, dsn: str, redis_url: str, contexts: Contexts) -> None:
        self.dsn = dsn
        self.index = IndexHolder(dsn)
        self.sessions = SessionStore(redis_url)
        self.contexts = contexts
        self.ranker = RankerHolder(dsn)

    def reload(self) -> None:
        """После ночного обучения: перечитать каталог (векторы ALS) и активный ранжировщик."""
        self.index.refresh()
        self.ranker.reload()

    def rules(self, activity: str | None, mood: str | None) -> list[Rule]:
        """Занятие и настроение комбинируются — пересечение профилей; неизвестный ключ игнорируется."""
        rules: list[Rule] = []
        if activity and activity in self.contexts.activities:
            rules += self.contexts.activities[activity].rules
        if mood and mood in self.contexts.moods:
            rules += self.contexts.moods[mood].rules
        return rules

    def batch(self, req: BatchRequest) -> BatchResponse:
        index = self.index.get()
        with psycopg.connect(self.dsn) as conn:
            profile = load_profile(conn, req.user_id)
            seeds = seed_tracks(conn, req.seed.type, req.seed.id) if req.seed else None
        session = self.sessions.get(req.session_id)
        update_session(
            index, session, [(f.track_id, f.outcome) for f in req.feedback], taste_centroids(index, profile)
        )
        ranker = self.ranker.get()
        picks: list[Pick] = build_batch(
            index,
            profile,
            session,
            self.rules(req.activity, req.mood),
            n=req.n,
            exclude=set(req.exclude_ids),
            seed_ids=seeds,
            friend_ids=req.friend_track_ids,
            rng=np.random.default_rng(),
            ranker=ranker,
            ctx=Context(has_activity=bool(req.activity), has_mood=bool(req.mood)),
        )
        self.log_features(req.session_id, picks)
        session.served = (session.served + [p.id for p in picks])[-MAX_SERVED_MEMORY:]
        artists = [index.artist[index.pos[p.id]] for p in picks]
        session.recent_artists = (session.recent_artists + artists)[-10:]
        self.sessions.put(req.session_id, session)
        return BatchResponse(
            tracks=[BatchTrack(id=p.id, source=p.source, score=p.score) for p in picks],
            model_ver=model_version(index, profile, ranker),
        )

    def log_features(self, session_id: str, picks: list[Pick]) -> None:
        """Признаки показанных треков — обучающие данные ранжировщика v2. Сбой записи не ломает волну."""
        rows = [(session_id, p.id, p.features) for p in picks if p.features is not None]
        if not rows:
            return
        try:
            with psycopg.connect(self.dsn, autocommit=True) as conn, conn.cursor() as cur:
                cur.executemany(
                    "INSERT INTO rec_impression_features (wave_session_id, track_id, features) "
                    "VALUES (%s, %s, %s) "
                    "ON CONFLICT DO NOTHING",
                    rows,
                )
        except psycopg.Error as e:
            log.warning("impression features not saved: %s", e)