Claude's Corner: Envariant - The Control Plane Foundation Models Never Had

# How to Build an AI Interpretability SDK Like Envariant

A step-by-step guide for building an interpretability and model-control SDK for foundation models, using Claude Code.

---

## Step 1: Project Setup and Architecture

Create a Python package with PyTorch as the core dependency. The SDK needs three main modules: `probe` (behavioral detection), `steer` (activation intervention), and `explain` (principle extraction).

```
envariant-sdk/
  envariant/
    __init__.py
    probe.py          # Probing classifiers and detectors
    steer.py          # Activation patching and steering
    explain.py        # SAE decomposition and principle extraction
    synthesize.py     # Edge case generation
    hooks.py          # PyTorch forward hook management
    utils.py
  examples/
  tests/
  pyproject.toml
```

**DB schema for storing behavioral probes:**

```sql
CREATE TABLE probes (
  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  model_family TEXT NOT NULL,          -- e.g. 'llama3', 'mistral'
  model_version TEXT NOT NULL,
  layer_index INTEGER NOT NULL,
  behavior_name TEXT NOT NULL,         -- e.g. 'hallucination', 'safety_violation'
  probe_weights BYTEA NOT NULL,        -- serialized sklearn LogisticRegression or similar
  accuracy FLOAT,
  f1_score FLOAT,
  threshold FLOAT NOT NULL DEFAULT 0.5,
  training_samples INTEGER,
  created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE behavioral_events (
  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  probe_id UUID REFERENCES probes(id),
  input_hash TEXT,
  detected_at_layer INTEGER,
  confidence FLOAT,
  intervened BOOLEAN DEFAULT FALSE,
  outcome TEXT,
  created_at TIMESTAMPTZ DEFAULT NOW()
);
```

---

## Step 2: Hook Management System

The core primitive is attaching to transformer layers without modifying the model. Use PyTorch's `register_forward_hook` API.

```python
# hooks.py
import torch
from typing import Dict, List, Callable, Optional
from contextlib import contextmanager

class HookManager:
    def __init__(self, model):
        self.model = model
        self._hooks = []
        self._activations: Dict[str, torch.Tensor] = {}
    
    def capture_layer(self, layer_name: str, layer_module):
        """Attach a hook that captures activations at a named layer."""
        def hook_fn(module, input, output):
            # For transformer blocks, capture the residual stream
            if isinstance(output, tuple):
                self._activations[layer_name] = output[0].detach()
            else:
                self._activations[layer_name] = output.detach()
        
        handle = layer_module.register_forward_hook(hook_fn)
        self._hooks.append(handle)
        return handle
    
    @contextmanager
    def capture_context(self, layers: Dict[str, any]):
        """Context manager for clean capture and teardown."""
        handles = []
        try:
            for name, module in layers.items():
                handles.append(self.capture_layer(name, module))
            yield self._activations
        finally:
            for h in handles:
                h.remove()
            self._activations.clear()
    
    def patch_activation(self, layer_module, patch_tensor: torch.Tensor):
        """Replace activations at a layer with a patched version (for causal tracing)."""
        def patch_hook(module, input, output):
            if isinstance(output, tuple):
                return (patch_tensor,) + output[1:]
            return patch_tensor
        return layer_module.register_forward_hook(patch_hook)
```

The key design decision here: use context managers to guarantee hook cleanup. Leaked hooks accumulate and silently corrupt inference results.

---

## Step 3: Probing Classifiers

Train lightweight classifiers on layer activations to predict behavioral properties. Use mean-pooled residual stream representations as features.

```python
# probe.py
import torch
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import pickle

class BehavioralProbe:
    def __init__(self, behavior_name: str, layer_index: int):
        self.behavior_name = behavior_name
        self.layer_index = layer_index
        self.classifier = LogisticRegression(C=1.0, max_iter=1000)
        self.scaler = StandardScaler()
        self._fitted = False
    
    def extract_features(self, activations: torch.Tensor) -> np.ndarray:
        """Mean-pool across sequence length to get fixed-size representation."""
        # activations: (batch, seq_len, hidden_dim)
        pooled = activations.mean(dim=1)  # (batch, hidden_dim)
        return pooled.cpu().float().numpy()
    
    def fit(self, activation_samples: list, labels: list):
        """Train probe on (activations, binary_label) pairs."""
        X = np.vstack([self.extract_features(a) for a in activation_samples])
        y = np.array(labels)
        X_scaled = self.scaler.fit_transform(X)
        self.classifier.fit(X_scaled, y)
        self._fitted = True
    
    def predict(self, activations: torch.Tensor) -> tuple[float, bool]:
        """Return (confidence, detected) for a single forward pass."""
        if not self._fitted:
            raise RuntimeError("Probe not fitted")
        X = self.extract_features(activations.unsqueeze(0))
        X_scaled = self.scaler.transform(X)
        proba = self.classifier.predict_proba(X_scaled)[0][1]
        return float(proba), proba > 0.5
    
    def serialize(self) -> bytes:
        return pickle.dumps({'classifier': self.classifier, 'scaler': self.scaler})
    
    @classmethod
    def deserialize(cls, behavior_name: str, layer_index: int, data: bytes) -> 'BehavioralProbe':
        probe = cls(behavior_name, layer_index)
        obj = pickle.loads(data)
        probe.classifier = obj['classifier']
        probe.scaler = obj['scaler']
        probe._fitted = True
        return probe
```

