全量物品池 (数十亿)
    ↓  候选生成 (Candidate Generation)
候选集合 (~1000 items)
    ↓  粗排 (Pre-ranking)
粗排结果 (~100 items)
    ↓  精排 (Ranking)
精排结果 (~20 items)
    ↓  重排 / 业务逻辑 (Re-ranking)
最终展示 (~10 items)

# two_tower_retrieval.py
# The Two-Tower (Dual Encoder) model architecture used by YouTube DNN, TikTok, etc.
# User tower and item tower are trained jointly but served independently.

import numpy as np
from dataclasses import dataclass
from typing import List
import faiss  # pip install faiss-cpu

@dataclass
class UserEmbedding:
    user_id: str
    vector: np.ndarray  # shape: (128,)
    updated_at: float   # unix timestamp

@dataclass 
class ItemEmbedding:
    item_id: str
    vector: np.ndarray  # shape: (128,)

class TwoTowerRetriever:
    """
    Manages ANN index for fast approximate nearest neighbor retrieval.
    
    In production: user embeddings updated every 5 min (stream processing),
    item embeddings updated every 6 hours (batch job).
    """
    
    def __init__(self, embedding_dim: int = 128):
        self.embedding_dim = embedding_dim
        
        # FAISS IVF (Inverted File) index — production choice for >10M items
        # IVFFlat: split vectors into n_clusters, only search relevant clusters
        # Tradeoff: ~10% accuracy loss vs exhaustive search, 100x faster
        n_clusters = 1000  # rule of thumb: sqrt(n_items) for 10M items
        quantizer = faiss.IndexFlatIP(embedding_dim)  # Inner product = cosine similarity (on normalized vectors)
        self.index = faiss.IndexIVFFlat(quantizer, embedding_dim, n_clusters, faiss.METRIC_INNER_PRODUCT)
        
        self.item_ids: List[str] = []  # parallel array: index position → item_id
        self.is_trained = False
        
    def build_index(self, item_embeddings: List[ItemEmbedding]):
        """
        Offline batch job: (re)build the ANN index.
        Called every 6 hours, swapped atomically (blue-green pattern).
        """
        vectors = np.array([e.vector for e in item_embeddings], dtype='float32')
        
        # Normalize to unit vectors — converts cosine similarity to inner product
        faiss.normalize_L2(vectors)
        
        # Train the IVF quantizer (learns cluster centroids via k-means)
        # Requires at least 30 * n_clusters training samples
        self.index.train(vectors)
        self.index.add(vectors)
        
        self.item_ids = [e.item_id for e in item_embeddings]
        self.is_trained = True
        
        # Production note: after building, serialize index to S3
        # faiss.write_index(self.index, '/tmp/item_index.faiss')
        
    def retrieve(self, user_embedding: UserEmbedding, top_k: int = 500) -> List[str]:
        """
        Online serving: retrieve top_k candidates for a given user.
        Target latency: <5ms for 10M items.
        """
        if not self.is_trained:
            raise RuntimeError("Index not built. Run build_index() first.")
        
        query = user_embedding.vector.reshape(1, -1).astype('float32')
        faiss.normalize_L2(query)  # same normalization as at index time
        
        # nprobe: how many clusters to search. Higher = more accurate, slower.
        # Production: nprobe=50 gives ~95% recall@500 with <5ms latency
        self.index.nprobe = 50
        
        distances, indices = self.index.search(query, top_k)
        
        # Convert flat indices back to item IDs; filter -1 (empty slots)
        return [
            self.item_ids[idx] 
            for idx in indices[0] 
            if idx != -1
        ]


class RankingModel:
    """
    Simplified ranking model interface.
    In production: DLRM (Deep Learning Recommendation Model) or 
    a gradient boosted tree (XGBoost/LightGBM) for pre-ranking.
    """
    
    def score(self, user_features: dict, item_features: dict) -> float:
        """
        Returns P(engagement | user, item).
        
        Features include:
        - User: embedding, historical CTR, session context, demographic
        - Item: embedding, global CTR, freshness, category, creator follower count
        - Cross features: user × item category match, historical interactions
        """
        # Placeholder — real model would call a TF Serving / TorchServe endpoint
        # or run a local ONNX model for sub-millisecond latency
        raise NotImplementedError("Load your model weights here")


