"""Deterministic contracts for analysis-informed orchestration boundaries.""" from __future__ import annotations import json from pathlib import Path from types import SimpleNamespace import pytest import torch from obliteratus.informed_pipeline import InformedAbliterationPipeline @pytest.fixture def pipeline(tmp_path): return InformedAbliterationPipeline( model_name="fixture/model", output_dir=str(tmp_path / "output"), on_log=lambda _message: None, ) def test_run_informed_executes_the_documented_stage_order(pipeline, monkeypatch): calls = [] output = pipeline.output_dir for name in ( "_summon", "_probe", "_analyze", "_distill_informed", "_excise_informed", "_verify_and_compensate", ): monkeypatch.setattr(pipeline, name, lambda name=name: calls.append(name)) monkeypatch.setattr( pipeline, "_rebirth_informed", lambda: calls.append("_rebirth_informed") or output, ) ticks = iter((10.0, 12.5)) monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks)) result, report = pipeline.run_informed() assert result == output assert calls == [ "_summon", "_probe", "_analyze", "_distill_informed", "_excise_informed", "_verify_and_compensate", "_rebirth_informed", ] assert report.total_duration == 2.5 def test_analyze_runs_only_enabled_modules_and_records_duration(pipeline, monkeypatch): pipeline._run_cone = False pipeline._run_defense = False calls = [] events = [] monkeypatch.setattr(pipeline, "_analyze_alignment_imprint", lambda: calls.append("alignment")) monkeypatch.setattr(pipeline, "_analyze_cone_geometry", lambda: calls.append("cone")) monkeypatch.setattr(pipeline, "_analyze_cross_layer", lambda: calls.append("cross-layer")) monkeypatch.setattr(pipeline, "_analyze_defense_robustness", lambda: calls.append("defense")) monkeypatch.setattr(pipeline, "_analyze_sparsity", lambda: calls.append("sparsity")) monkeypatch.setattr(pipeline, "_derive_configuration", lambda: calls.append("derive")) monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs))) ticks = iter((20.0, 21.25)) monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks)) pipeline._analyze() assert calls == ["alignment", "cross-layer", "sparsity", "derive"] assert pipeline._report.analysis_duration == 1.25 assert events[0][0][:2] == ("analyze", "running") assert events[-1][0][:2] == ("analyze", "done") assert events[-1][1]["duration"] == 1.25 def test_alignment_imprint_skips_zero_directions(pipeline, monkeypatch): detector_calls = [] class Detector: def detect_imprint(self, directions): detector_calls.append(directions) monkeypatch.setattr( "obliteratus.analysis.alignment_imprint.AlignmentImprintDetector", Detector, ) pipeline._harmful_means = {0: torch.ones(1, 3)} pipeline._harmless_means = {0: torch.ones(1, 3)} pipeline._analyze_alignment_imprint() assert detector_calls == [] assert pipeline._insights.detected_alignment_method == "unknown" def test_alignment_imprint_normalizes_directions_and_copies_result(pipeline, monkeypatch): captured = {} result = SimpleNamespace( predicted_method="dpo", confidence=0.8, dpo_probability=0.8, rlhf_probability=0.1, cai_probability=0.05, sft_probability=0.05, gini_coefficient=0.3, effective_rank=2.0, cross_layer_smoothness=0.7, tail_layer_bias=0.2, ) class Detector: def detect_imprint(self, directions): captured.update(directions) return result monkeypatch.setattr( "obliteratus.analysis.alignment_imprint.AlignmentImprintDetector", Detector, ) pipeline._harmful_means = {1: torch.tensor([[3.0, 4.0]])} pipeline._harmless_means = {1: torch.zeros(1, 2)} pipeline._analyze_alignment_imprint() assert torch.allclose(captured[1], torch.tensor([0.6, 0.8])) assert pipeline._insights.detected_alignment_method == "dpo" assert pipeline._insights.alignment_confidence == 0.8 assert pipeline._insights.alignment_probabilities == { "dpo": 0.8, "rlhf": 0.1, "cai": 0.05, "sft": 0.05, } def test_cone_geometry_aggregates_layers_and_keeps_strongest_categories( pipeline, monkeypatch, ): category = SimpleNamespace( category="safety", direction=torch.tensor([1.0, 0.0]), specificity=0.7, strength=2.0, ) results = iter( ( SimpleNamespace( is_polyhedral=True, cone_dimensionality=3.0, mean_pairwise_cosine=0.2, general_direction=torch.tensor([2.0, 0.0]), category_directions=[category], category_count=1, ), SimpleNamespace( is_polyhedral=True, cone_dimensionality=5.0, mean_pairwise_cosine=0.4, general_direction=torch.tensor([1.0, 0.0]), category_directions=[], category_count=0, ), ) ) class Analyzer: def analyze_layer(self, *_args, **_kwargs): return next(results) monkeypatch.setattr( "obliteratus.analysis.concept_geometry.ConceptConeAnalyzer", Analyzer, ) pipeline._harmful_acts = {idx: [torch.ones(1, 2)] for idx in range(4)} pipeline._harmless_acts = {idx: [torch.zeros(1, 2)] for idx in range(4)} pipeline._analyze_cone_geometry() assert pipeline._insights.cone_is_polyhedral is True assert pipeline._insights.cone_dimensionality == 4.0 assert pipeline._insights.mean_pairwise_cosine == pytest.approx(0.3) assert torch.equal( pipeline._insights.per_category_directions["safety"], torch.tensor([1.0, 0.0]), ) assert pipeline._insights.direction_specificity == {"safety": 0.7} def test_cone_geometry_preserves_defaults_when_no_layers_are_eligible(pipeline, monkeypatch): class Analyzer: def analyze_layer(self, *_args, **_kwargs): raise AssertionError("no layer should be analyzed") monkeypatch.setattr( "obliteratus.analysis.concept_geometry.ConceptConeAnalyzer", Analyzer, ) pipeline._harmful_acts = {0: [torch.ones(1, 2)]} pipeline._harmless_acts = {} pipeline._analyze_cone_geometry() assert pipeline._insights.cone_is_polyhedral is False assert pipeline._insights.cone_dimensionality == 1.0 def test_cross_layer_analysis_selects_the_strongest_cluster_representatives( pipeline, monkeypatch, ): result = SimpleNamespace( clusters=[[0, 1], [2]], cluster_count=2, direction_persistence_score=0.75, mean_adjacent_cosine=0.6, ) class Analyzer: def __init__(self, cluster_threshold): assert cluster_threshold == 0.85 def analyze(self, directions): assert set(directions) == {0, 1, 2} return result monkeypatch.setattr( "obliteratus.analysis.cross_layer.CrossLayerAlignmentAnalyzer", Analyzer, ) pipeline._harmful_means = { 0: torch.tensor([[1.0, 0.0]]), 1: torch.tensor([[3.0, 0.0]]), 2: torch.tensor([[0.0, 2.0]]), } pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(3)} pipeline._analyze_cross_layer() assert pipeline._insights.direction_clusters == [[0, 1], [2]] assert pipeline._insights.cluster_count == 2 assert pipeline._insights.direction_persistence == 0.75 assert pipeline._insights.cluster_representative_layers == [1, 2] def test_cross_layer_analysis_requires_two_nonzero_directions(pipeline, monkeypatch): class Analyzer: def __init__(self, **_kwargs): raise AssertionError("analyzer should not be constructed") monkeypatch.setattr( "obliteratus.analysis.cross_layer.CrossLayerAlignmentAnalyzer", Analyzer, ) pipeline._harmful_means = {0: torch.ones(1, 2)} pipeline._harmless_means = {0: torch.zeros(1, 2)} pipeline._analyze_cross_layer() assert pipeline._insights.cluster_count == 0 def test_defense_analysis_restores_directions_and_copies_profile(pipeline, monkeypatch): original = {9: torch.tensor([9.0])} observed = {} class Evaluator: def __init__(self, candidate): observed["temporary"] = dict(candidate.refusal_directions) def profile_defense(self): return SimpleNamespace( estimated_robustness="high", self_repair_estimate=0.6, entanglement_score=0.4, ) def map_entanglement(self): return SimpleNamespace( most_entangled_layers=[2], least_entangled_layers=[0], ) monkeypatch.setattr( "obliteratus.analysis.defense_robustness.DefenseRobustnessEvaluator", Evaluator, ) pipeline.refusal_directions = original pipeline._harmful_means = {0: torch.tensor([[0.0, 2.0]])} pipeline._harmless_means = {0: torch.zeros(1, 2)} pipeline._analyze_defense_robustness() assert torch.equal(observed["temporary"][0], torch.tensor([0.0, 1.0])) assert pipeline.refusal_directions is original assert pipeline._insights.estimated_robustness == "high" assert pipeline._insights.self_repair_estimate == 0.6 assert pipeline._insights.entangled_layers == [2] assert pipeline._insights.clean_layers == [0] def test_sparsity_analysis_builds_a_plan_from_compatible_ffn_weights( pipeline, monkeypatch, ): projection = SimpleNamespace(weight=SimpleNamespace(data=torch.eye(2))) layer = SimpleNamespace() plan = SimpleNamespace( mean_refusal_sparsity_index=0.65, recommended_sparsity=0.2, most_sparse_layer=0, most_dense_layer=0, ) captured = {} class Surgeon: def __init__(self, auto_sparsity): assert auto_sparsity is True def plan_surgery(self, weights, directions): captured["weights"] = weights captured["directions"] = directions return plan monkeypatch.setattr( "obliteratus.analysis.sparse_surgery.SparseDirectionSurgeon", Surgeon, ) monkeypatch.setattr( "obliteratus.strategies.utils.get_layer_modules", lambda _handle: [layer], ) monkeypatch.setattr( "obliteratus.strategies.utils.get_ffn_module", lambda _layer, _arch: SimpleNamespace(down_proj=projection), ) pipeline.handle = SimpleNamespace(architecture="gpt2") pipeline._harmful_means = {0: torch.tensor([[0.0, 2.0]])} pipeline._harmless_means = {0: torch.zeros(1, 2)} pipeline._analyze_sparsity() assert set(captured["weights"]) == {0} assert torch.equal(captured["directions"][0], torch.tensor([0.0, 1.0])) assert pipeline._insights.mean_refusal_sparsity_index == 0.65 assert pipeline._insights.recommended_sparsity == 0.2 @pytest.mark.parametrize( ("method", "budget"), [("dpo", 0.5), ("rlhf", 0.3), ("cai", 0.2), ("sft", 0.4), ("unknown", 0.35)], ) def test_bayesian_warm_start_sets_alignment_specific_kl_budget( pipeline, method, budget, ): pipeline._insights.detected_alignment_method = method pipeline._configure_bayesian_warm_start() assert pipeline.kl_budget == budget assert pipeline._bayesian_trials == 50 assert pipeline.layer_adaptive_strength is True assert pipeline.float_layer_interpolation is True assert pipeline.use_kl_optimization is True def test_bayesian_warm_start_uses_strongest_cluster_and_entanglement(pipeline): pipeline._harmful_means = { 0: torch.tensor([[1.0, 0.0]]), 1: torch.tensor([[4.0, 0.0]]), 2: torch.tensor([[2.0, 0.0]]), 3: torch.tensor([[1.0, 0.0]]), } pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(4)} pipeline._insights.cluster_representative_layers = [0, 1] pipeline._insights.direction_clusters = [[0, 1], [2, 3]] pipeline._insights.direction_persistence = 0.5 pipeline._insights.entanglement_score = 0.8 pipeline._configure_bayesian_warm_start() warm = pipeline._informed_warm_start assert warm["peak_position"] == pytest.approx(1 / 3) assert warm["spread"] == pytest.approx(1 / 3) assert warm["min_weight"] == 0.1 assert warm["attn_scale"] == 0.7 assert warm["mlp_scale"] == 0.4 def test_excise_informed_routes_sparse_and_dense_paths(pipeline, monkeypatch): calls = [] monkeypatch.setattr(pipeline, "_excise_sparse", lambda: calls.append("sparse")) monkeypatch.setattr( pipeline, "_configure_bayesian_warm_start", lambda: calls.append("warm-start"), ) monkeypatch.setattr(pipeline, "_excise", lambda: calls.append("dense")) pipeline._insights.use_sparse_surgery = True pipeline._excise_informed() pipeline._insights.use_sparse_surgery = False pipeline._excise_informed() assert calls == ["sparse", "warm-start", "dense"] def test_verify_compensation_stops_when_no_residual_layers(pipeline, monkeypatch): calls = [] def verify(): calls.append("verify") pipeline._quality_metrics = {"refusal_rate": 0.9, "kl_divergence": 0.1} monkeypatch.setattr(pipeline, "_verify", verify) monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe")) def distill(): calls.append("distill") pipeline._strong_layers = [] monkeypatch.setattr(pipeline, "_distill_inner", distill) monkeypatch.setattr(pipeline, "_excise_informed", lambda: calls.append("excise")) pipeline._verify_and_compensate() assert calls == ["verify", "probe", "distill"] assert pipeline._report.ouroboros_passes == 1 assert pipeline._report.final_refusal_rate == 0.9 def test_verify_compensation_stops_at_kl_ceiling(pipeline, monkeypatch): outcomes = iter( ( {"refusal_rate": 0.9, "kl_divergence": 0.1}, {"refusal_rate": 0.8, "kl_divergence": 0.8}, ) ) calls = [] def verify(): calls.append("verify") pipeline._quality_metrics = next(outcomes) monkeypatch.setattr(pipeline, "_verify", verify) monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe")) monkeypatch.setattr(pipeline, "_distill_inner", lambda: setattr(pipeline, "_strong_layers", [1])) monkeypatch.setattr(pipeline, "_excise_informed", lambda: calls.append("excise")) pipeline.kl_budget = 0.3 pipeline._verify_and_compensate() assert calls == ["verify", "probe", "excise", "verify"] assert pipeline._report.ouroboros_passes == 1 assert pipeline._report.final_refusal_rate == 0.8 def test_rebirth_writes_model_tokenizer_and_research_metadata(pipeline, monkeypatch): saved = [] class Artifact: def __init__(self, name): self.name = name def save_pretrained(self, path): saved.append((self.name, Path(path))) pipeline.handle = SimpleNamespace(model=Artifact("model"), tokenizer=Artifact("tokenizer")) pipeline._strong_layers = [1, 3] pipeline._quality_metrics = {"refusal_rate": 0.1} pipeline._insights.detected_alignment_method = "dpo" pipeline._insights.recommended_layers = [1, 3] pipeline._report.analysis_duration = 1.2 pipeline._report.total_duration = 3.4 pipeline._report.ouroboros_passes = 1 pipeline._report.final_refusal_rate = 0.1 events = [] monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs))) ticks = iter((5.0, 5.5)) monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks)) result = pipeline._rebirth_informed() assert result == pipeline.output_dir assert saved == [("model", result), ("tokenizer", result)] metadata = json.loads((result / "abliteration_metadata.json").read_text()) assert metadata["analysis_insights"]["detected_alignment_method"] == "dpo" assert metadata["derived_config"]["layers_used"] == [1, 3] assert metadata["pipeline_stats"]["ouroboros_passes"] == 1 assert metadata["quality_metrics"] == {"refusal_rate": 0.1} assert events[0][0][:2] == ("rebirth", "running") assert events[-1][0][:2] == ("rebirth", "done") def test_distill_single_direction_honors_recommended_and_skipped_layers( pipeline, monkeypatch, ): pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000) pipeline.n_directions = 1 pipeline._harmful_means = { 0: torch.tensor([[3.0, 4.0]]), 1: torch.tensor([[0.0, 2.0]]), } pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(2)} pipeline._insights.recommended_layers = [0, 1, 99] pipeline._insights.skip_layers = [1] events = [] monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs))) ticks = iter((1.0, 1.5)) monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks)) pipeline._distill_informed() assert torch.allclose(pipeline.refusal_directions[0], torch.tensor([0.6, 0.8])) assert pipeline.refusal_subspaces[0].shape == (1, 2) assert pipeline._strong_layers == [0] assert events[0][0][:2] == ("distill", "running") assert events[-1][1]["strong_layers"] == [0] def test_distill_svd_sanitizes_nonfinite_input_and_enriches_category_directions( pipeline, monkeypatch, ): pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000) pipeline.n_directions = 3 pipeline.use_whitened_svd = False pipeline._harmful_means = {0: torch.tensor([[2.0, 0.0, 0.0]])} pipeline._harmless_means = {0: torch.zeros(1, 3)} pipeline._harmful_acts = { 0: [ torch.tensor([[2.0, 0.0, 0.0]]), torch.tensor([[0.0, float("nan"), 0.0]]), torch.tensor([[0.0, 0.0, 1.0]]), ] } pipeline._harmless_acts = {0: [torch.zeros(1, 3) for _ in range(3)]} pipeline._insights.cone_is_polyhedral = True pipeline._insights.per_category_directions = { "one": torch.tensor([0.0, 1.0, 0.0]), "two": torch.tensor([0.0, 0.0, 1.0]), } monkeypatch.setattr(pipeline, "_select_layers_knee", lambda ranked: [ranked[0][0]]) monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None) pipeline._distill_informed() assert pipeline._strong_layers == [0] assert pipeline.refusal_subspaces[0].shape[0] >= 2 assert torch.isfinite(pipeline.refusal_subspaces[0]).all() def test_distill_uses_whitened_extractor_for_multi_direction_models( pipeline, monkeypatch, ): result = SimpleNamespace( directions=torch.tensor([[1.0, 0.0], [0.0, 1.0]]), singular_values=torch.tensor([3.0, 1.0]), ) class Extractor: def extract(self, harmful, harmless, *, n_directions, layer_idx): assert harmful is pipeline._harmful_acts[0] assert harmless is pipeline._harmless_acts[0] assert (n_directions, layer_idx) == (2, 0) return result monkeypatch.setattr( "obliteratus.analysis.whitened_svd.WhitenedSVDExtractor", Extractor, ) pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000) pipeline.n_directions = 2 pipeline.use_whitened_svd = True pipeline._harmful_means = {0: torch.tensor([[1.0, 0.0]])} pipeline._harmless_means = {0: torch.zeros(1, 2)} pipeline._harmful_acts = {0: [torch.ones(1, 2)]} pipeline._harmless_acts = {0: [torch.zeros(1, 2)]} monkeypatch.setattr(pipeline, "_select_layers_knee", lambda _ranked: [0]) monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None) pipeline._distill_informed() assert torch.equal(pipeline.refusal_subspaces[0], result.directions) assert torch.equal(pipeline.refusal_directions[0], result.directions[0]) def test_distill_leace_falls_back_per_layer_after_extractor_failure( pipeline, monkeypatch, ): leace_result = SimpleNamespace( direction=torch.tensor([0.0, 1.0]), generalized_eigenvalue=4.0, erasure_loss=0.2, ) class Extractor: def extract(self, _harmful, _harmless, *, layer_idx): if layer_idx == 0: return leace_result raise RuntimeError("singular fixture") monkeypatch.setattr("obliteratus.analysis.leace.LEACEExtractor", Extractor) pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000) pipeline.direction_method = "leace" pipeline.n_directions = 1 pipeline._harmful_means = { 0: torch.tensor([[0.0, 1.0]]), 1: torch.tensor([[1.0, 0.0]]), } pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(2)} pipeline._harmful_acts = {idx: [torch.ones(1, 2)] for idx in range(2)} pipeline._harmless_acts = {idx: [torch.zeros(1, 2)] for idx in range(2)} monkeypatch.setattr(pipeline, "_select_layers_knee", lambda _ranked: [0, 1]) monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None) pipeline._distill_informed() assert torch.equal(pipeline.refusal_directions[0], leace_result.direction) assert torch.equal(pipeline.refusal_directions[1], torch.tensor([1.0, 0.0])) def test_sparse_excision_projects_attention_and_ffn_with_iterative_reprobe( pipeline, monkeypatch, ): class Layer(torch.nn.Module): def __init__(self): super().__init__() self.anchor = torch.nn.Parameter(torch.ones(1)) layer = Layer() attention = SimpleNamespace(o_proj=torch.nn.Linear(2, 2, bias=False)) ffn = SimpleNamespace(down_proj=torch.nn.Linear(2, 2, bias=False)) calls = [] class Surgeon: def __init__(self, *, sparsity, auto_sparsity): assert sparsity == 0.25 assert auto_sparsity is True def apply_sparse_projection(self, weight, direction): calls.append((weight.clone(), direction.clone())) return weight * 0.5 monkeypatch.setattr( "obliteratus.analysis.sparse_surgery.SparseDirectionSurgeon", Surgeon, ) monkeypatch.setattr( "obliteratus.strategies.utils.get_layer_modules", lambda _handle: [layer], ) monkeypatch.setattr( "obliteratus.strategies.utils.get_attention_module", lambda _layer, _arch: attention, ) monkeypatch.setattr( "obliteratus.strategies.utils.get_ffn_module", lambda _layer, _arch: ffn, ) pipeline.handle = SimpleNamespace(architecture="gpt2") pipeline._insights.recommended_sparsity = 0.25 pipeline._strong_layers = [0] pipeline.refusal_subspaces = {0: torch.eye(2)} pipeline.refinement_passes = 2 pipeline.true_iterative_refinement = True monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe")) monkeypatch.setattr(pipeline, "_distill_inner", lambda: calls.append("distill")) events = [] monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs))) pipeline._excise_sparse() projection_calls = [call for call in calls if isinstance(call, tuple)] assert len(projection_calls) == 8 assert calls.count("probe") == 1 assert calls.count("distill") == 1 assert events[-1][1]["modified_count"] == 8 def test_verify_compensation_stops_when_kl_rises_sharply(pipeline, monkeypatch): outcomes = iter( ( {"refusal_rate": 0.9, "kl_divergence": 0.1}, {"refusal_rate": 0.8, "kl_divergence": 0.12}, {"refusal_rate": 0.7, "kl_divergence": 0.2}, ) ) def verify(): pipeline._quality_metrics = next(outcomes) monkeypatch.setattr(pipeline, "_verify", verify) monkeypatch.setattr(pipeline, "_probe", lambda: None) monkeypatch.setattr(pipeline, "_distill_inner", lambda: setattr(pipeline, "_strong_layers", [0])) monkeypatch.setattr(pipeline, "_excise_informed", lambda: None) pipeline.kl_budget = 1.0 pipeline._verify_and_compensate() assert pipeline._report.ouroboros_passes == 2 assert pipeline._report.final_refusal_rate == 0.7