For performance: logistic regression inference is microseconds. The bottleneck is extracting activations during the model's forward pass, not the probe itself.

---

## Step 4: Causal Tracing (Activation Patching)

Implement the activation patching algorithm to identify which model components cause a behavior.

```python
# steer.py
import torch
from typing import Dict, Callable

class CausalTracer:
    def __init__(self, model, hook_manager):
        self.model = model
        self.hooks = hook_manager
    
    def run_with_cache(self, inputs, layers_to_capture: list) -> tuple:
        """Run model forward pass and cache specified layer activations."""
        layer_map = {f"layer_{i}": self.model.layers[i] for i in layers_to_capture}
        with self.hooks.capture_context(layer_map) as activations:
            with torch.no_grad():
                output = self.model(**inputs)
        return output, dict(activations)
    
    def patch_and_measure(self, clean_inputs, corrupted_inputs, 
                           patch_layer_idx: int, metric_fn: Callable) -> float:
        """
        Patch corrupted run with clean activations at one layer.
        Returns change in metric (higher = this layer matters more).
        """
        # Get clean activations
        _, clean_cache = self.run_with_cache(clean_inputs, [patch_layer_idx])
        clean_acts = clean_cache[f"layer_{patch_layer_idx}"]
        
        # Run corrupted model with patch at this layer
        patch_handle = self.hooks.patch_activation(
            self.model.layers[patch_layer_idx], 
            clean_acts
        )
        try:
            with torch.no_grad():
                patched_output = self.model(**corrupted_inputs)
        finally:
            patch_handle.remove()
        
        return metric_fn(patched_output)
    
    def full_trace(self, clean_inputs, corrupted_inputs, metric_fn: Callable) -> Dict[int, float]:
        """Run patching experiment across all layers. Returns layer -> effect size."""
        n_layers = len(self.model.layers)
        effects = {}
        for i in range(n_layers):
            effects[i] = self.patch_and_measure(clean_inputs, corrupted_inputs, i, metric_fn)
        return effects

class BehavioralSteerer:
    def __init__(self, model, hook_manager):
        self.model = model
        self.hooks = hook_manager
    
    def add_steering_vector(self, layer_idx: int, direction: torch.Tensor, 
                             coefficient: float = 1.0):
        """
        Add a steering vector to a layer's output at inference time.
        direction: (hidden_dim,) tensor pointing toward desired behavior.
        """
        def steer_hook(module, input, output):
            if isinstance(output, tuple):
                hidden = output[0]
                hidden = hidden + coefficient * direction.to(hidden.device)
                return (hidden,) + output[1:]
            return output + coefficient * direction.to(output.device)
        
        return self.model.layers[layer_idx].register_forward_hook(steer_hook)
```

---

## Step 5: Sparse Autoencoder for Feature Decomposition

SAEs decompose the dense activation space into interpretable sparse features. Train offline, run at inference.

```python
# explain.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseAutoencoder(nn.Module):
    """
    Standard ReLU sparse autoencoder for mechanistic interpretability.
    Input: residual stream activations (hidden_dim,)
    Output: sparse feature activations (n_features,) and reconstruction
    """
    def __init__(self, hidden_dim: int, n_features: int, l1_coeff: float = 1e-3):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_features = n_features
        self.l1_coeff = l1_coeff
        
        self.encoder = nn.Linear(hidden_dim, n_features, bias=True)
        self.decoder = nn.Linear(n_features, hidden_dim, bias=True)
        
        # Normalize decoder columns to unit norm (key for interpretability)
        with torch.no_grad():
            self.decoder.weight.data = F.normalize(self.decoder.weight.data, dim=0)
    
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Map activations to sparse feature space."""
        return F.relu(self.encoder(x))
    
    def decode(self, features: torch.Tensor) -> torch.Tensor:
        return self.decoder(features)
    
    def forward(self, x: torch.Tensor):
        features = self.encode(x)
        reconstruction = self.decode(features)
        l1_loss = features.abs().sum(dim=-1).mean()
        recon_loss = F.mse_loss(reconstruction, x)
        return reconstruction, features, recon_loss + self.l1_coeff * l1_loss
    
    def get_top_features(self, x: torch.Tensor, k: int = 10) -> list:
        """Return indices of top-k activated features for an input."""
        with torch.no_grad():
            features = self.encode(x)
        top_k = features.topk(k, dim=-1)
        return list(zip(top_k.indices.tolist(), top_k.values.tolist()))

def extract_human_readable_principles(sae: SparseAutoencoder, 
                                       feature_idx: int,
                                       example_activations: list,
                                       tokenizer,
                                       top_n: int = 20) -> str:
    """
    Find examples that maximally activate a feature.
    Use an LLM to generate a human-readable description.
    """
    scores = []
    for text, acts in example_activations:
        feat_val = sae.encode(acts)[feature_idx].item()
        scores.append((feat_val, text))
    scores.sort(reverse=True)
    top_examples = [text for _, text in scores[:top_n]]
    # Pass top_examples to an LLM with prompt: "What concept do these share?"
    return top_examples  # caller passes to LLM for description
```

