Merge: F1 backend — expose query_position and query_token_id (#4)

backend/model_service.py

@@ -140,10 +140,35 @@ class MatrixCache:
         for older requests.
         """
         with self._lock:
-            self._request_meta[request_id] = {
+            existing = self._request_meta.get(request_id, {})
+            existing.update({
                 "num_heads": num_heads,
                 "num_layers": num_layers,
                 "model_id": model_id,
+            })
+            self._request_meta[request_id] = existing
+
+    def set_step_query_info(self, request_id: str, step: int,
+                            query_position: int, query_token_id: int):
+        """
+        Record the predicting query position for a given generation step.
+
+        At step `s`, the model's forward pass runs over the prompt + first
+        `s` generated tokens; the logits at the LAST query position select
+        the step-`s` token. The cached attention[step][layer][head][-1] is
+        therefore the attention pattern of *that* query position, not of
+        the predicted token (which doesn't exist in the sequence yet).
+        Storing the query position and token id makes the provenance of
+        attention rows explicit when the Why this token? panel surfaces
+        them — supporting the "grounded hypotheses, not causal claims"
+        framing required by the project.
+        """
+        with self._lock:
+            meta = self._request_meta.setdefault(request_id, {})
+            steps = meta.setdefault("steps", {})
+            steps[step] = {
+                "query_position": query_position,
+                "query_token_id": query_token_id,
             }
 
     def get_request_metadata(self, request_id: str) -> Optional[dict]:
@@ -152,6 +177,18 @@ class MatrixCache:
             meta = self._request_meta.get(request_id)
             return dict(meta) if meta else None
 
+    def get_step_query_info(self, request_id: str, step: int) -> Optional[dict]:
+        """Return {query_position, query_token_id} for a step, or None."""
+        with self._lock:
+            meta = self._request_meta.get(request_id)
+            if not meta:
+                return None
+            steps = meta.get("steps")
+            if not steps:
+                return None
+            entry = steps.get(step)
+            return dict(entry) if entry else None
+
     def get(self, request_id: str, step: int, layer: int, head: int) -> Optional[dict]:
         """Retrieve matrix data, returning None if expired or not found."""
         key = f"{request_id}:{step}:{layer}:{head}"
@@ -2048,6 +2085,16 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                     output_hidden_states=True
                 )
 
+                # Record the query position whose logits will select this
+                # step's token. attention[layer][head][-1] is the row of
+                # this query — exposing it in /attention/row makes the
+                # provenance of "Why this token?" attention explicit.
+                query_position = current_ids.shape[1] - 1
+                query_token_id = current_ids[0, -1].item()
+                matrix_cache.set_step_query_info(
+                    request_id, step, query_position, query_token_id
+                )
+
                 # Get logits for next token
                 raw_logits = outputs.logits[0, -1, :].clone()  # Clone raw logits before any scaling
 
@@ -2952,6 +2999,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         output_hidden_states=True
                     )
 
+                    # Record the query position whose logits will select this
+                    # step's token (see non-SSE site for rationale).
+                    query_position = current_ids.shape[1] - 1
+                    query_token_id = current_ids[0, -1].item()
+                    matrix_cache.set_step_query_info(
+                        request_id, step, query_position, query_token_id
+                    )
+
                     # Get logits for next token
                     raw_logits = outputs.logits[0, -1, :].clone()  # Clone raw logits before any scaling
 
@@ -3913,12 +3968,21 @@ async def get_attention_row(
     - layer: Layer index
     - head: Head index (null if aggregated)
     - aggregate_mode: Mode used if aggregated (null otherwise)
+    - query_position: Sequence index whose attention this row represents.
+      The model selects the step-N token from logits at this position;
+      the cached attention[-1] is therefore *this* position's attention,
+      not the predicted token's self-attention.
+    - query_token_id: Token id at query_position (last context token).
     """
     # Aggregation parameters (num_heads in particular) come from per-request
     # metadata recorded at analysis time, not from the live model config.
     # This keeps aggregate rows correct even after the user switches models
     # mid-session — the row is averaged over the heads of the model that
     # actually produced the cached attention.
+    query_info = matrix_cache.get_step_query_info(request_id, step)
+    query_position = query_info.get("query_position") if query_info else None
+    query_token_id = query_info.get("query_token_id") if query_info else None
+
     if head is not None:
         # Fetch specific head — no aggregation, no metadata lookup needed.
         attention_row = matrix_cache.get_attention_row(request_id, step, layer, head)
@@ -3934,7 +3998,9 @@ async def get_attention_row(
             "seq_len": len(attention_row),
             "layer": layer,
             "head": head,
-            "aggregate_mode": None
+            "aggregate_mode": None,
+            "query_position": query_position,
+            "query_token_id": query_token_id,
         }
     else:
         # Aggregate across all heads. num_heads is derived inside
@@ -3954,7 +4020,9 @@ async def get_attention_row(
             "seq_len": len(attention_row),
             "layer": layer,
             "head": None,
-            "aggregate_mode": aggregate_mode
+            "aggregate_mode": aggregate_mode,
+            "query_position": query_position,
+            "query_token_id": query_token_id,
         }