Add SSE streaming endpoint for real-time analysis progress

Add /analyze/research/attention/stream endpoint that streams progress
events during attention analysis via Server-Sent Events (SSE).

Progress events are emitted at each stage:
- tokenizing: After prompt tokenization, with token count
- generating: Per-token, with step index and total steps
- extracting: Per-layer (on last token), with layer/head counts
- serializing: During response building, with size estimate
- complete: When analysis finishes, with timing and size metadata
- result: Final payload with complete analysis data

Each event includes:
- type, stage, totalStages, progress, stageProgress, detail
- metadata object with stage-specific info (counts, sizes, timing)

Benefits:
- Frontend gets real progress updates instead of fake animation
- Users see exactly what's happening: "Processing layer 15/40"
- Response size displayed during serialization
- No more "stuck at 95%" - progress reflects actual work

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

backend/model_service.py

@@ -4,6 +4,7 @@ Combines model loading, generation, and trace extraction into a single service
 """
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, BackgroundTasks, HTTPException, Depends
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import asyncio
@@ -1920,6 +1921,460 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
         logger.error(traceback.format_exc())
         raise HTTPException(status_code=500, detail=str(e))
 
+
+def sse_event(event_type: str, **kwargs) -> str:
+    """Format data as SSE event"""
+    data = {'type': event_type, 'timestamp': int(time.time() * 1000), **kwargs}
+    return f"data: {json.dumps(data)}\n\n"
+
+
+@app.post("/analyze/research/attention/stream")
+async def analyze_research_attention_stream(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
+    """
+    SSE Streaming version of Research-Grade Attention Analysis
+
+    Emits progress events during each stage:
+    - tokenizing: Initial tokenization
+    - generating: Per-token generation progress
+    - extracting: Per-layer attention extraction
+    - serializing: Building response
+    - complete: Analysis finished
+    - result: Final data payload
+    """
+    async def event_generator():
+        try:
+            import time
+            start_time = time.time()
+
+            # Get parameters
+            prompt = request.get("prompt", "def quicksort(arr):")
+            max_tokens = request.get("max_tokens", 8)
+            temperature = request.get("temperature", 0.7)
+
+            logger.info(f"[SSE] Research attention analysis: prompt_len={len(prompt)}, max_tokens={max_tokens}")
+
+            # === STAGE 1: TOKENIZING ===
+            yield sse_event('tokenizing', stage=1, totalStages=5, progress=2,
+                           stageProgress=0, detail=f'Tokenizing {len(prompt)} characters...')
+
+            # Get model config for prompt formatting
+            from .model_config import get_model_config
+            from .prompt_formatter import format_prompt
+            model_config = get_model_config(manager.model_id)
+
+            # Get optional system prompt override from request
+            system_prompt_override = request.get("system_prompt")
+
+            # Format prompt using the unified formatter
+            formatted_prompt = format_prompt(
+                prompt=prompt,
+                model_config=model_config or {},
+                tokenizer=manager.tokenizer,
+                system_prompt_override=system_prompt_override
+            )
+
+            prompt_style = model_config.get("prompt_style", "completion") if model_config else "completion"
+
+            # Use model's recommended temperature for instruction models
+            if model_config and "recommended_temperature" in model_config:
+                temperature = model_config["recommended_temperature"]
+
+            # Tokenize and prepare - use MistralTokenizer for Devstral
+            if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
+                system_prompt = system_prompt_override or (model_config.get("system_prompt") if model_config else "")
+                prompt_token_ids = manager.mistral_tokenizer.encode_chat(system_prompt, prompt)
+                inputs = {"input_ids": torch.tensor([prompt_token_ids]).to(manager.device)}
+                prompt_length = len(prompt_token_ids)
+                prompt_tokens = [manager.mistral_tokenizer.decode_token(tid) for tid in prompt_token_ids]
+            else:
+                inputs = manager.tokenizer(formatted_prompt, return_tensors="pt").to(manager.device)
+                prompt_length = inputs["input_ids"].shape[1]
+                prompt_token_ids = inputs["input_ids"][0].tolist()
+                prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]
+
+            yield sse_event('tokenizing', stage=1, totalStages=5, progress=8,
+                           stageProgress=100, detail=f'Tokenized into {prompt_length} tokens',
+                           metadata={'tokenCount': prompt_length})
+            await asyncio.sleep(0)  # Yield to event loop
+
+            # Storage for generation
+            generated_token_ids = []
+            generated_tokens = []
+
+            # Model info
+            n_layers = len(list(manager.model.parameters()))
+            if hasattr(manager.model.config, 'n_layer'):
+                n_layers = manager.model.config.n_layer
+            elif hasattr(manager.model.config, 'num_hidden_layers'):
+                n_layers = manager.model.config.num_hidden_layers
+
+            n_heads = manager.model.config.n_head if hasattr(manager.model.config, 'n_head') else manager.model.config.num_attention_heads
+            d_model = manager.model.config.n_embd if hasattr(manager.model.config, 'n_embd') else manager.model.config.hidden_size
+            head_dim = d_model // n_heads
+
+            # === STAGE 2: GENERATING ===
+            layer_data_by_token = []
+            token_alternatives_by_step = []
+
+            # Hook system to capture Q/K/V matrices
+            qkv_captures = {}
+            hooks = []
+
+            def make_qkv_hook(layer_idx):
+                def hook(module, input, output):
+                    try:
+                        if output.dim() != 3:
+                            return
+                        batch_size, seq_len, hidden = output.shape
+                        expected_hidden = 3 * n_heads * head_dim
+                        if hidden != expected_hidden:
+                            return
+                        qkv = output.reshape(batch_size, seq_len, 3, n_heads, head_dim)
+                        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+                        qkv_captures[layer_idx] = {
+                            'q': q[0].detach().cpu(),
+                            'k': k[0].detach().cpu(),
+                            'v': v[0].detach().cpu()
+                        }
+                    except Exception:
+                        pass
+                return hook
+
+            # Register hooks
+            try:
+                if hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'h'):
+                    for layer_idx, layer in enumerate(manager.model.transformer.h):
+                        if hasattr(layer, 'attn') and hasattr(layer.attn, 'qkv_proj'):
+                            hook = layer.attn.qkv_proj.register_forward_hook(make_qkv_hook(layer_idx))
+                            hooks.append(hook)
+                        elif hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
+                            hook = layer.attn.c_attn.register_forward_hook(make_qkv_hook(layer_idx))
+                            hooks.append(hook)
+            except Exception as hook_error:
+                logger.warning(f"Could not register QKV hooks: {hook_error}")
+
+            with torch.no_grad():
+                current_ids = inputs["input_ids"]
+
+                for step in range(max_tokens):
+                    # Emit progress for this generation step
+                    step_progress = (step / max_tokens) * 100
+                    overall_progress = 10 + (step / max_tokens) * 20  # 10-30%
+                    yield sse_event('generating', stage=2, totalStages=5, progress=overall_progress,
+                                   stageProgress=step_progress,
+                                   detail=f'Generating token {step + 1}/{max_tokens}',
+                                   metadata={'stepIndex': step, 'totalSteps': max_tokens})
+                    await asyncio.sleep(0)
+
+                    qkv_captures.clear()
+
+                    # Forward pass with full outputs
+                    outputs = manager.model(
+                        current_ids,
+                        output_attentions=True,
+                        output_hidden_states=True
+                    )
+
+                    # Get logits for next token
+                    logits = outputs.logits[0, -1, :]
+
+                    # Apply temperature and sample
+                    if temperature > 0:
+                        logits = logits / temperature
+                    probs = torch.softmax(logits, dim=0)
+
+                    if temperature == 0:
+                        next_token_id = torch.argmax(probs, dim=-1).item()
+                    else:
+                        next_token_id = torch.multinomial(probs, 1).item()
+                    next_token_text = manager.tokenizer.decode([next_token_id], skip_special_tokens=False)
+
+                    generated_token_ids.append(next_token_id)
+                    generated_tokens.append(next_token_text)
+
+                    # Capture top-k token alternatives
+                    import math as math_module
+                    top_k = 5
+                    top_probs, top_indices = torch.topk(probs, k=min(top_k, len(probs)))
+                    alternatives = []
+                    for prob, idx in zip(top_probs.tolist(), top_indices.tolist()):
+                        token_text = manager.tokenizer.decode([idx], skip_special_tokens=False)
+                        alternatives.append({
+                            "token": token_text,
+                            "token_id": idx,
+                            "probability": prob,
+                            "log_probability": math_module.log(prob) if prob > 0 else float('-inf')
+                        })
+                    token_alternatives_by_step.append({
+                        "step": step,
+                        "selected_token": next_token_text,
+                        "selected_token_id": next_token_id,
+                        "alternatives": alternatives
+                    })
+
+                    # === STAGE 3: EXTRACTING (per layer within each token) ===
+                    layer_data_this_token = []
+
+                    for layer_idx in range(len(outputs.attentions)):
+                        # Emit extraction progress (within generating stage for combined progress)
+                        if step == max_tokens - 1:  # Only emit detailed layer progress on last token
+                            layer_progress = (layer_idx / len(outputs.attentions)) * 100
+                            overall_progress = 30 + (layer_idx / len(outputs.attentions)) * 40  # 30-70%
+                            yield sse_event('extracting', stage=3, totalStages=5, progress=overall_progress,
+                                           stageProgress=layer_progress,
+                                           detail=f'Processing layer {layer_idx + 1}/{len(outputs.attentions)}',
+                                           metadata={'layerIndex': layer_idx, 'totalLayers': len(outputs.attentions),
+                                                    'headsPerLayer': n_heads, 'stepIndex': step, 'totalSteps': max_tokens})
+                            if layer_idx % 5 == 0:  # Yield every 5 layers to avoid too many events
+                                await asyncio.sleep(0)
+
+                        layer_attn = outputs.attentions[layer_idx][0]
+                        current_hidden = outputs.hidden_states[layer_idx + 1]
+                        if current_hidden.dim() == 3:
+                            current_hidden = current_hidden[0]
+
+                        if layer_idx > 0:
+                            prev_hidden = outputs.hidden_states[layer_idx]
+                            if prev_hidden.dim() == 3:
+                                prev_hidden = prev_hidden[0]
+                            delta_norm = torch.norm(current_hidden - prev_hidden).item()
+                        else:
+                            delta_norm = None
+
+                        activation_magnitude = torch.norm(current_hidden).item()
+                        last_token_hidden = current_hidden[-1]
+                        activation_entropy = torch.std(last_token_hidden).item()
+                        hidden_state_norm = torch.norm(last_token_hidden).item()
+
+                        # Sanitize
+                        activation_magnitude = 0.0 if math.isnan(activation_magnitude) or math.isinf(activation_magnitude) else activation_magnitude
+                        activation_entropy = 0.0 if math.isnan(activation_entropy) or math.isinf(activation_entropy) else activation_entropy
+                        hidden_state_norm = 0.0 if math.isnan(hidden_state_norm) or math.isinf(hidden_state_norm) else hidden_state_norm
+                        if delta_norm is not None:
+                            delta_norm = 0.0 if math.isnan(delta_norm) or math.isinf(delta_norm) else delta_norm
+
+                        # Process heads
+                        critical_heads = []
+                        for head_idx in range(layer_attn.shape[0]):
+                            head_weights = layer_attn[head_idx, -1, :]
+                            max_weight = head_weights.max().item()
+                            entropy = -(head_weights * torch.log(head_weights + 1e-10)).sum().item()
+
+                            max_weight = 0.0 if math.isnan(max_weight) or math.isinf(max_weight) else max_weight
+                            entropy = 0.0 if math.isnan(entropy) or math.isinf(entropy) else entropy
+
+                            pattern_type = None
+                            confidence = 0.0
+
+                            if step > 0 and max_weight > 0.8:
+                                pattern_type = "induction"
+                                confidence = max_weight
+                            elif entropy < 1.0:
+                                pattern_type = "positional"
+                                confidence = 1.0 - entropy
+                            elif 1.0 <= entropy < 2.5:
+                                pattern_type = "semantic"
+                                confidence = min(1.0, entropy / 2.5)
+                            elif max_weight > 0.9 and head_weights[-2].item() > 0.85:
+                                pattern_type = "previous_token"
+                                confidence = head_weights[-2].item()
+
+                            confidence = 0.0 if math.isnan(confidence) or math.isinf(confidence) else confidence
+
+                            attention_matrix = layer_attn[head_idx].cpu().float().numpy().tolist()
+
+                            q_matrix = None
+                            k_matrix = None
+                            v_matrix = None
+                            if layer_idx in qkv_captures:
+                                q_matrix = qkv_captures[layer_idx]['q'][:, head_idx, :].float().numpy().tolist()
+                                k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy().tolist()
+                                v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy().tolist()
+
+                            critical_heads.append({
+                                "head_idx": head_idx,
+                                "entropy": entropy,
+                                "max_weight": max_weight,
+                                "attention_weights": attention_matrix,
+                                "q_matrix": q_matrix,
+                                "k_matrix": k_matrix,
+                                "v_matrix": v_matrix,
+                                "pattern": {"type": pattern_type, "confidence": confidence} if pattern_type else None
+                            })
+
+                        critical_heads.sort(key=lambda h: h["max_weight"], reverse=True)
+
+                        layer_pattern = None
+                        layer_fraction = (layer_idx + 1) / n_layers
+                        if layer_idx == 0:
+                            layer_pattern = {"type": "positional", "confidence": 0.78}
+                        elif layer_fraction <= 0.25 and step > 0:
+                            layer_pattern = {"type": "previous_token", "confidence": 0.65}
+                        elif layer_fraction <= 0.75:
+                            layer_pattern = {"type": "induction", "confidence": 0.87}
+                        else:
+                            layer_pattern = {"type": "semantic", "confidence": 0.92}
+
+                        layer_data_this_token.append({
+                            "layer_idx": layer_idx,
+                            "pattern": layer_pattern,
+                            "critical_heads": critical_heads,
+                            "activation_magnitude": activation_magnitude,
+                            "activation_entropy": activation_entropy,
+                            "hidden_state_norm": hidden_state_norm,
+                            "delta_norm": delta_norm
+                        })
+
+                    layer_data_by_token.append(layer_data_this_token)
+
+                    # Update inputs
+                    next_token_tensor = torch.tensor([[next_token_id]], dtype=torch.long, device=manager.device)
+                    current_ids = torch.cat([current_ids, next_token_tensor], dim=1)
+
+                    # Stop on EOS
+                    if next_token_id == manager.tokenizer.eos_token_id:
+                        break
+
+            # Clean up hooks
+            for hook in hooks:
+                hook.remove()
+
+            # === STAGE 4: SERIALIZING ===
+            yield sse_event('serializing', stage=4, totalStages=5, progress=75,
+                           stageProgress=0, detail='Building response data...')
+            await asyncio.sleep(0)
+
+            qkv_by_layer_head = {}
+            generation_time = time.time() - start_time
+
+            # Calculate token section boundaries
+            total_tokens = prompt_length + len(generated_token_ids)
+            system_prompt_text = system_prompt_override or (model_config.get("system_prompt") if model_config else None)
+
+            system_prompt_end = 0
+            if prompt_style == "instruction" and system_prompt_text:
+                if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
+                    try:
+                        no_system_tokens = manager.mistral_tokenizer.encode_chat("", prompt)
+                        system_prompt_end = prompt_length - len(no_system_tokens)
+                        system_prompt_end = max(0, min(system_prompt_end, prompt_length))
+                    except Exception:
+                        system_prompt_end = 0
+                else:
+                    total_chars = len(system_prompt_text or "") + len(prompt)
+                    if total_chars > 0:
+                        system_ratio = len(system_prompt_text or "") / total_chars
+                        system_prompt_end = int(prompt_length * system_ratio)
+
+            token_sections = {
+                "systemPrompt": {
+                    "start": 0,
+                    "end": system_prompt_end,
+                    "text": system_prompt_text,
+                    "tokenCount": system_prompt_end
+                },
+                "userPrompt": {
+                    "start": system_prompt_end,
+                    "end": prompt_length,
+                    "text": prompt,
+                    "tokenCount": prompt_length - system_prompt_end
+                },
+                "output": {
+                    "start": prompt_length,
+                    "end": total_tokens,
+                    "text": "".join(generated_tokens),
+                    "tokenCount": len(generated_token_ids)
+                }
+            }
+
+            yield sse_event('serializing', stage=4, totalStages=5, progress=82,
+                           stageProgress=50, detail='Building token metadata...')
+            await asyncio.sleep(0)
+
+            # Build token metadata
+            from .tokenizer_utils import TokenizerMetadata
+            token_metadata_builder = TokenizerMetadata(manager.tokenizer)
+
+            special_token_ids_set = {
+                manager.tokenizer.eos_token_id,
+                manager.tokenizer.bos_token_id,
+                manager.tokenizer.pad_token_id,
+                manager.tokenizer.unk_token_id
+            }
+
+            def build_token_data(token_ids, token_texts, token_type):
+                multi_split_flags = token_metadata_builder.is_multi_split_identifier(token_ids)
+                result = []
+                for i, (tid, t) in enumerate(zip(token_ids, token_texts)):
+                    bpe_pieces = token_metadata_builder.get_subword_pieces(tid)
+                    result.append({
+                        "text": t,
+                        "idx": tid,
+                        "bytes": len(t.encode('utf-8')),
+                        "type": token_type,
+                        "bpe_pieces": bpe_pieces,
+                        "is_special": tid in special_token_ids_set,
+                        "is_multi_split": multi_split_flags[i] if i < len(multi_split_flags) else False,
+                        "num_pieces": len(bpe_pieces),
+                    })
+                return result
+
+            # Build response
+            response = {
+                "prompt": prompt,
+                "promptTokens": build_token_data(prompt_token_ids, prompt_tokens, "prompt"),
+                "generatedTokens": build_token_data(generated_token_ids, generated_tokens, "generated"),
+                "tokenSections": token_sections,
+                "tokenAlternatives": token_alternatives_by_step,
+                "layersDataByStep": layer_data_by_token,
+                "layersData": layer_data_by_token[-1] if layer_data_by_token else [],
+                "qkvData": qkv_by_layer_head,
+                "modelInfo": {
+                    "numLayers": n_layers,
+                    "numHeads": n_heads,
+                    "modelDimension": d_model,
+                    "headDim": head_dim,
+                    "vocabSize": manager.model.config.vocab_size
+                },
+                "generationTime": generation_time,
+                "numTokensGenerated": len(generated_tokens)
+            }
+
+            # Estimate response size
+            response_json = json.dumps(sanitize_for_json(response))
+            response_size_bytes = len(response_json.encode('utf-8'))
+
+            yield sse_event('serializing', stage=4, totalStages=5, progress=90,
+                           stageProgress=100, detail=f'Response ready ({response_size_bytes / 1024 / 1024:.1f}MB)',
+                           metadata={'responseSizeBytes': response_size_bytes})
+            await asyncio.sleep(0)
+
+            # === STAGE 5: COMPLETE ===
+            yield sse_event('complete', stage=5, totalStages=5, progress=95,
+                           stageProgress=0, detail='Transferring data...',
+                           metadata={'responseSizeBytes': response_size_bytes, 'generationTimeMs': int(generation_time * 1000)})
+
+            logger.info(f"✅ [SSE] Research attention analysis complete: {len(generated_tokens)} tokens, {generation_time:.2f}s, {response_size_bytes / 1024 / 1024:.1f}MB")
+
+            # Send final result
+            yield sse_event('result', data=sanitize_for_json(response))
+
+        except Exception as e:
+            logger.error(f"[SSE] Research attention analysis error: {e}")
+            logger.error(traceback.format_exc())
+            yield sse_event('error', detail=str(e), stage=0, totalStages=5, progress=0, stageProgress=0)
+
+    return StreamingResponse(
+        event_generator(),
+        media_type='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache, no-store, must-revalidate',
+            'Connection': 'keep-alive',
+            'X-Accel-Buffering': 'no',  # Disable nginx/proxy buffering
+        }
+    )
+
+
 @app.post("/analyze/study")
 async def analyze_study(request: StudyRequest, authenticated: bool = Depends(verify_api_key)):
     """