---

## Step 6: Edge Case Synthesis

Generate targeted inputs that probe model behavior near behavioral boundaries.

```python
# synthesize.py
import torch
from typing import Optional

class EdgeCaseSynthesizer:
    def __init__(self, model, tokenizer, probe, hook_manager):
        self.model = model
        self.tokenizer = tokenizer
        self.probe = probe
        self.hooks = hook_manager
    
    def find_boundary_inputs(self, seed_texts: list, 
                              target_confidence: float = 0.5,
                              n_candidates: int = 100) -> list:
        """
        Given seed texts near a behavioral boundary, generate variants
        that cluster around probe confidence = target_confidence.
        Strategy: gradient-guided token perturbation.
        """
        boundary_cases = []
        
        for seed in seed_texts:
            inputs = self.tokenizer(seed, return_tensors='pt')
            embeddings = self.model.get_input_embeddings()(inputs['input_ids'])
            embeddings.requires_grad_(True)
            
            # Forward pass to get probe confidence
            # Gradient tells us which embedding directions increase/decrease confidence
            layer_module = self.model.layers[self.probe.layer_index]
            captured = {}
            
            def capture_fn(module, inp, out):
                captured['acts'] = out[0] if isinstance(out, tuple) else out
            
            handle = layer_module.register_forward_hook(capture_fn)
            output = self.model(inputs_embeds=embeddings, 
                               attention_mask=inputs['attention_mask'])
            handle.remove()
            
            confidence, _ = self.probe.predict(captured['acts'].squeeze(0))
            
            # For actual synthesis: perturb in gradient direction, decode, filter
            # This simplified version returns the seed with its confidence
            boundary_cases.append({
                'text': seed,
                'confidence': confidence,
                'distance_to_boundary': abs(confidence - target_confidence)
            })
        
        return sorted(boundary_cases, key=lambda x: x['distance_to_boundary'])
    
    def generate_adversarial(self, base_text: str, target_behavior: bool, 
                              max_edits: int = 5) -> Optional[str]:
        """
        Minimal edit to base_text that flips probe prediction.
        Returns None if no flip found within max_edits.
        """
        # Implementation: beam search over token replacements
        # guided by probe confidence gradient
        pass  # Full implementation requires model-specific tokenizer integration
```

---

## Step 7: API Design and Deployment

Package the SDK for enterprise use with a clean Python API and optional dashboard.

```python
# envariant/__init__.py - Main SDK interface
from .probe import BehavioralProbe
from .steer import CausalTracer, BehavioralSteerer
from .explain import SparseAutoencoder
from .hooks import HookManager

class Envariant:
    def __init__(self, model, model_family: str = "auto"):
        self.model = model
        self.model_family = model_family
        self.hooks = HookManager(model)
        self.probes = {}
        self.saes = {}
        self.steerers = []
    
    def add_probe(self, behavior: str, layer: int) -> BehavioralProbe:
        probe = BehavioralProbe(behavior, layer)
        self.probes[behavior] = probe
        return probe
    
    def detect(self, inputs, behavior: str) -> dict:
        """Run detection for a registered behavior during inference."""
        probe = self.probes[behavior]
        layer = self.model.layers[probe.layer_index]
        captured = {}
        
        def fn(module, inp, out):
            captured['acts'] = out[0] if isinstance(out, tuple) else out
        
        handle = layer.register_forward_hook(fn)
        with torch.no_grad():
            output = self.model(**inputs)
        handle.remove()
        
        conf, detected = probe.predict(captured['acts'].squeeze(0))
        return {'behavior': behavior, 'confidence': conf, 'detected': detected, 'output': output}
    
    def steer(self, layer: int, direction: torch.Tensor, coefficient: float = 1.0):
        """Add a persistent steering hook for subsequent inference calls."""
        steerer = BehavioralSteerer(self.model, self.hooks)
        handle = steerer.add_steering_vector(layer, direction, coefficient)
        self.steerers.append(handle)
        return handle
    
    def clear_steering(self):
        for h in self.steerers:
            h.remove()
        self.steerers.clear()
```

**Deployment architecture:**
- Publish as a pip package (`pip install envariant`)
- Optional SaaS dashboard: Next.js frontend, FastAPI backend, PostgreSQL for probe storage
- Enterprise self-hosted option: Docker Compose with Postgres + Redis for caching probe results
- The probe weights are small (KB-range) and can be stored in Supabase/S3, fetched at SDK init
- For real-time dashboard: stream behavioral events via WebSocket from the inference server

**Key deployment decision:** the SDK runs in-process with the model, not as a sidecar. This keeps latency overhead minimal (sub-millisecond for probe inference) but means the customer needs to install the SDK in their inference environment. Cloud-native packaging (NVIDIA Triton custom backend, or a vLLM plugin) is the path to enterprise adoption.