class RecommendationPipeline:
    """
    Orchestrates the full retrieval + ranking funnel.
    """
    
    def __init__(self, retriever: TwoTowerRetriever, ranker: RankingModel):
        self.retriever = retriever
        self.ranker = ranker
        
    def get_recommendations(
        self, 
        user_embedding: UserEmbedding,
        user_features: dict,
        item_features_store: dict,  # item_id → features
        n_final: int = 20
    ) -> List[str]:
        
        # Stage 1: Candidate Generation (~1000 items)
        candidates = self.retriever.retrieve(user_embedding, top_k=500)
        # (In practice, merge with CF-based and trending pools here)
        
        # Stage 2: Ranking (score all candidates)
        scored = []
        for item_id in candidates:
            if item_id not in item_features_store:
                continue  # item not in catalog (e.g., deleted)
            score = self.ranker.score(user_features, item_features_store[item_id])
            scored.append((item_id, score))
        
        # Sort by score descending
        scored.sort(key=lambda x: x[1], reverse=True)
        
        # Stage 3: Re-ranking (diversity injection)
        final = self._inject_diversity(scored, n_final)
        return final
    
    def _inject_diversity(self, scored: List[tuple], n: int) -> List[str]:
        """
        Ensure we don't show all items from the same creator/category.
        Simple approach: cap same-creator items to 2 in the final slate.
        """
        # In production: use MMR (Maximal Marginal Relevance) or DPP (Determinantal Point Process)
        result = []
        seen_count = {}  # creator_id → count
        
        for item_id, score in scored:
            if len(result) >= n:
                break
            # placeholder diversity check
            result.append(item_id)
            
        return result

# feature_store_client.py
# Production systems use a feature store (Feast, Tecton, or custom) 
# to serve features with consistent definitions between training and serving.

import redis
import json
from typing import Dict, Any, Optional

class FeatureStoreClient:
    """
    Unified interface for online (Redis) and offline (S3/Hive) features.
    
    Key insight: Training/serving skew is the #1 cause of recommendation 
    quality degradation. A feature store ensures the same feature computation 
    logic is used at both training time and serving time.
    """
    
    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        
    def get_user_features(self, user_id: str) -> Optional[Dict[str, Any]]:
        """
        Fetch real-time user features from Redis.
        
        What's stored:
        - user_embedding: 128-dim float vector (updated every 5 min by Flink)
        - recent_categories: top-3 categories clicked in last hour (stream)
        - session_count_today: integer counter (atomic increment)
        - historical_ctr: click-through rate across all sessions (batch, updated hourly)
        """
        raw = self.redis.get(f"user_features:{user_id}")
        if raw is None:
            # Cold start: new user, return default features
            return self._default_user_features()
        return json.loads(raw)
    
    def get_item_features(self, item_id: str) -> Optional[Dict[str, Any]]:
        """
        Fetch item features.
        
        What's stored:
        - item_embedding: 128-dim float vector (updated every 6h by batch job)
        - global_ctr_7d: click-through rate over last 7 days (batch)
        - freshness_score: decay function of item age (recomputed every 15 min)
        - creator_id: string (static)
        - category: enum (static)
        """
        raw = self.redis.get(f"item_features:{item_id}")
        if raw is None:
            return None  # item not in catalog
        return json.loads(raw)
    
    def _default_user_features(self) -> Dict[str, Any]:
        """
        For new users: use global popular items (popularity-based fallback).
        This is the standard cold-start solution.
        """
        return {
            "user_embedding": [0.0] * 128,  # zero vector → will match popular items
            "recent_categories": [],
            "session_count_today": 0,
            "historical_ctr": 0.05  # global average CTR
        }
    
    def update_user_embedding(self, user_id: str, new_embedding: list, ttl_seconds: int = 86400):
        """
        Called by Flink stream job after processing user events.
        TTL ensures stale embeddings don't persist indefinitely.
        """
        current = self.get_user_features(user_id) or {}
        current["user_embedding"] = new_embedding
        self.redis.setex(
            f"user_features:{user_id}",
            ttl_seconds,
            json.dumps(current)
        )