feat: add two-stage embedding retrieval and candidate ranking

app/api/routes.py

@@ -25,6 +25,7 @@
     recommendations_total,
 )
 from app.core.model import DeepSequenceModel
+from app.core.retrieval import ExactEmbeddingRetriever
 from app.core.security import api_key_is_valid
 from app.core.serving import AdmissionController, RateLimiter, RecommendationCache
 
@@ -38,6 +39,7 @@ class ModelRuntime:
     model_version: str
     trained: bool
     popular_items: list[str]
+    retriever: ExactEmbeddingRetriever
 
 
 _runtime: ModelRuntime | None = None
@@ -62,6 +64,7 @@ def init_model(
         model_version=model_version,
         trained=trained,
         popular_items=popular_items or list(processor.export_vocabulary())[: settings.max_top_k],
+        retriever=ExactEmbeddingRetriever(settings.retrieval_candidate_pool_size),
     )
 
 
@@ -163,10 +166,17 @@ def recommend(
     try:
         tensor = _runtime.processor.to_tensor(known_items)
         infer_started = time.perf_counter()
-        indices = _runtime.model.recommend(
+        excluded_ids = [_runtime.processor.item_to_idx(item) for item in known_items]
+        candidates = _runtime.retriever.retrieve(
+            _runtime.model,
             tensor,
             top_k=req.top_k,
-            exclude_ids=[_runtime.processor.item_to_idx(item) for item in known_items],
+            exclude_ids=excluded_ids,
+        )
+        indices = _runtime.model.rank_candidates(
+            tensor,
+            candidates.candidate_ids,
+            top_k=req.top_k,
         )
         inference_ms = (time.perf_counter() - infer_started) * 1_000
         model_inference_latency.observe(inference_ms / 1_000)

app/core/config.py

@@ -15,6 +15,7 @@ class Settings(BaseSettings):
     max_sequence_length: int = 50
     top_k: int = 10
     max_top_k: int = 50
+    retrieval_candidate_pool_size: int = 100
     model_bundle_path: str = "models/current"
     max_inference_ms: float = 250.0
     max_concurrent_inferences: int = 8

app/core/model.py

@@ -83,6 +83,36 @@ def forward(self, item_seq: torch.Tensor) -> torch.Tensor:
         logits[:, self.padding_idx] = float("-inf")
         return logits
 
+    @torch.no_grad()
+    def rank_candidates(
+        self,
+        item_seq: torch.Tensor,
+        candidate_ids: list[int],
+        *,
+        top_k: int,
+    ) -> list[int]:
+        """Apply the sequence ranker only to retrieval-stage candidates."""
+
+        if item_seq.ndim != 2 or not item_seq.ne(self.padding_idx).any():
+            raise ValueError("A recommendation requires at least one known item")
+        unique_candidates = list(dict.fromkeys(candidate_ids))
+        if not unique_candidates:
+            raise ValueError("candidate_ids must not be empty")
+        if any(candidate < 1 or candidate > self.num_items for candidate in unique_candidates):
+            raise ValueError("candidate_ids must be known non-padding item IDs")
+        if not 1 <= top_k <= len(unique_candidates):
+            raise ValueError("top_k must not exceed the candidate pool")
+
+        logits = self.forward(item_seq)
+        candidates = torch.tensor(
+            unique_candidates,
+            dtype=torch.long,
+            device=logits.device,
+        )
+        candidate_scores = logits[0, candidates]
+        positions = torch.topk(candidate_scores, k=top_k).indices
+        return candidates[positions].tolist()
+
     @torch.no_grad()
     def recommend(
         self,

app/core/retrieval.py

@@ -0,0 +1,79 @@
+"""Candidate retrieval for the two-stage recommendation serving path."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from app.core.model import DeepSequenceModel
+
+
+@dataclass(frozen=True)
+class RetrievalResult:
+    """Candidate IDs emitted by a retriever before sequence-aware ranking."""
+
+    candidate_ids: list[int]
+
+
+class ExactEmbeddingRetriever:
+    """Retrieve a bounded candidate pool with normalized item-embedding similarity.
+
+    This is an exact in-memory retriever, not an ANN or FAISS implementation.  It
+    establishes a replaceable retrieval boundary while keeping the default bundle
+    self-contained and deterministic for the current catalogue scale.
+    """
+
+    def __init__(self, candidate_pool_size: int) -> None:
+        if candidate_pool_size < 1:
+            raise ValueError("candidate_pool_size must be at least one")
+        self.candidate_pool_size = candidate_pool_size
+
+    @torch.no_grad()
+    def retrieve(
+        self,
+        model: DeepSequenceModel,
+        item_sequence: torch.Tensor,
+        *,
+        top_k: int,
+        exclude_ids: list[int] | None = None,
+    ) -> RetrievalResult:
+        """Return eligible IDs ordered by embedding-similarity score."""
+
+        if item_sequence.ndim != 2 or item_sequence.shape[0] != 1:
+            raise ValueError("Retrieval requires one padded recommendation sequence")
+        if not 1 <= top_k <= model.num_items:
+            raise ValueError(f"top_k must be between 1 and {model.num_items}")
+
+        history_ids = item_sequence[0]
+        known_mask = history_ids.ne(model.padding_idx)
+        if not known_mask.any():
+            raise ValueError("Retrieval requires at least one known item")
+
+        excluded = {
+            item_id
+            for item_id in (exclude_ids or [])
+            if 1 <= item_id <= model.num_items
+        }
+        eligible_count = model.num_items - len(excluded)
+        if top_k > eligible_count:
+            raise ValueError("top_k exceeds the remaining eligible catalogue")
+
+        query = model.embedding(history_ids[known_mask]).mean(dim=0)
+        query = F.normalize(query, dim=0, eps=1e-12)
+        catalogue = F.normalize(
+            model.embedding.weight[1 : model.num_items + 1],
+            dim=1,
+            eps=1e-12,
+        )
+        scores = torch.matmul(catalogue, query)
+        if excluded:
+            scores[[item_id - 1 for item_id in excluded]] = float("-inf")
+
+        candidate_count = min(
+            max(top_k, self.candidate_pool_size),
+            eligible_count,
+        )
+        retrieved = torch.topk(scores, k=candidate_count).indices.add(1).tolist()
+        return RetrievalResult(candidate_ids=retrieved)

docs/architecture-walkthrough.md

@@ -0,0 +1,91 @@
+# Architecture Walkthrough: Two-Stage Serving
+
+This is an 8-minute walkthrough script and diagram set for the current implementation. It describes
+the code in this repository; it does not claim that an ANN index, external feature store, or runtime
+reasoning service exists.
+
+## 0:00–1:00 — Problem and decision
+
+The original serving path applied the sequence model's dense output projection to every catalogue
+item. That is simple and correct for the current small bundle, but its work grows with catalogue
+size and provides no replaceable candidate boundary. The two-stage design separates *recall*
+from *precision*: retrieval produces a bounded set of plausible item IDs, then ranking spends the
+sequence model's richer scoring capacity only on that set.
+
+```mermaid
+flowchart LR
+    H["Known interaction history"] --> Q["Embedding query"]
+    Q --> R["Stage 1: exact embedding retrieval"]
+    R --> C["Bounded candidate IDs"]
+    C --> K["Stage 2: BiLSTM + attention ranking"]
+    H --> K
+    K --> O["Top-K recommendations"]
+```
+
+## 1:00–3:00 — Retrieval is not ranking
+
+`app/core/retrieval.py` implements `ExactEmbeddingRetriever`. It averages embeddings of known
+history items, normalizes that query, compares it with normalized catalogue embeddings, excludes
+items already seen, and returns a bounded candidate list. Its default is an exact in-memory vector
+scan. That makes behavior reproducible and keeps model bundles self-contained, but it is **not**
+FAISS, HNSW, or another approximate-nearest-neighbor index.
+
+Retrieval optimizes candidate recall and cost. It can return items that are semantically or
+behaviorally near the history, but it has no access to the full sequence-ordering signal used by
+the ranker. The stable `RetrievalResult` contract is the seam where an ANN implementation can
+later be added after index lifecycle, recall, freshness, and operational evidence are available.
+
+## 3:00–5:00 — Candidate-only ranking
+
+`DeepSequenceModel.rank_candidates` runs the existing padding-aware bidirectional LSTM and
+attention model, then gathers scores only for the retrieved IDs. This preserves the existing model
+and training compatibility while making the ranking stage explicit. The API uses retrieval after
+admission control and before decoding recommendations; cache, authentication, rate limits, and
+fallback behavior remain unchanged.
+
+```mermaid
+sequenceDiagram
+    participant Client
+    participant API as FastAPI route
+    participant Retriever as ExactEmbeddingRetriever
+    participant Ranker as DeepSequenceModel
+
+    Client->>API: history + top_k
+    API->>API: validate, authorize, rate-limit, cache check
+    API->>Retriever: retrieve(history, exclusions, top_k)
+    Retriever-->>API: candidate_ids
+    API->>Ranker: rank_candidates(history, candidate_ids, top_k)
+    Ranker-->>API: ordered item IDs
+    API-->>Client: recommendations + model version + latency
+```
+
+The tradeoff is intentional: the first stage is logically separated but still scans all embeddings,
+so it does not yet deliver the latency or memory profile of a production ANN system. Candidate-pool
+size is configurable with `RETRIEVAL_CANDIDATE_POOL_SIZE`; it should be measured against
+Recall@K and latency before changing it in production.
+
+## 5:00–6:30 — Reasoning and explanations
+
+No runtime `reasoning/` package or per-recommendation explanation endpoint exists in the current
+repository. That is deliberate in this walkthrough: an LSTM score is not a causal explanation, and
+the service should not invent a user-facing reason from hidden states. Issue #16 records the
+separate evidence-backed explanation contract, including privacy review and insufficient-history
+handling. Until that work is implemented and tested, the only trustworthy response-level evidence
+is model version, fallback state, cache state, and bounded request latency.
+
+## 6:30–8:00 — Engineering tradeoffs and next decisions
+
+- **Exact retrieval now:** easy to test and bundle; unsuitable for large catalogues without an ANN
+  index and index-refresh lifecycle.
+- **Sequence ranker retained:** preserves current training artifacts; a future ranker change needs
+  temporal evaluation against the popularity baseline.
+- **No fabricated confidence:** offline ranking quality and per-recommendation confidence are
+  different measurements.
+- **No hidden reasoning:** user-facing explanations must be constrained to permitted evidence,
+  not chain-of-thought or unvalidated causal language.
+- **Safety preserved:** existing authentication, credential-derived rate limiting, admission
+  control, cache keying, fallback, and model-bundle checks remain the API boundary.
+
+Before a large-catalogue deployment, add an evaluated ANN backend, version and validate its index
+with the model bundle, measure candidate recall and end-to-end p95/p99 latency, and keep the
+candidate-ranker contract stable during rollout.

tests/test_retrieval.py

@@ -0,0 +1,67 @@
+import pytest
+import torch
+
+from app.core.model import DeepSequenceModel
+from app.core.retrieval import ExactEmbeddingRetriever
+
+
+def _model() -> DeepSequenceModel:
+    model = DeepSequenceModel(
+        num_items=4,
+        embedding_dim=2,
+        hidden_dim=2,
+        num_layers=1,
+        dropout=0.0,
+    ).eval()
+    with torch.no_grad():
+        model.embedding.weight.copy_(
+            torch.tensor(
+                [
+                    [0.0, 0.0],
+                    [1.0, 0.0],
+                    [0.9, 0.1],
+                    [0.0, 1.0],
+                    [-1.0, 0.0],
+                ]
+            )
+        )
+        model.output_proj.weight.zero_()
+        model.output_proj.bias.copy_(
+            torch.tensor([float("-inf"), 0.1, 0.9, 0.2, 0.8])
+        )
+    return model
+
+
+def test_embedding_retriever_excludes_history_and_bounds_candidate_pool() -> None:
+    retriever = ExactEmbeddingRetriever(candidate_pool_size=3)
+
+    result = retriever.retrieve(
+        _model(),
+        torch.tensor([[0, 0, 1]]),
+        top_k=1,
+        exclude_ids=[1],
+    )
+
+    assert result.candidate_ids == [2, 3, 4]
+
+
+def test_ranker_only_orders_retrieved_candidates() -> None:
+    ranked = _model().rank_candidates(
+        torch.tensor([[0, 0, 1]]),
+        [1, 3, 4],
+        top_k=2,
+    )
+
+    assert ranked == [4, 3]
+
+
+def test_retriever_rejects_an_impossible_request() -> None:
+    retriever = ExactEmbeddingRetriever(candidate_pool_size=2)
+
+    with pytest.raises(ValueError, match="remaining eligible"):
+        retriever.retrieve(
+            _model(),
+            torch.tensor([[0, 0, 1]]),
+            top_k=4,
+            exclude_ids=[1],
+        )