From c83c8972801ed403e376a479f9336ef00add3f51 Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:37:16 +0100 Subject: [PATCH 1/6] feat: add fail-closed route study control plane --- source/route_policy_explorer.py | 1479 ++++++++++++++ source/route_policy_lab.py | 1328 +++++++++++++ source/route_policy_ledger.py | 2511 ++++++++++++++++++++++++ source/route_policy_protocol.py | 1541 +++++++++++++++ source/route_policy_protocol_cli.py | 163 ++ source/route_policy_shadow_cli.py | 352 ++++ source/route_policy_shadow_registry.py | 2268 +++++++++++++++++++++ source/route_policy_study_cli.py | 174 ++ test_route_policy_explorer.py | 646 ++++++ test_route_policy_lab.py | 1008 ++++++++++ test_route_policy_ledger.py | 1201 ++++++++++++ test_route_policy_protocol.py | 435 ++++ test_route_policy_protocol_cli.py | 105 + test_route_policy_shadow_cli.py | 310 +++ test_route_policy_shadow_registry.py | 853 ++++++++ test_route_policy_study_cli.py | 89 + 16 files changed, 14463 insertions(+) create mode 100644 source/route_policy_explorer.py create mode 100644 source/route_policy_lab.py create mode 100644 source/route_policy_ledger.py create mode 100644 source/route_policy_protocol.py create mode 100644 source/route_policy_protocol_cli.py create mode 100644 source/route_policy_shadow_cli.py create mode 100644 source/route_policy_shadow_registry.py create mode 100644 source/route_policy_study_cli.py create mode 100644 test_route_policy_explorer.py create mode 100644 test_route_policy_lab.py create mode 100644 test_route_policy_ledger.py create mode 100644 test_route_policy_protocol.py create mode 100644 test_route_policy_protocol_cli.py create mode 100644 test_route_policy_shadow_cli.py create mode 100644 test_route_policy_shadow_registry.py create mode 100644 test_route_policy_study_cli.py diff --git a/source/route_policy_explorer.py b/source/route_policy_explorer.py new file mode 100644 index 000000000..9c9b2b8e4 --- /dev/null +++ b/source/route_policy_explorer.py @@ -0,0 +1,1479 @@ +"""Pure planning primitives for a bounded-exposure adjacent-route rehearsal. + +This module deliberately stops before runtime integration. It can describe a +versioned, prompt-free exploration charter and deterministically assign an +action from caller-supplied nonce material, but it performs no I/O, writes no +ledger rows, starts no inference, estimates no policy value, and authorizes no +promotion. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +from decimal import Decimal, localcontext +from typing import Any, Dict, List, Mapping, Sequence, Tuple + +try: + from .route_policy_ledger import OUTCOME_CONTRACT_SCHEMA_VERSION, SUPPORT_SCHEMA_VERSION +except ImportError: # pragma: no cover - direct source/ execution compatibility + from route_policy_ledger import OUTCOME_CONTRACT_SCHEMA_VERSION, SUPPORT_SCHEMA_VERSION + + +STUDY_PLAN_SCHEMA_VERSION = "route-exploration-plan-v1" +STUDY_ASSIGNMENT_SCHEMA_VERSION = "route-exploration-rehearsal-assignment-v1" +REHEARSAL_RECEIPT_SCHEMA_VERSION = "route-exploration-rehearsal-receipt-v1" +REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION = "route-exploration-rehearsal-support-proposal-v1" +STUDY_ID = "auto-route-adjacent-explorer-v1" +STUDY_VERSION = "1.0.0" +STUDY_LABEL = "Bounded-Exposure Adjacent-Route Rehearsal v1" +ASSIGNMENT_ALGORITHM = "sha256-cdf-v1" + +AUTO_AGENT_MODE_ORDER: Tuple[str, ...] = ( + "off", + "collective", + "loop", + "collective_loop", +) +DEFAULT_EXPLORATION_RATE = 0.10 +MAX_EXPLORATION_RATE = 0.20 +MIN_POSITIVE_EXPLORATION_PROBABILITY = 0.05 +MAX_ADJACENT_NEIGHBORS = 2 +DEFAULT_PLANNED_ROUTES = 100 +DEFAULT_SCENARIO_CONFIDENCE = 0.95 +DEFAULT_ASSUMED_FEEDBACK_RATE = 0.30 +DEFAULT_TARGET_OBSERVED_LABELS = 20 +MAX_PLANNED_ROUTES = 100_000 +MAX_TARGET_OBSERVED_LABELS = 1_000 +MAX_EXACT_BINOMIAL_ROUTE_FORECAST = 10_000_000_000 + +_DESIGN_HASH_DOMAIN = b"supermix.route-explorer.design.v1\x00" +_NONCE_HASH_DOMAIN = b"supermix.route-explorer.nonce.v1\x00" +_DRAW_HASH_DOMAIN = b"supermix.route-explorer.draw.v1\x00" +_SOURCE_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,159}$") + +ACTIVATION_BLOCKERS: Tuple[str, ...] = ( + "target_policy_class_not_precommitted", + "outcome_observation_process_not_validated", + "population_definition_not_precommitted", + "session_carryover_not_addressed", + "interference_not_addressed", + "stopping_rules_not_precommitted", + "preassignment_seed_commitment_not_sealed", + "external_ope_not_validated", +) + +_SOURCE_CONTRACT_KEYS = { + "policy_id", + "policy_version", + "feature_schema_version", + "support_schema_version", + "candidate_set_hash", + "distribution_hash", + "outcome_contract_schema_version", +} + +_CANDIDATE_KEYS = { + "action", + "estimated_cost_units", + "estimated_model_calls", + "planned_loop_steps", + "latency_tier", + "selected", +} +_REQUIRED_CANDIDATE_KEYS = { + "action", + "estimated_cost_units", + "estimated_model_calls", + "planned_loop_steps", + "latency_tier", +} +_EXCLUSION_KEYS = {"action", "reasons"} +_ALLOWED_EXCLUSION_REASONS = { + "action_mode_unsupported", + "capability_or_policy_filter", + "session_budget_post_filter", +} +_LATENCY_TIERS = {"low", "moderate", "high", "frontier", "unknown"} +_PLAN_KEYS = {"schema_version", "study", "charter", "design_hash"} +_STUDY_KEYS = {"study_id", "study_version", "label", "design_kind"} +_CHARTER_KEYS = { + "source_contract", + "prompt_free_contract", + "enrollment", + "probability_design", + "candidates", + "exclusions", + "traffic_scenario", + "resource_forecast", + "causal_boundaries", +} +_ENROLLMENT_KEYS = { + "eligible", + "reason", + "baseline_action", + "adjacent_feasible_actions", + "maximum_adjacent_neighbors", +} +_PROBABILITY_DESIGN_KEYS = { + "decision_type", + "probability_stage", + "requested_exploration_rate", + "applied_exploration_rate", + "minimum_positive_exploration_probability", + "eligible_actions", + "action_probabilities", + "assignment_algorithm", + "assignment_unit", + "assignment_performed", +} +_STUDY_CANDIDATE_KEYS = { + "action", + "baseline", + "adjacent_distance", + "estimated_cost_units", + "estimated_model_calls", + "planned_loop_steps", + "latency_tier", +} +_STUDY_EXCLUSION_KEYS = { + "action", + "reasons", + "post_filter_feasible", +} + + +def _canonical_json(value: Any) -> str: + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError("study charter must be canonical JSON without non-finite values") from exc + + +def _domain_hash(domain: bytes, value: Any) -> str: + return hashlib.sha256(domain + _canonical_json(value).encode("utf-8")).hexdigest() + + +def _strict_keys(value: Mapping[str, Any], allowed: set[str], name: str) -> None: + unknown = set(value) - allowed + if unknown: + raise ValueError(f"{name} contains unsupported fields: {', '.join(sorted(map(str, unknown)))}") + + +def _source_identifier(value: Any, name: str) -> str: + cooked = str(value or "").strip() + if not _SOURCE_IDENTIFIER_RE.fullmatch(cooked): + raise ValueError(f"source_contract {name} must be a versioned identifier") + return cooked + + +def _sha256_hex(value: Any, name: str) -> str: + cooked = str(value or "").strip() + if len(cooked) != 64 or any(character not in "0123456789abcdef" for character in cooked): + raise ValueError(f"source_contract {name} must be a lowercase SHA-256 hex digest") + return cooked + + +def _normalize_source_contract(value: Any) -> Dict[str, str]: + if not isinstance(value, Mapping): + raise ValueError("source_contract must be a JSON object") + _strict_keys(value, _SOURCE_CONTRACT_KEYS, "source_contract") + if set(value) != _SOURCE_CONTRACT_KEYS: + raise ValueError("source_contract must contain every required provenance field") + support_schema = _source_identifier( + value.get("support_schema_version"), "support_schema_version" + ) + outcome_schema = _source_identifier( + value.get("outcome_contract_schema_version"), + "outcome_contract_schema_version", + ) + if support_schema != SUPPORT_SCHEMA_VERSION: + raise ValueError( + f"source_contract support_schema_version must be {SUPPORT_SCHEMA_VERSION}" + ) + if outcome_schema != OUTCOME_CONTRACT_SCHEMA_VERSION: + raise ValueError( + "source_contract outcome_contract_schema_version must be " + f"{OUTCOME_CONTRACT_SCHEMA_VERSION}" + ) + return { + "policy_id": _source_identifier(value.get("policy_id"), "policy_id"), + "policy_version": _source_identifier( + value.get("policy_version"), "policy_version" + ), + "feature_schema_version": _source_identifier( + value.get("feature_schema_version"), "feature_schema_version" + ), + "support_schema_version": support_schema, + "candidate_set_hash": _sha256_hex( + value.get("candidate_set_hash"), "candidate_set_hash" + ), + "distribution_hash": _sha256_hex( + value.get("distribution_hash"), "distribution_hash" + ), + "outcome_contract_schema_version": outcome_schema, + } + + +def _finite_nonnegative(value: Any, name: str) -> float: + if isinstance(value, bool): + raise ValueError(f"{name} must be a finite non-negative number") + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"{name} must be a finite non-negative number") from exc + if not math.isfinite(cooked) or cooked < 0.0: + raise ValueError(f"{name} must be a finite non-negative number") + return round(cooked, 6) + + +def _nonnegative_int(value: Any, name: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be an integer of at least {minimum}") + if value < minimum: + raise ValueError(f"{name} must be an integer of at least {minimum}") + return int(value) + + +def _validate_rate(value: Any) -> float: + if isinstance(value, bool): + raise ValueError("exploration_rate must be a finite number in (0, 0.20]") + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("exploration_rate must be a finite number in (0, 0.20]") from exc + if not math.isfinite(cooked) or not 0.0 < cooked <= MAX_EXPLORATION_RATE: + raise ValueError("exploration_rate must be a finite number in (0, 0.20]") + return round(cooked, 12) + + +def _validate_planned_routes(value: Any) -> int: + routes = _nonnegative_int(value, "planned_routes", minimum=1) + if routes > MAX_PLANNED_ROUTES: + raise ValueError(f"planned_routes must not exceed {MAX_PLANNED_ROUTES}") + return routes + + +def _validate_confidence(value: Any) -> float: + if isinstance(value, bool): + raise ValueError("scenario_confidence must be a finite number in (0.5, 1)") + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("scenario_confidence must be a finite number in (0.5, 1)") from exc + if not math.isfinite(cooked) or not 0.5 < cooked < 1.0: + raise ValueError("scenario_confidence must be a finite number in (0.5, 1)") + return round(cooked, 12) + + +def _validate_feedback_rate(value: Any) -> float: + if isinstance(value, bool): + raise ValueError("assumed_feedback_rate must be a finite number in (0, 1]") + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("assumed_feedback_rate must be a finite number in (0, 1]") from exc + if not math.isfinite(cooked) or not 0.0 < cooked <= 1.0: + raise ValueError("assumed_feedback_rate must be a finite number in (0, 1]") + return round(cooked, 12) + + +def _validate_target_labels(value: Any) -> int: + target = _nonnegative_int(value, "target_observed_labels", minimum=1) + if target > MAX_TARGET_OBSERVED_LABELS: + raise ValueError( + f"target_observed_labels must not exceed {MAX_TARGET_OBSERVED_LABELS}" + ) + return target + + +def _normalize_candidates(value: Any, baseline_mode: str) -> List[Dict[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError("post_filter_candidates must be a non-empty sequence") + normalized: List[Dict[str, Any]] = [] + seen: set[str] = set() + for raw in value: + if not isinstance(raw, Mapping): + raise ValueError("post_filter_candidates must contain JSON objects") + _strict_keys(raw, _CANDIDATE_KEYS, "post-filter candidate") + missing = _REQUIRED_CANDIDATE_KEYS - set(raw) + if missing: + raise ValueError( + "post-filter candidate is missing required fields: " + + ", ".join(sorted(missing)) + ) + action = str(raw.get("action") or "").strip() + if action not in AUTO_AGENT_MODE_ORDER: + raise ValueError(f"unsupported route action: {action or ''}") + if action in seen: + raise ValueError(f"duplicate post-filter candidate action: {action}") + selected = raw.get("selected", action == baseline_mode) + if not isinstance(selected, bool) or selected is not (action == baseline_mode): + raise ValueError("candidate selected flags must identify exactly the baseline action") + latency_tier = str(raw.get("latency_tier") or "").strip().lower() + if latency_tier not in _LATENCY_TIERS: + raise ValueError(f"unsupported latency_tier: {latency_tier or ''}") + normalized.append( + { + "action": action, + "estimated_cost_units": _finite_nonnegative( + raw.get("estimated_cost_units"), "estimated_cost_units" + ), + "estimated_model_calls": _nonnegative_int( + raw.get("estimated_model_calls"), "estimated_model_calls", minimum=1 + ), + "planned_loop_steps": _nonnegative_int( + raw.get("planned_loop_steps"), "planned_loop_steps" + ), + "latency_tier": latency_tier, + "selected": selected, + } + ) + seen.add(action) + if not normalized: + raise ValueError("post_filter_candidates must be a non-empty sequence") + if baseline_mode not in seen: + raise ValueError("baseline_mode must be present in post_filter_candidates") + return sorted(normalized, key=lambda row: AUTO_AGENT_MODE_ORDER.index(row["action"])) + + +def _normalize_exclusions(value: Any) -> List[Dict[str, Any]]: + if value is None: + value = [] + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError("post_filter_exclusions must be a sequence") + normalized: List[Dict[str, Any]] = [] + seen: set[str] = set() + for raw in value: + if not isinstance(raw, Mapping): + raise ValueError("post_filter_exclusions must contain JSON objects") + _strict_keys(raw, _EXCLUSION_KEYS, "post-filter exclusion") + action = str(raw.get("action") or "").strip() + if action not in AUTO_AGENT_MODE_ORDER: + raise ValueError(f"unsupported excluded route action: {action or ''}") + if action in seen: + raise ValueError(f"duplicate post-filter exclusion action: {action}") + reasons_raw = raw.get("reasons") + if not isinstance(reasons_raw, Sequence) or isinstance(reasons_raw, (str, bytes)): + raise ValueError("post-filter exclusion reasons must be a non-empty sequence") + reasons = [str(reason or "").strip() for reason in reasons_raw] + if not reasons or any(reason not in _ALLOWED_EXCLUSION_REASONS for reason in reasons): + raise ValueError("post-filter exclusions must use only versioned reason codes") + if len(set(reasons)) != len(reasons): + raise ValueError("post-filter exclusion reasons must be unique") + normalized.append({"action": action, "reasons": reasons}) + seen.add(action) + return sorted(normalized, key=lambda row: AUTO_AGENT_MODE_ORDER.index(row["action"])) + + +def _binomial_upper_quantile(trials: int, probability: float, confidence: float) -> int: + """Return the exact smallest k with Binomial(n, p) CDF(k) >= confidence.""" + + if probability <= 0.0: + return 0 + if probability >= 1.0: + return trials + log_p = math.log(probability) + log_q = math.log1p(-probability) + log_n_factorial = math.lgamma(trials + 1) + log_probabilities = [ + log_n_factorial + - math.lgamma(k + 1) + - math.lgamma(trials - k + 1) + + (k * log_p) + + ((trials - k) * log_q) + for k in range(trials + 1) + ] + pivot = max(log_probabilities) + scaled = [math.exp(item - pivot) for item in log_probabilities] + total = math.fsum(scaled) + threshold = confidence * total + cumulative = 0.0 + for k, mass in enumerate(scaled): + cumulative += mass + if cumulative >= threshold: + return k + return trials + + +def _binomial_below_probability( + trials: int, + probability: float, + target_successes: int, +) -> float: + """Return exact P[Binomial(trials, probability) < target_successes].""" + + if target_successes <= 0: + return 0.0 + if trials < target_successes or probability <= 0.0: + return 1.0 + if probability >= 1.0: + return 0.0 + log_q = math.log1p(-probability) + log_ratio_constant = math.log(probability) - log_q + log_probability = trials * log_q + lower_tail_logs = [] + for successes in range(target_successes): + lower_tail_logs.append(log_probability) + if successes + 1 < target_successes: + log_probability += ( + math.log(trials - successes) + - math.log(successes + 1) + + log_ratio_constant + ) + pivot = max(lower_tail_logs) + lower_tail = math.exp(pivot) * math.fsum( + math.exp(item - pivot) for item in lower_tail_logs + ) + return max(0.0, min(1.0, lower_tail)) + + +def _binomial_at_least_probability( + trials: int, + probability: float, + target_successes: int, +) -> float: + """Return exact P[Binomial(trials, probability) >= target_successes].""" + + return 1.0 - _binomial_below_probability( + trials, probability, target_successes + ) + + +def _binomial_log_pmf(trials: int, probability: float, successes: int) -> float: + if successes < 0 or successes > trials or not 0.0 < probability < 1.0: + return -math.inf + log_q = math.log1p(-probability) + return ( + trials * log_q + + successes * (math.log(probability) - log_q) + + math.fsum( + math.log(trials - index) - math.log(index + 1) + for index in range(successes) + ) + ) + + +def _logaddexp(left: float, right: float) -> float: + if left == -math.inf: + return right + if right == -math.inf: + return left + pivot = max(left, right) + return pivot + math.log1p(math.exp(-abs(left - right))) + + +def _multinomial_at_least_each_probability( + trials: int, + probabilities: Sequence[float], + target_successes: int, +) -> float: + """Exact P[X1 >= t and X2 >= t] for a three-cell multinomial.""" + + if len(probabilities) != 2: + raise ValueError("joint multinomial rehearsal requires exactly two alternates") + first, second = (float(probabilities[0]), float(probabilities[1])) + if trials < (2 * target_successes) or first <= 0.0 or second <= 0.0: + return 0.0 + remainder = 1.0 - first - second + if remainder <= 0.0: + raise ValueError("joint multinomial label probabilities must sum to less than 1") + first_lower = _binomial_below_probability(trials, first, target_successes) + second_lower = _binomial_below_probability(trials, second, target_successes) + + # Condition on X1=i. Given X1=i, X2 is Binomial(n-i, p2/(1-p1)). + # Summing i -math.inf + else 0.0 + ) + marginal_log_pmf = trials * math.log1p(-first) + marginal_log_ratio_constant = math.log(first) - math.log1p(-first) + rectangle_log_probability = -math.inf + conditional_trials = trials + for first_count in range(target_successes): + if conditional_cdf > 0.0: + rectangle_log_probability = _logaddexp( + rectangle_log_probability, + marginal_log_pmf + math.log(conditional_cdf), + ) + if first_count + 1 >= target_successes: + break + next_boundary_pmf = conditional_boundary_pmf * ( + (conditional_trials - boundary_successes) + / (conditional_trials * (1.0 - conditional_probability)) + ) + conditional_cdf = min( + 1.0, + conditional_cdf + (conditional_probability * next_boundary_pmf), + ) + conditional_boundary_pmf = next_boundary_pmf + conditional_trials -= 1 + marginal_log_pmf += ( + math.log(trials - first_count) + - math.log(first_count + 1) + + marginal_log_ratio_constant + ) + both_lower = ( + math.exp(rectangle_log_probability) + if rectangle_log_probability > -math.inf + else 0.0 + ) + joint = 1.0 - first_lower - second_lower + both_lower + return max(0.0, min(1.0, joint)) + + +def _simultaneous_label_probability( + trials: int, + label_probabilities: Sequence[float], + target_successes: int, +) -> float: + if len(label_probabilities) == 1: + return _binomial_at_least_probability( + trials, float(label_probabilities[0]), target_successes + ) + if len(label_probabilities) == 2: + return _multinomial_at_least_each_probability( + trials, label_probabilities, target_successes + ) + raise ValueError("study supports simultaneous targets for one or two alternates") + + +def _minimum_simultaneous_label_routes( + *, + label_probabilities: Sequence[float], + target_labels: int, + confidence: float, +) -> int: + """Invert the exact binomial or joint multinomial tail.""" + + if not label_probabilities or any(probability <= 0.0 for probability in label_probabilities): + raise ValueError("alternate label probabilities must be positive for exact inversion") + lower = target_labels * len(label_probabilities) + weakest = min(float(probability) for probability in label_probabilities) + upper = max(lower, int(math.ceil(target_labels / weakest))) + while _simultaneous_label_probability( + upper, label_probabilities, target_labels + ) < confidence: + if upper >= MAX_EXACT_BINOMIAL_ROUTE_FORECAST: + raise ValueError("exact simultaneous target route forecast exceeds the supported limit") + upper = min(MAX_EXACT_BINOMIAL_ROUTE_FORECAST, max(upper + 1, upper * 2)) + while lower < upper: + midpoint = (lower + upper) // 2 + if _simultaneous_label_probability( + midpoint, label_probabilities, target_labels + ) >= confidence: + upper = midpoint + else: + lower = midpoint + 1 + return lower + + +def _traffic_scenario( + probabilities: Mapping[str, float], + *, + planned_routes: int, + confidence: float, + exploration_rate: float, + baseline_mode: str, + assumed_feedback_rate: float, + target_observed_labels: int, +) -> Dict[str, Any]: + actions = list(probabilities) + marginal_confidence = 1.0 - ((1.0 - confidence) / max(1, len(actions))) + expected_by_action = { + action: round(planned_routes * float(probability), 6) + for action, probability in probabilities.items() + } + upper_by_action = { + action: _binomial_upper_quantile( + planned_routes, float(probability), marginal_confidence + ) + for action, probability in probabilities.items() + } + alternate_propensities = { + action: float(probability) + for action, probability in probabilities.items() + if action != baseline_mode + } + weakest_alternate_propensity = ( + min(alternate_propensities.values()) if alternate_propensities else None + ) + label_probabilities_by_action = { + action: round(propensity * assumed_feedback_rate, 12) + for action, propensity in alternate_propensities.items() + } + label_probabilities = list(label_probabilities_by_action.values()) + expected_labels_by_action = { + action: round(planned_routes * probability, 6) + for action, probability in label_probabilities_by_action.items() + } + expected_routes_by_action = { + action: round(target_observed_labels / probability, 6) + for action, probability in label_probabilities_by_action.items() + } + minimum_routes = ( + _minimum_simultaneous_label_routes( + label_probabilities=label_probabilities, + target_labels=target_observed_labels, + confidence=confidence, + ) + if label_probabilities + else None + ) + simultaneous_method = ( + "exact_binomial_tail_inversion_single_alternate" + if len(label_probabilities) == 1 + else ( + "exact_joint_multinomial_tail_inversion_two_alternates" + if len(label_probabilities) == 2 + else "not_applicable_no_alternate_actions" + ) + ) + probability_at_minimum = ( + _simultaneous_label_probability( + minimum_routes, label_probabilities, target_observed_labels + ) + if minimum_routes is not None + else None + ) + probability_at_previous = ( + _simultaneous_label_probability( + minimum_routes - 1, label_probabilities, target_observed_labels + ) + if minimum_routes is not None and minimum_routes > 0 + else None + ) + label_scenario = { + "analysis_type": "simultaneous_alternate_label_traffic_not_power_or_mnar_correction", + "target_scope": "at_least_target_observed_labels_on_every_alternate_action", + "alternate_actions": list(alternate_propensities), + "alternate_action_propensities": alternate_propensities, + "weakest_alternate_propensity": weakest_alternate_propensity, + "assumed_feedback_rate": assumed_feedback_rate, + "per_route_observed_label_probability_by_alternate_action": ( + label_probabilities_by_action + ), + "target_observed_labels_per_alternate_action": target_observed_labels, + "expected_observed_labels_at_planned_routes_by_alternate_action": ( + expected_labels_by_action + ), + "expected_routes_for_target_by_alternate_action": expected_routes_by_action, + "exact_simultaneous_target": { + "method": simultaneous_method, + "confidence_level": confidence, + "minimum_routes_for_target_on_every_alternate_action": minimum_routes, + "probability_at_minimum_routes": ( + round(probability_at_minimum, 12) + if probability_at_minimum is not None + else None + ), + "probability_at_previous_route_count": ( + round(probability_at_previous, 12) + if probability_at_previous is not None + else None + ), + }, + "feedback_rate_source": "caller_assumption_for_capacity_planning", + "missingness_identification_performed": False, + "mnar_correction_performed": False, + "not_power_analysis": True, + "interpretation": ( + "simultaneous per-alternate traffic planning under an assumed feedback rate; " + "not statistical power, causal identification, or correction for selective feedback" + ), + } + return { + "analysis_type": "traffic_capacity_scenario_not_statistical_power", + "planned_routes": planned_routes, + "assignment_model": "fixed_n_independent_route_assignments", + "expected": { + "routes_by_action": expected_by_action, + "exploration_routes": round(planned_routes * exploration_rate, 6), + }, + "high_probability": { + "confidence_level": confidence, + "method": "exact_binomial_marginal_upper_quantiles_with_bonferroni_union_bound", + "marginal_confidence_level": round(marginal_confidence, 12), + "simultaneous_confidence_at_least": confidence, + "upper_routes_by_action": upper_by_action, + "upper_exploration_routes": _binomial_upper_quantile( + planned_routes, exploration_rate, confidence + ), + "interpretation": "capacity scenario only; not an effect-detection guarantee", + }, + "not_power_analysis": True, + "power_analysis_performed": False, + "minimum_detectable_effect": None, + "sample_size_recommendation": None, + "policy_value_estimate": None, + "observed_label_scenario": label_scenario, + } + + +def _weighted_metric( + candidates: Mapping[str, Mapping[str, Any]], + probabilities: Mapping[str, float], + metric: str, +) -> float: + return round( + math.fsum( + float(probabilities[action]) * float(candidates[action][metric]) + for action in probabilities + ), + 6, + ) + + +def _resource_forecast( + candidates: Sequence[Mapping[str, Any]], + probabilities: Mapping[str, float], + baseline_mode: str, + traffic: Mapping[str, Any], +) -> Dict[str, Any]: + candidate_map = {str(row["action"]): row for row in candidates} + planned_routes = int(traffic["planned_routes"]) + expected_per_route = { + "cost_units": _weighted_metric(candidate_map, probabilities, "estimated_cost_units"), + "model_calls": _weighted_metric(candidate_map, probabilities, "estimated_model_calls"), + "loop_steps": _weighted_metric(candidate_map, probabilities, "planned_loop_steps"), + } + baseline = candidate_map[baseline_mode] + baseline_per_route = { + "cost_units": float(baseline["estimated_cost_units"]), + "model_calls": int(baseline["estimated_model_calls"]), + "loop_steps": int(baseline["planned_loop_steps"]), + } + expected_totals = { + key: round(float(value) * planned_routes, 6) + for key, value in expected_per_route.items() + } + baseline_totals = { + key: round(float(value) * planned_routes, 6) + for key, value in baseline_per_route.items() + } + upper_counts = traffic["high_probability"]["upper_routes_by_action"] + conservative_upper = { + "cost_units": round( + math.fsum( + int(upper_counts[action]) * float(candidate_map[action]["estimated_cost_units"]) + for action in probabilities + ), + 6, + ), + "model_calls": round( + math.fsum( + int(upper_counts[action]) * int(candidate_map[action]["estimated_model_calls"]) + for action in probabilities + ), + 6, + ), + "loop_steps": round( + math.fsum( + int(upper_counts[action]) * int(candidate_map[action]["planned_loop_steps"]) + for action in probabilities + ), + 6, + ), + } + return { + "basis": "caller_supplied_post_filter_route_economics", + "by_action": { + action: { + "probability": float(probabilities[action]), + "estimated_cost_units": float(candidate_map[action]["estimated_cost_units"]), + "estimated_model_calls": int(candidate_map[action]["estimated_model_calls"]), + "planned_loop_steps": int(candidate_map[action]["planned_loop_steps"]), + "latency_tier": str(candidate_map[action]["latency_tier"]), + } + for action in probabilities + }, + "baseline_per_route": baseline_per_route, + "expected_per_route": expected_per_route, + "expected_for_planned_routes": expected_totals, + "baseline_for_planned_routes": baseline_totals, + "expected_increment_vs_baseline": { + key: round(expected_totals[key] - baseline_totals[key], 6) + for key in expected_totals + }, + "high_probability_conservative_capacity_upper": { + "confidence_level": traffic["high_probability"]["confidence_level"], + "count_bound_method": traffic["high_probability"]["method"], + **conservative_upper, + }, + "forecast_only": True, + "budget_guarantee": False, + } + + +def _causal_boundaries() -> Dict[str, Any]: + return { + "deployment": "shadow_only", + "execution_enabled": False, + "execution_performed": False, + "io_performed": False, + "ledger_write_performed": False, + "model_inference_performed": False, + "off_policy_estimate_computed": False, + "causal_identification_performed": False, + "missingness_identification_performed": False, + "mnar_correction_performed": False, + "assumed_feedback_rate_is_not_an_observation_model": True, + "power_analysis_performed": False, + "automatic_promotion_allowed": False, + "baseline_performance_guarantee": False, + "logging_policy_optimality_claim": False, + "ledger_eligible": False, + "preassignment_commitment_sealed": False, + "nonce_grinding_resistant": False, + "activation_blockers": list(ACTIVATION_BLOCKERS), + "requires_separate_runtime_integration": True, + "requires_validated_external_ope": True, + "interpretation": "design, audit, and capacity-planning artifact only", + } + + +def plan_adjacent_route_study( + baseline_mode: str, + post_filter_candidates: Sequence[Mapping[str, Any]], + post_filter_exclusions: Sequence[Mapping[str, Any]], + *, + source_contract: Mapping[str, Any], + exploration_rate: float = DEFAULT_EXPLORATION_RATE, + planned_routes: int = DEFAULT_PLANNED_ROUTES, + scenario_confidence: float = DEFAULT_SCENARIO_CONFIDENCE, + assumed_feedback_rate: float = DEFAULT_ASSUMED_FEEDBACK_RATE, + target_observed_labels: int = DEFAULT_TARGET_OBSERVED_LABELS, +) -> Dict[str, Any]: + """Build a canonical, prompt-free, non-executing adjacent-route charter. + + The candidate/exclusion inputs must describe the complete final action + partition after capability, policy, and budget filters. Only immediate + feasible neighbors of the deterministic baseline receive exploration mass. + """ + + cooked_baseline = str(baseline_mode or "").strip() + if cooked_baseline not in AUTO_AGENT_MODE_ORDER: + raise ValueError("baseline_mode must be a canonical route action") + cooked_rate = _validate_rate(exploration_rate) + cooked_routes = _validate_planned_routes(planned_routes) + cooked_confidence = _validate_confidence(scenario_confidence) + cooked_feedback_rate = _validate_feedback_rate(assumed_feedback_rate) + cooked_target_labels = _validate_target_labels(target_observed_labels) + cooked_source_contract = _normalize_source_contract(source_contract) + candidates = _normalize_candidates(post_filter_candidates, cooked_baseline) + exclusions = _normalize_exclusions(post_filter_exclusions) + + candidate_actions = {row["action"] for row in candidates} + exclusion_actions = {row["action"] for row in exclusions} + if candidate_actions & exclusion_actions: + raise ValueError("post-filter candidates and exclusions must be disjoint") + if candidate_actions | exclusion_actions != set(AUTO_AGENT_MODE_ORDER): + raise ValueError("post-filter candidates and exclusions must partition all route actions") + + baseline_index = AUTO_AGENT_MODE_ORDER.index(cooked_baseline) + neighbor_actions = [ + row["action"] + for row in candidates + if abs(AUTO_AGENT_MODE_ORDER.index(row["action"]) - baseline_index) == 1 + ] + if len(neighbor_actions) > MAX_ADJACENT_NEIGHBORS: # defensive against mode-order changes + raise ValueError("study supports at most two adjacent feasible neighbors") + enrolled = bool(neighbor_actions) + if enrolled: + per_neighbor = round(cooked_rate / len(neighbor_actions), 12) + if per_neighbor + 1e-12 < MIN_POSITIVE_EXPLORATION_PROBABILITY: + raise ValueError( + "exploration_rate split across feasible neighbors must keep each at or above 0.05" + ) + study_actions = [ + mode + for mode in AUTO_AGENT_MODE_ORDER + if mode == cooked_baseline or mode in neighbor_actions + ] + probabilities = { + action: ( + round(1.0 - cooked_rate, 12) + if action == cooked_baseline + else per_neighbor + ) + for action in study_actions + } + applied_rate = cooked_rate + enrollment_reason = "eligible_adjacent_post_filter_support" + else: + study_actions = [cooked_baseline] + probabilities = {cooked_baseline: 1.0} + applied_rate = 0.0 + enrollment_reason = "no_feasible_adjacent_action" + + candidate_by_action = {row["action"]: row for row in candidates} + study_candidates = [ + { + "action": action, + "baseline": action == cooked_baseline, + "adjacent_distance": abs(AUTO_AGENT_MODE_ORDER.index(action) - baseline_index), + "estimated_cost_units": candidate_by_action[action]["estimated_cost_units"], + "estimated_model_calls": candidate_by_action[action]["estimated_model_calls"], + "planned_loop_steps": candidate_by_action[action]["planned_loop_steps"], + "latency_tier": candidate_by_action[action]["latency_tier"], + } + for action in study_actions + ] + study_exclusions = [ + { + "action": row["action"], + "reasons": list(row["reasons"]), + "post_filter_feasible": False, + } + for row in exclusions + ] + for row in candidates: + if row["action"] not in study_actions: + study_exclusions.append( + { + "action": row["action"], + "reasons": ["nonadjacent_exploration_guard"], + "post_filter_feasible": True, + } + ) + study_exclusions.sort(key=lambda row: AUTO_AGENT_MODE_ORDER.index(row["action"])) + + traffic = _traffic_scenario( + probabilities, + planned_routes=cooked_routes, + confidence=cooked_confidence, + exploration_rate=applied_rate, + baseline_mode=cooked_baseline, + assumed_feedback_rate=cooked_feedback_rate, + target_observed_labels=cooked_target_labels, + ) + resource_forecast = _resource_forecast( + study_candidates, + probabilities, + cooked_baseline, + traffic, + ) + charter = { + "source_contract": cooked_source_contract, + "prompt_free_contract": { + "prompt_free": True, + "raw_prompt_included": False, + "raw_session_id_included": False, + "free_form_text_fields_allowed": False, + "canonical_json": "sorted_keys_compact_utf8_no_nan", + }, + "enrollment": { + "eligible": enrolled, + "reason": enrollment_reason, + "baseline_action": cooked_baseline, + "adjacent_feasible_actions": list(neighbor_actions), + "maximum_adjacent_neighbors": MAX_ADJACENT_NEIGHBORS, + }, + "probability_design": { + "decision_type": "randomized" if enrolled else "deterministic_not_enrolled", + "probability_stage": "post_filter", + "requested_exploration_rate": cooked_rate, + "applied_exploration_rate": applied_rate, + "minimum_positive_exploration_probability": ( + min(probabilities[action] for action in neighbor_actions) + if neighbor_actions + else None + ), + "eligible_actions": list(study_actions), + "action_probabilities": dict(probabilities), + "assignment_algorithm": ASSIGNMENT_ALGORITHM, + "assignment_unit": "route", + "assignment_performed": False, + }, + "candidates": study_candidates, + "exclusions": study_exclusions, + "traffic_scenario": traffic, + "resource_forecast": resource_forecast, + "causal_boundaries": _causal_boundaries(), + } + payload = { + "schema_version": STUDY_PLAN_SCHEMA_VERSION, + "study": { + "study_id": STUDY_ID, + "study_version": STUDY_VERSION, + "label": STUDY_LABEL, + "design_kind": "bounded_exposure_adjacent_rehearsal", + }, + "charter": charter, + } + return {**payload, "design_hash": _domain_hash(_DESIGN_HASH_DOMAIN, payload)} + + +def _validated_assignment_plan(plan: Any) -> Dict[str, Any]: + if not isinstance(plan, Mapping): + raise ValueError("plan must be a canonical adjacent-route study object") + _strict_keys(plan, _PLAN_KEYS, "plan") + if set(plan) != _PLAN_KEYS: + raise ValueError("plan is missing required top-level fields") + if plan.get("schema_version") != STUDY_PLAN_SCHEMA_VERSION: + raise ValueError("unsupported study plan schema_version") + study = plan.get("study") + charter = plan.get("charter") + if not isinstance(study, Mapping) or not isinstance(charter, Mapping): + raise ValueError("plan study and charter must be JSON objects") + _strict_keys(study, _STUDY_KEYS, "plan study") + _strict_keys(charter, _CHARTER_KEYS, "plan charter") + if set(study) != _STUDY_KEYS or set(charter) != _CHARTER_KEYS: + raise ValueError("plan study or charter is missing required fields") + if study.get("study_id") != STUDY_ID or study.get("study_version") != STUDY_VERSION: + raise ValueError("plan does not identify Bounded-Exposure Adjacent-Route Rehearsal v1") + supplied_hash = str(plan.get("design_hash") or "").strip() + payload = { + "schema_version": plan["schema_version"], + "study": dict(study), + "charter": dict(charter), + } + expected_hash = _domain_hash(_DESIGN_HASH_DOMAIN, payload) + if len(supplied_hash) != 64 or supplied_hash != expected_hash: + raise ValueError("plan design_hash does not match its canonical charter") + enrollment = charter.get("enrollment") + probability_design = charter.get("probability_design") + source_contract = _normalize_source_contract(charter.get("source_contract")) + candidates = charter.get("candidates") + exclusions = charter.get("exclusions") + boundaries = charter.get("causal_boundaries") + if not all( + isinstance(item, Mapping) + for item in (enrollment, probability_design, boundaries) + ) or not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)): + raise ValueError("plan charter has invalid assignment fields") + if not isinstance(exclusions, Sequence) or isinstance(exclusions, (str, bytes)): + raise ValueError("plan charter exclusions must be a sequence") + if enrollment.get("eligible") is not True: + raise ValueError("plan is not eligible for randomized assignment") + _strict_keys(enrollment, _ENROLLMENT_KEYS, "plan enrollment") + _strict_keys(probability_design, _PROBABILITY_DESIGN_KEYS, "plan probability design") + if set(enrollment) != _ENROLLMENT_KEYS or set(probability_design) != _PROBABILITY_DESIGN_KEYS: + raise ValueError("plan enrollment or probability design is missing required fields") + if probability_design.get("decision_type") != "randomized": + raise ValueError("eligible plan must declare a randomized decision type") + if probability_design.get("probability_stage") != "post_filter": + raise ValueError("eligible plan probabilities must be post-filter") + if probability_design.get("assignment_algorithm") != ASSIGNMENT_ALGORITHM: + raise ValueError("plan uses an unsupported assignment algorithm") + if probability_design.get("assignment_unit") != "route": + raise ValueError("plan assignment_unit must be route") + if probability_design.get("assignment_performed") is not False: + raise ValueError("plan must not claim an assignment was already performed") + if boundaries.get("execution_enabled") is not False: + raise ValueError("planner assignments cannot enable execution") + if dict(boundaries) != _causal_boundaries(): + raise ValueError("plan causal boundaries do not match the non-executing study contract") + + baseline_action = str(enrollment.get("baseline_action") or "") + if baseline_action not in AUTO_AGENT_MODE_ORDER: + raise ValueError("plan baseline action is invalid") + neighbors_raw = enrollment.get("adjacent_feasible_actions") + if not isinstance(neighbors_raw, Sequence) or isinstance(neighbors_raw, (str, bytes)): + raise ValueError("plan adjacent feasible actions must be a sequence") + neighbors = [str(action) for action in neighbors_raw] + if not 1 <= len(neighbors) <= MAX_ADJACENT_NEIGHBORS or len(set(neighbors)) != len(neighbors): + raise ValueError("eligible plan must contain one or two unique adjacent actions") + baseline_index = AUTO_AGENT_MODE_ORDER.index(baseline_action) + if any( + action not in AUTO_AGENT_MODE_ORDER + or abs(AUTO_AGENT_MODE_ORDER.index(action) - baseline_index) != 1 + for action in neighbors + ): + raise ValueError("plan contains a nonadjacent exploration action") + + eligible_raw = probability_design.get("eligible_actions") + probabilities_raw = probability_design.get("action_probabilities") + if not isinstance(eligible_raw, Sequence) or isinstance(eligible_raw, (str, bytes)): + raise ValueError("plan eligible actions must be a sequence") + if not isinstance(probabilities_raw, Mapping): + raise ValueError("plan action probabilities must be a mapping") + eligible_actions = [str(action) for action in eligible_raw] + expected_eligible = [ + action + for action in AUTO_AGENT_MODE_ORDER + if action == baseline_action or action in neighbors + ] + if eligible_actions != expected_eligible or set(probabilities_raw) != set(expected_eligible): + raise ValueError("plan eligible actions do not match its adjacent enrollment set") + requested_rate = _validate_rate(probability_design.get("requested_exploration_rate")) + applied_rate = _validate_rate(probability_design.get("applied_exploration_rate")) + if not math.isclose(requested_rate, applied_rate, rel_tol=0.0, abs_tol=1e-12): + raise ValueError("eligible plan must apply its requested exploration rate") + expected_neighbor_probability = round(applied_rate / len(neighbors), 12) + if expected_neighbor_probability + 1e-12 < MIN_POSITIVE_EXPLORATION_PROBABILITY: + raise ValueError("plan alternate probabilities fall below the study floor") + cooked_probabilities: Dict[str, float] = {} + for action in eligible_actions: + raw_probability = probabilities_raw.get(action) + if isinstance(raw_probability, bool): + raise ValueError("plan action probabilities must be finite numbers") + try: + probability = float(raw_probability) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("plan action probabilities must be finite numbers") from exc + expected_probability = ( + round(1.0 - applied_rate, 12) + if action == baseline_action + else expected_neighbor_probability + ) + if not math.isfinite(probability) or not math.isclose( + probability, expected_probability, rel_tol=0.0, abs_tol=1e-12 + ): + raise ValueError("plan action probabilities do not match its versioned design") + cooked_probabilities[action] = probability + if not math.isclose(sum(cooked_probabilities.values()), 1.0, rel_tol=0.0, abs_tol=1e-9): + raise ValueError("plan action probabilities must sum to 1") + minimum_probability = probability_design.get("minimum_positive_exploration_probability") + if not isinstance(minimum_probability, (int, float)) or isinstance(minimum_probability, bool): + raise ValueError("plan minimum exploration probability is invalid") + if not math.isclose( + float(minimum_probability), + expected_neighbor_probability, + rel_tol=0.0, + abs_tol=1e-12, + ): + raise ValueError("plan minimum exploration probability does not match its distribution") + + normalized_candidates: List[Dict[str, Any]] = [] + candidate_actions: List[str] = [] + for raw_candidate in candidates: + if not isinstance(raw_candidate, Mapping): + raise ValueError("plan candidates must contain JSON objects") + _strict_keys(raw_candidate, _STUDY_CANDIDATE_KEYS, "plan candidate") + if set(raw_candidate) != _STUDY_CANDIDATE_KEYS: + raise ValueError("plan candidate is missing required fields") + action = str(raw_candidate.get("action") or "") + if action not in expected_eligible or action in candidate_actions: + raise ValueError("plan candidates must exactly match unique eligible actions") + baseline_flag = raw_candidate.get("baseline") + if not isinstance(baseline_flag, bool) or baseline_flag is not (action == baseline_action): + raise ValueError("plan candidate baseline flags are invalid") + distance = _nonnegative_int( + raw_candidate.get("adjacent_distance"), "adjacent_distance" + ) + expected_distance = abs(AUTO_AGENT_MODE_ORDER.index(action) - baseline_index) + if distance != expected_distance or distance not in {0, 1}: + raise ValueError("plan candidate adjacent distance is invalid") + latency_tier = str(raw_candidate.get("latency_tier") or "") + if latency_tier not in _LATENCY_TIERS: + raise ValueError("plan candidate latency tier is invalid") + normalized_candidates.append( + { + "action": action, + "baseline": baseline_flag, + "adjacent_distance": distance, + "estimated_cost_units": _finite_nonnegative( + raw_candidate.get("estimated_cost_units"), "estimated_cost_units" + ), + "estimated_model_calls": _nonnegative_int( + raw_candidate.get("estimated_model_calls"), + "estimated_model_calls", + minimum=1, + ), + "planned_loop_steps": _nonnegative_int( + raw_candidate.get("planned_loop_steps"), "planned_loop_steps" + ), + "latency_tier": latency_tier, + } + ) + candidate_actions.append(action) + if candidate_actions != expected_eligible: + raise ValueError("plan candidate order must match canonical eligible action order") + + normalized_exclusions: List[Dict[str, Any]] = [] + exclusion_actions: List[str] = [] + for raw_exclusion in exclusions: + if not isinstance(raw_exclusion, Mapping): + raise ValueError("plan exclusions must contain JSON objects") + _strict_keys(raw_exclusion, _STUDY_EXCLUSION_KEYS, "plan exclusion") + if set(raw_exclusion) != _STUDY_EXCLUSION_KEYS: + raise ValueError("plan exclusion is missing required fields") + action = str(raw_exclusion.get("action") or "") + if action not in AUTO_AGENT_MODE_ORDER or action in expected_eligible or action in exclusion_actions: + raise ValueError("plan exclusions must be unique and outside eligible actions") + reasons_raw = raw_exclusion.get("reasons") + if not isinstance(reasons_raw, Sequence) or isinstance(reasons_raw, (str, bytes)): + raise ValueError("plan exclusion reasons must be a sequence") + reasons = [str(reason) for reason in reasons_raw] + allowed_reasons = _ALLOWED_EXCLUSION_REASONS | {"nonadjacent_exploration_guard"} + if not reasons or any(reason not in allowed_reasons for reason in reasons): + raise ValueError("plan exclusion reasons are invalid") + feasible = raw_exclusion.get("post_filter_feasible") + if not isinstance(feasible, bool): + raise ValueError("plan exclusion feasibility flag must be boolean") + if feasible is not (reasons == ["nonadjacent_exploration_guard"]): + raise ValueError("plan exclusion feasibility contradicts its reason") + normalized_exclusions.append( + { + "action": action, + "reasons": reasons, + "post_filter_feasible": feasible, + } + ) + exclusion_actions.append(action) + if set(candidate_actions) | set(exclusion_actions) != set(AUTO_AGENT_MODE_ORDER): + raise ValueError("plan candidates and exclusions must partition all route actions") + if exclusion_actions != sorted( + exclusion_actions, key=lambda action: AUTO_AGENT_MODE_ORDER.index(action) + ): + raise ValueError("plan exclusions must use canonical action order") + return { + "plan": dict(plan), + "study": dict(study), + "charter": dict(charter), + "design_hash": supplied_hash, + "enrollment": dict(enrollment), + "probability_design": dict(probability_design), + "source_contract": source_contract, + "candidates": normalized_candidates, + "exclusions": normalized_exclusions, + } + + +def validate_adjacent_route_study(plan: Any) -> Dict[str, Any]: + """Validate an eligible rehearsal plan without drawing an assignment. + + The returned projection is a defensive, normalized view for other + prompt-free control-plane tools. It performs no randomness, I/O, ledger + writes, model inference, or activation. + """ + + try: + return _validated_assignment_plan(plan) + except ValueError as exc: + if str(exc) != "plan design_hash does not match its canonical charter": + raise + # JSON/JavaScript transports erase the lexical distinction between + # integral floats (``1.0``) and integers (``1``). Rebuild from the + # closed v1 fields and accept only exact semantic equality with the + # canonical builder result; this restores transport portability without + # accepting a changed plan or a caller-supplied replacement hash. + if not isinstance(plan, Mapping): + raise + charter = plan.get("charter") + if not isinstance(charter, Mapping): + raise + enrollment = charter.get("enrollment") + probability = charter.get("probability_design") + traffic = charter.get("traffic_scenario") + if not all(isinstance(item, Mapping) for item in (enrollment, probability, traffic)): + raise + observed = traffic.get("observed_label_scenario") + high_probability = traffic.get("high_probability") + candidates = charter.get("candidates") + exclusions = charter.get("exclusions") + if ( + not isinstance(observed, Mapping) + or not isinstance(high_probability, Mapping) + or not isinstance(candidates, Sequence) + or isinstance(candidates, (str, bytes)) + or not isinstance(exclusions, Sequence) + or isinstance(exclusions, (str, bytes)) + ): + raise + baseline = str(enrollment.get("baseline_action") or "") + reconstructed_candidates: List[Dict[str, Any]] = [] + for raw in candidates: + if not isinstance(raw, Mapping): + raise + reconstructed_candidates.append( + { + "action": raw.get("action"), + "estimated_cost_units": raw.get("estimated_cost_units"), + "estimated_model_calls": raw.get("estimated_model_calls"), + "planned_loop_steps": raw.get("planned_loop_steps"), + "latency_tier": raw.get("latency_tier"), + "selected": str(raw.get("action") or "") == baseline, + } + ) + reconstructed_exclusions: List[Dict[str, Any]] = [] + for raw in exclusions: + if not isinstance(raw, Mapping): + raise + if raw.get("post_filter_feasible") is True: + reconstructed_candidates.append( + { + "action": raw.get("action"), + "estimated_cost_units": 0.0, + "estimated_model_calls": 1, + "planned_loop_steps": 0, + "latency_tier": "unknown", + "selected": False, + } + ) + else: + reconstructed_exclusions.append( + {"action": raw.get("action"), "reasons": raw.get("reasons")} + ) + rebuilt = plan_adjacent_route_study( + baseline, + reconstructed_candidates, + reconstructed_exclusions, + source_contract=charter.get("source_contract"), + exploration_rate=probability.get("requested_exploration_rate"), + planned_routes=traffic.get("planned_routes"), + scenario_confidence=high_probability.get("confidence_level"), + assumed_feedback_rate=observed.get("assumed_feedback_rate"), + target_observed_labels=observed.get( + "target_observed_labels_per_alternate_action" + ), + ) + if dict(plan) != rebuilt: + raise + return _validated_assignment_plan(rebuilt) + + +def _nonce_bytes(value: Any) -> bytes: + if isinstance(value, bytes): + cooked = value + elif isinstance(value, str): + if not value.strip(): + raise ValueError("assignment_nonce must not be empty") + cooked = value.encode("utf-8") + else: + raise ValueError("assignment_nonce must be bytes or a string") + if not 16 <= len(cooked) <= 240: + raise ValueError("assignment_nonce must contain between 16 and 240 bytes") + return cooked + + +def assign_adjacent_route(plan: Mapping[str, Any], assignment_nonce: Any) -> Dict[str, Any]: + """Produce a deterministic rehearsal receipt without ledger-ready support.""" + + validated = _validated_assignment_plan(plan) + nonce = _nonce_bytes(assignment_nonce) + nonce_hash = hashlib.sha256(_NONCE_HASH_DOMAIN + nonce).hexdigest() + design_bytes = bytes.fromhex(validated["design_hash"]) + nonce_hash_bytes = bytes.fromhex(nonce_hash) + draw_digest = hashlib.sha256( + _DRAW_HASH_DOMAIN + design_bytes + nonce_hash_bytes + ).digest() + draw_int = int.from_bytes(draw_digest, "big", signed=False) + with localcontext() as context: + context.prec = 90 + draw = Decimal(draw_int) / Decimal(1 << 256) + cumulative = Decimal(0) + chosen_action = "" + probabilities = validated["probability_design"].get("action_probabilities") + eligible_actions = validated["probability_design"].get("eligible_actions") + if not isinstance(probabilities, Mapping) or not isinstance(eligible_actions, Sequence): + raise ValueError("plan probability design is invalid") + for action in eligible_actions: + cooked_action = str(action) + if cooked_action not in probabilities: + raise ValueError("eligible action is missing its assignment probability") + probability = Decimal(str(probabilities[cooked_action])) + if not probability.is_finite() or probability <= 0: + raise ValueError("eligible assignment probabilities must be positive and finite") + cumulative += probability + if not chosen_action and draw < cumulative: + chosen_action = cooked_action + if abs(cumulative - Decimal(1)) > Decimal("0.000001"): + raise ValueError("plan assignment probabilities must sum to 1") + if not chosen_action: + chosen_action = str(eligible_actions[-1]) + draw_decimal = format(draw, ".24f") + + applied_rate = float(validated["probability_design"]["applied_exploration_rate"]) + cooked_probabilities = { + str(action): float(validated["probability_design"]["action_probabilities"][action]) + for action in validated["probability_design"]["eligible_actions"] + } + baseline_action = str(validated["enrollment"]["baseline_action"]) + receipt = { + "schema_version": REHEARSAL_RECEIPT_SCHEMA_VERSION, + "rehearsal_only": True, + "study_design_hash": validated["design_hash"], + "assignment_algorithm": ASSIGNMENT_ALGORITHM, + "assignment_nonce_hash": nonce_hash, + "assignment_draw_hash": draw_digest.hex(), + "assignment_draw_unit_interval": draw_decimal, + "baseline_action": baseline_action, + "proposed_action": chosen_action, + "is_exploration_proposal": chosen_action != baseline_action, + "proposed_action_probability": float(cooked_probabilities[chosen_action]), + } + support_proposal = { + "schema_version": REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION, + "decision_type": "rehearsal_randomized_non_ledger", + "probability_stage": "post_filter_rehearsal_only", + "rehearsal_only": True, + "study_design_hash": validated["design_hash"], + "source_contract": dict(validated["source_contract"]), + "baseline_action": baseline_action, + "eligible_actions": list(cooked_probabilities), + "action_probabilities": cooked_probabilities, + "proposed_action": chosen_action, + "proposed_action_probability": float(cooked_probabilities[chosen_action]), + "exploration_rate": applied_rate, + "candidates": [dict(row) for row in validated["candidates"]], + "exclusions": [dict(row) for row in validated["exclusions"]], + "ledger_eligible": False, + "preassignment_commitment_sealed": False, + "nonce_grinding_resistant": False, + "activation_blockers": list(ACTIVATION_BLOCKERS), + } + return { + "schema_version": STUDY_ASSIGNMENT_SCHEMA_VERSION, + "study_id": STUDY_ID, + "study_version": STUDY_VERSION, + "design_hash": validated["design_hash"], + "rehearsal_assignment_receipt": receipt, + "rehearsal_support_proposal": support_proposal, + "ledger_eligible": False, + "preassignment_commitment_sealed": False, + "nonce_grinding_resistant": False, + "activation_blockers": list(ACTIVATION_BLOCKERS), + "side_effects": { + "io_performed": False, + "ledger_write_performed": False, + "execution_started": False, + "model_inference_started": False, + }, + "causal_boundaries": _causal_boundaries(), + } + + +__all__ = [ + "ACTIVATION_BLOCKERS", + "ASSIGNMENT_ALGORITHM", + "AUTO_AGENT_MODE_ORDER", + "DEFAULT_ASSUMED_FEEDBACK_RATE", + "DEFAULT_EXPLORATION_RATE", + "DEFAULT_PLANNED_ROUTES", + "DEFAULT_SCENARIO_CONFIDENCE", + "DEFAULT_TARGET_OBSERVED_LABELS", + "MAX_ADJACENT_NEIGHBORS", + "MIN_POSITIVE_EXPLORATION_PROBABILITY", + "REHEARSAL_RECEIPT_SCHEMA_VERSION", + "REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION", + "STUDY_ASSIGNMENT_SCHEMA_VERSION", + "STUDY_ID", + "STUDY_LABEL", + "STUDY_PLAN_SCHEMA_VERSION", + "STUDY_VERSION", + "assign_adjacent_route", + "plan_adjacent_route_study", + "validate_adjacent_route_study", +] diff --git a/source/route_policy_lab.py b/source/route_policy_lab.py new file mode 100644 index 000000000..d5be76414 --- /dev/null +++ b/source/route_policy_lab.py @@ -0,0 +1,1328 @@ +"""Honest, deterministic evidence summaries for Supermix route policies. + +This module deliberately does *not* estimate counterfactual policy value. It +can replay a threshold profile against historical route decisions and describe +outcomes for the subset where the candidate agrees with the observed action. +Those matched outcomes are associational: deterministic product logs do not +identify what would have happened under a different route. + +The propensity checks are intentionally strict. They describe whether a row +contains the minimum logging metadata a separate, validated off-policy +evaluator would need; passing them is not itself an off-policy estimate. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Union + +try: + from .route_policy_ledger import SUPPORT_SCHEMA_VERSION, build_logging_support_envelope +except ImportError: # Runtime entry points add ``source`` directly to sys.path. + from route_policy_ledger import SUPPORT_SCHEMA_VERSION, build_logging_support_envelope + + +AUTO_AGENT_MODE_ORDER: Tuple[str, ...] = ( + "off", + "collective", + "loop", + "collective_loop", +) + + +@dataclass(frozen=True) +class RoutePolicyProfile: + """Immutable score thresholds for the three non-trivial route modes.""" + + name: str + collective_min_score: int + loop_min_score: int + collective_loop_min_score: int + + def __post_init__(self) -> None: + thresholds = ( + self.collective_min_score, + self.loop_min_score, + self.collective_loop_min_score, + ) + if any(isinstance(value, bool) or not isinstance(value, int) for value in thresholds): + raise TypeError("profile thresholds must be integers") + if not (0 <= thresholds[0] <= thresholds[1] <= thresholds[2]): + raise ValueError("profile thresholds must be non-negative and monotonic") + + @property + def thresholds(self) -> Mapping[str, int]: + return MappingProxyType( + { + "off": 0, + "collective": self.collective_min_score, + "loop": self.loop_min_score, + "collective_loop": self.collective_loop_min_score, + } + ) + + def as_dict(self) -> Dict[str, Any]: + return {"name": self.name, "thresholds": dict(self.thresholds)} + + +POLICY_PROFILES: Mapping[str, RoutePolicyProfile] = MappingProxyType( + { + "efficiency": RoutePolicyProfile("efficiency", 3, 5, 6), + "balanced": RoutePolicyProfile("balanced", 2, 4, 5), + "quality_first": RoutePolicyProfile("quality_first", 1, 3, 4), + } +) + + +_MISSING = object() +_POSITIVE_RATINGS = frozenset(("up", "good", "approve", "approved", "positive")) +_NEGATIVE_RATINGS = frozenset(("down", "bad", "reject", "rejected", "negative")) +READINESS_SCHEMA_VERSION = "route-readiness-v2" +READINESS_MIN_TARGET_PROBABILITY = 0.05 +READINESS_MIN_GLOBAL_EFFECTIVE_SAMPLE_SIZE = 20.0 +READINESS_MIN_PER_ACTION_EFFECTIVE_SAMPLE_SIZE = 10.0 + + +def get_policy_profile(profile: Union[str, RoutePolicyProfile]) -> RoutePolicyProfile: + if isinstance(profile, RoutePolicyProfile): + return profile + key = str(profile or "").strip().lower() + try: + return POLICY_PROFILES[key] + except KeyError as exc: + choices = ", ".join(POLICY_PROFILES) + raise ValueError(f"unknown route policy profile {profile!r}; choose one of: {choices}") from exc + + +def ordered_allowed_modes(modes: Optional[Iterable[Any]]) -> Tuple[str, ...]: + """Return valid modes once each, in the canonical shallow-to-deep order. + + ``off`` is always retained as the safe fallback, including for empty or + malformed eligibility lists. + """ + + supplied = set() + if modes is not None and not isinstance(modes, (str, bytes)): + supplied = {str(mode).strip() for mode in modes if str(mode).strip() in AUTO_AGENT_MODE_ORDER} + supplied.add("off") + return tuple(mode for mode in AUTO_AGENT_MODE_ORDER if mode in supplied) + + +def select_profile_action( + score: Any, + allowed_modes: Optional[Iterable[Any]], + profile: Union[str, RoutePolicyProfile] = "balanced", +) -> str: + """Choose the deepest allowed mode whose immutable threshold is met.""" + + selected_profile = get_policy_profile(profile) + try: + cooked_score = int(score) + except (TypeError, ValueError, OverflowError): + cooked_score = 0 + cooked_score = max(0, cooked_score) + allowed = set(ordered_allowed_modes(allowed_modes)) + thresholds = selected_profile.thresholds + for mode in reversed(AUTO_AGENT_MODE_ORDER): + if mode in allowed and cooked_score >= thresholds[mode]: + return mode + return "off" + + +def profile_decision( + score: Any, + allowed_modes: Optional[Iterable[Any]], + profile: Union[str, RoutePolicyProfile] = "balanced", +) -> Dict[str, Any]: + selected_profile = get_policy_profile(profile) + ordered = ordered_allowed_modes(allowed_modes) + try: + cooked_score = max(0, int(score)) + except (TypeError, ValueError, OverflowError): + cooked_score = 0 + return { + "profile": selected_profile.name, + "thresholds": dict(selected_profile.thresholds), + "score": cooked_score, + "allowed_agent_modes": list(ordered), + "selected_agent_mode": select_profile_action(cooked_score, ordered, selected_profile), + } + + +def _finite_probability(value: Any) -> Optional[float]: + if isinstance(value, bool): + return None + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError): + return None + if not math.isfinite(cooked) or cooked < 0.0 or cooked > 1.0: + return None + return cooked + + +def propensity_readiness( + policy: Any, + *, + chosen_action: str, + target_action: Optional[str] = None, + expected_policy_id: Optional[str] = None, + expected_policy_version: Optional[str] = None, + expected_feature_schema: Optional[str] = None, + expected_context: Any = _MISSING, +) -> Dict[str, Any]: + """Validate strict post-filter behavior-policy logging for one decision. + + A ready row must be explicitly stochastic, use the exact requested policy + version and decision context, list at least two eligible actions, and carry + a normalized post-filter probability vector with positive probability for + the chosen action (and target action, when supplied). + """ + + reasons = [] + row = policy if isinstance(policy, Mapping) else {} + + expected_id = str(expected_policy_id or "").strip() + actual_id = str(row.get("policy_id") or "").strip() + if expected_id and actual_id != expected_id: + reasons.append("policy_id_mismatch") + + expected_version = str(expected_policy_version or "").strip() + actual_version = str(row.get("policy_version") or "").strip() + if not expected_version: + reasons.append("expected_policy_version_required") + elif actual_version != expected_version: + reasons.append("policy_version_mismatch") + + if expected_context is _MISSING: + reasons.append("expected_context_required") + elif "decision_context" not in row or row.get("decision_context") != expected_context: + reasons.append("decision_context_mismatch") + + decision_type = str(row.get("decision_type") or "").strip().lower() + if decision_type not in {"stochastic", "randomized"}: + reasons.append("decision_not_explicitly_stochastic") + + expected_schema = str(expected_feature_schema or "").strip() + actual_schema = str(row.get("feature_schema_version") or "").strip() + if expected_schema and actual_schema != expected_schema: + reasons.append("feature_schema_mismatch") + + if str(row.get("probability_stage") or "").strip() != "post_filter": + reasons.append("probabilities_not_post_filter") + + eligible_raw = row.get("eligible_actions") + eligible: Tuple[str, ...] = () + if not isinstance(eligible_raw, (list, tuple)) or isinstance(eligible_raw, (str, bytes)): + reasons.append("eligible_actions_missing") + else: + eligible_values = [str(action).strip() for action in eligible_raw] + if ( + len(eligible_values) < 2 + or len(set(eligible_values)) != len(eligible_values) + or any(action not in AUTO_AGENT_MODE_ORDER for action in eligible_values) + ): + reasons.append("eligible_actions_invalid") + else: + eligible = tuple(eligible_values) + + probabilities_raw = row.get("post_filter_action_probabilities") + probabilities: Dict[str, float] = {} + if not isinstance(probabilities_raw, Mapping): + reasons.append("post_filter_probability_vector_missing") + else: + invalid_probability = False + for action, value in probabilities_raw.items(): + action_key = str(action).strip() + probability = _finite_probability(value) + if action_key not in AUTO_AGENT_MODE_ORDER or probability is None: + invalid_probability = True + continue + probabilities[action_key] = probability + if invalid_probability: + reasons.append("post_filter_probability_vector_invalid") + if eligible and set(probabilities) != set(eligible): + reasons.append("probability_vector_not_exactly_eligible_actions") + if probabilities and not math.isclose(sum(probabilities.values()), 1.0, rel_tol=0.0, abs_tol=1e-6): + reasons.append("probabilities_do_not_sum_to_one") + + positive_actions = tuple( + action for action in AUTO_AGENT_MODE_ORDER if probabilities.get(action, 0.0) > 0.0 + ) + if len(positive_actions) < 2: + reasons.append("insufficient_randomized_support") + + chosen = str(chosen_action or "").strip() + if chosen not in eligible or probabilities.get(chosen, 0.0) <= 0.0: + reasons.append("chosen_action_has_no_positive_probability") + + target = str(target_action or "").strip() + if target_action is not None and (target not in eligible or probabilities.get(target, 0.0) <= 0.0): + reasons.append("target_action_has_no_positive_probability") + + support = row.get("logging_support") + if not isinstance(support, Mapping): + reasons.append("logging_support_missing") + else: + support_schema = str(support.get("schema_version") or "").strip() + if support_schema != SUPPORT_SCHEMA_VERSION: + reasons.append("support_schema_mismatch") + + support_decision_type = str(support.get("decision_type") or "").strip().lower() + if support_decision_type == "stochastic": + support_decision_type = "randomized" + normalized_row_decision_type = "randomized" if decision_type == "stochastic" else decision_type + if support_decision_type != normalized_row_decision_type: + reasons.append("support_decision_type_mismatch") + + support_stage = str(support.get("probability_stage") or "").strip() + if support_stage != "post_filter" or support_stage != str(row.get("probability_stage") or "").strip(): + reasons.append("support_probability_stage_mismatch") + + sampler = support.get("sampler") if isinstance(support.get("sampler"), Mapping) else {} + assignment_commitment = str(sampler.get("assignment_commitment") or "").strip() + if support_decision_type == "randomized" and not assignment_commitment: + reasons.append("assignment_commitment_missing") + + if eligible and set(probabilities) == set(eligible) and chosen in eligible: + try: + canonical_support = build_logging_support_envelope( + support, + eligible_modes=eligible, + action_probabilities=probabilities, + chosen_mode=chosen, + ) + except (TypeError, ValueError): + reasons.append("logging_support_invalid") + else: + canonical_chosen_probability = float(canonical_support["chosen_probability"]) + support_chosen_probability = _finite_probability(support.get("chosen_probability")) + row_logging_propensity = _finite_probability(row.get("logging_propensity")) + if ( + support_chosen_probability is None + or not math.isclose( + support_chosen_probability, + canonical_chosen_probability, + rel_tol=0.0, + abs_tol=1e-9, + ) + ): + reasons.append("chosen_probability_mismatch") + if ( + row_logging_propensity is None + or not math.isclose( + row_logging_propensity, + canonical_chosen_probability, + rel_tol=0.0, + abs_tol=1e-9, + ) + ): + reasons.append("logging_propensity_mismatch") + + canonical_candidate_hash = str(canonical_support["candidate_set_hash"]) + if ( + str(support.get("candidate_set_hash") or "") != canonical_candidate_hash + or str(row.get("candidate_set_hash") or "") != canonical_candidate_hash + ): + reasons.append("candidate_set_hash_mismatch") + + canonical_distribution_hash = str(canonical_support["distribution_hash"]) + if ( + str(support.get("distribution_hash") or "") != canonical_distribution_hash + or str(row.get("distribution_hash") or "") != canonical_distribution_hash + ): + reasons.append("distribution_hash_mismatch") + + # Preserve first-occurrence order while avoiding noisy duplicate failures. + reasons = list(dict.fromkeys(reasons)) + return { + "ready": not reasons, + "reasons": reasons, + "policy_id": actual_id or None, + "policy_version": actual_version or None, + "feature_schema_version": actual_schema or None, + "chosen_action": chosen or None, + "chosen_probability": probabilities.get(chosen), + "target_action": target or None, + "target_probability": probabilities.get(target) if target else None, + "eligible_actions": list(eligible), + "positive_support_actions": list(positive_actions), + "post_filter_action_probabilities": { + action: round(probabilities[action], 9) + for action in AUTO_AGENT_MODE_ORDER + if action in probabilities + }, + "causal_claim": False, + "off_policy_estimate": False, + } + + +def _rows(rows: Optional[Sequence[Any]]) -> Tuple[Mapping[str, Any], ...]: + return tuple(row for row in (rows or ()) if isinstance(row, Mapping)) + + +def _route_index(rows: Sequence[Mapping[str, Any]]) -> Tuple[Dict[str, Mapping[str, Any]], Dict[str, int]]: + indexed: Dict[str, Mapping[str, Any]] = {} + with_id = 0 + for row in rows: + route_id = str(row.get("route_id") or "").strip() + if not route_id: + continue + with_id += 1 + indexed[route_id] = row + return indexed, { + "rows": len(rows), + "rows_with_route_id": with_id, + "rows_without_route_id": len(rows) - with_id, + "unique_route_ids": len(indexed), + "duplicate_route_ids": with_id - len(indexed), + } + + +def _score(row: Mapping[str, Any]) -> Optional[int]: + policy = row.get("auto_agent_policy") if isinstance(row.get("auto_agent_policy"), Mapping) else {} + raw = policy.get("score") + if raw is None: + raw = row.get("score") + if isinstance(raw, bool): + return None + try: + cooked = int(raw) + except (TypeError, ValueError, OverflowError): + return None + if cooked < 0: + return None + if isinstance(raw, float) and (not math.isfinite(raw) or not raw.is_integer()): + return None + return cooked + + +def _observed_mode(row: Mapping[str, Any]) -> Optional[str]: + policy = row.get("auto_agent_policy") if isinstance(row.get("auto_agent_policy"), Mapping) else {} + mode = str( + row.get("chosen_agent_mode") + or row.get("chosen_mode") + or row.get("selected_agent_mode") + or policy.get("selected_agent_mode") + or "" + ).strip() + return mode if mode in AUTO_AGENT_MODE_ORDER else None + + +def _durable_execution_state(row: Mapping[str, Any]) -> Dict[str, Any]: + status = str(row.get("status") or row.get("route_status") or "").strip().lower() + status_valid = status in {"inflight", "completed", "failed"} + success = row.get("success") + if status == "completed": + status_success_consistent = success is True + elif status == "failed": + status_success_consistent = success is False + elif status == "inflight": + status_success_consistent = success is None + else: + status_success_consistent = False + + chosen_mode = _observed_mode(row) + executed_text = str( + row.get("executed_agent_mode") or row.get("executed_mode") or "" + ).strip() + executed_mode = executed_text if executed_text in AUTO_AGENT_MODE_ORDER else None + successful_terminal = bool(status == "completed" and success is True) + mode_mismatch = bool( + chosen_mode is not None + and executed_mode is not None + and chosen_mode != executed_mode + ) + execution_mode_verified = bool( + successful_terminal + and chosen_mode is not None + and executed_mode is not None + and not mode_mismatch + ) + return { + "status": status, + "status_valid": status_valid, + "status_success_consistent": status_success_consistent, + "successful_terminal": successful_terminal, + "chosen_mode": chosen_mode, + "executed_mode": executed_mode, + "executed_mode_present": bool(executed_text), + "mode_mismatch": mode_mismatch, + "execution_mode_verified": execution_mode_verified, + } + + +def _allowed_modes(row: Mapping[str, Any]) -> Tuple[Tuple[str, ...], bool]: + policy = row.get("auto_agent_policy") if isinstance(row.get("auto_agent_policy"), Mapping) else {} + raw = policy.get("allowed_agent_modes") + if raw is None: + raw = row.get("allowed_agent_modes") + explicit = isinstance(raw, (list, tuple)) and not isinstance(raw, (str, bytes)) + if not explicit: + return AUTO_AGENT_MODE_ORDER, False + return ordered_allowed_modes(raw), True + + +def _quality_outcome(row: Mapping[str, Any]) -> Optional[int]: + axes = row.get("feedback_axes") + if isinstance(axes, Mapping) and "quality" in axes: + quality = axes.get("quality") + if quality is None: + return None + try: + cooked = float(quality) + except (TypeError, ValueError, OverflowError): + return None + if cooked > 0: + return 1 + if cooked < 0: + return 0 + return None + + delta = row.get("score_delta") + if delta is not None and not isinstance(delta, bool): + try: + cooked_delta = float(delta) + except (TypeError, ValueError, OverflowError): + cooked_delta = 0.0 + if cooked_delta > 0: + return 1 + if cooked_delta < 0: + return 0 + + rating = str(row.get("rating") or "").strip().lower() + if rating in _POSITIVE_RATINGS: + return 1 + if rating in _NEGATIVE_RATINGS: + return 0 + return None + + +def _actual_economics( + usage: Mapping[str, Any], feedback: Mapping[str, Any] +) -> Mapping[str, Any]: + for row in (usage, feedback): + economics = row.get("route_economics") if isinstance(row.get("route_economics"), Mapping) else {} + actual = economics.get("actual") if isinstance(economics.get("actual"), Mapping) else {} + if actual: + return actual + return {} + + +def _finite_metric(value: Any) -> Optional[float]: + if isinstance(value, bool): + return None + try: + cooked = float(value) + except (TypeError, ValueError, OverflowError): + return None + if not math.isfinite(cooked) or cooked < 0.0: + return None + return cooked + + +def _ratio(numerator: int, denominator: int) -> Optional[float]: + return round(numerator / denominator, 6) if denominator else None + + +def _wilson_interval(approved: int, total: int, z: float = 1.959963984540054) -> Dict[str, Any]: + if total <= 0: + return { + "method": "descriptive_unweighted_wilson", + "confidence_level": 0.95, + "lower_bound": None, + "upper_bound": None, + "causal": False, + } + probability = approved / total + z2 = z * z + denominator = 1.0 + z2 / total + center = (probability + z2 / (2.0 * total)) / denominator + radius = z * math.sqrt((probability * (1.0 - probability) + z2 / (4.0 * total)) / total) / denominator + return { + "method": "descriptive_unweighted_wilson", + "confidence_level": 0.95, + "lower_bound": round(max(0.0, center - radius), 6), + "upper_bound": round(min(1.0, center + radius), 6), + "causal": False, + } + + +def _reason_counts(results: Sequence[Mapping[str, Any]]) -> Dict[str, int]: + counts: Dict[str, int] = {} + for result in results: + for reason in result.get("reasons") or (): + key = str(reason) + counts[key] = counts.get(key, 0) + 1 + return dict(sorted(counts.items())) + + +def _effective_sample_size(weights: Sequence[float]) -> float: + total = sum(float(weight) for weight in weights) + squared = sum(float(weight) * float(weight) for weight in weights) + return round((total * total) / squared, 6) if total > 0.0 and squared > 0.0 else 0.0 + + +def _expected_context_for_route( + route_id: str, + *, + expected_context: Any, + expected_context_by_route_id: Optional[Mapping[str, Any]], +) -> Any: + if expected_context_by_route_id is not None: + return expected_context_by_route_id.get(route_id, _MISSING) + return expected_context + + +def analyze_route_policy( + usage_rows: Optional[Sequence[Any]], + feedback_rows: Optional[Sequence[Any]], + *, + profile: Union[str, RoutePolicyProfile] = "balanced", + expected_policy_id: Optional[str] = None, + expected_policy_version: Optional[str] = None, + expected_feature_schema: Optional[str] = None, + expected_context: Any = _MISSING, + expected_context_by_route_id: Optional[Mapping[str, Any]] = None, + min_overlap_routes: int = 20, + min_target_probability: float = READINESS_MIN_TARGET_PROBABILITY, + min_global_effective_sample_size: float = READINESS_MIN_GLOBAL_EFFECTIVE_SAMPLE_SIZE, + min_per_action_effective_sample_size: float = READINESS_MIN_PER_ACTION_EFFECTIVE_SAMPLE_SIZE, + durable_lifecycle: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Build an associational, exact-route-id policy replay report. + + The report never imputes outcomes for changed actions. Only candidate + decisions that exactly match the historical action contribute to the + matched approval/cost/latency description. + """ + + if isinstance(min_overlap_routes, bool) or int(min_overlap_routes) < 1: + raise ValueError("min_overlap_routes must be at least 1") + min_overlap_routes = int(min_overlap_routes) + try: + min_target_probability = float(min_target_probability) + min_global_effective_sample_size = float(min_global_effective_sample_size) + min_per_action_effective_sample_size = float(min_per_action_effective_sample_size) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("readiness thresholds must be finite numbers") from exc + if not math.isfinite(min_target_probability) or not 0.0 < min_target_probability <= 1.0: + raise ValueError("min_target_probability must be in (0, 1]") + if ( + not math.isfinite(min_global_effective_sample_size) + or min_global_effective_sample_size <= 0.0 + or not math.isfinite(min_per_action_effective_sample_size) + or min_per_action_effective_sample_size <= 0.0 + ): + raise ValueError("effective sample size thresholds must be positive") + if expected_context_by_route_id is not None and not isinstance(expected_context_by_route_id, Mapping): + raise ValueError("expected_context_by_route_id must be a mapping") + selected_profile = get_policy_profile(profile) + usage = _rows(usage_rows) + feedback = _rows(feedback_rows) + usage_index, usage_counts = _route_index(usage) + feedback_index, feedback_counts = _route_index(feedback) + joined_ids = tuple(sorted(set(usage_index) & set(feedback_index))) + lifecycle_source = durable_lifecycle if isinstance(durable_lifecycle, Mapping) else None + durable_lifecycle_present = lifecycle_source is not None + + def fingerprint_valid(row: Mapping[str, Any]) -> bool: + if row.get("decision_record_fingerprint_valid") is True: + return True + policy = row.get("auto_agent_policy") + return bool( + isinstance(policy, Mapping) + and policy.get("decision_record_fingerprint_valid") is True + ) + + def fingerprint_reason(row: Mapping[str, Any]) -> str: + reason = row.get("decision_record_fingerprint_reason") + if not reason and isinstance(row.get("auto_agent_policy"), Mapping): + reason = row["auto_agent_policy"].get("decision_record_fingerprint_reason") + return str(reason or "missing_unverifiable") + + def outcome_contract_precommitted(row: Mapping[str, Any]) -> bool: + if row.get("outcome_contracts_precommitted_at_begin") is True: + return True + policy = row.get("auto_agent_policy") + return bool( + isinstance(policy, Mapping) + and policy.get("outcome_contracts_precommitted_at_begin") is True + ) + + fingerprint_valid_routes = sum(1 for row in usage_index.values() if fingerprint_valid(row)) + fingerprint_invalid_reason_counts: Dict[str, int] = {} + for row in usage_index.values(): + if fingerprint_valid(row): + continue + reason = fingerprint_reason(row) + fingerprint_invalid_reason_counts[reason] = fingerprint_invalid_reason_counts.get(reason, 0) + 1 + decision_record_fingerprint_complete = bool( + not durable_lifecycle_present + or ( + usage_counts["unique_route_ids"] > 0 + and fingerprint_valid_routes == usage_counts["unique_route_ids"] + ) + ) + precommitted_outcome_contract_routes = sum( + 1 for row in usage_index.values() if outcome_contract_precommitted(row) + ) + outcome_contract_integrity_complete = bool( + not durable_lifecycle_present + or ( + usage_counts["unique_route_ids"] > 0 + and precommitted_outcome_contract_routes + == usage_counts["unique_route_ids"] + ) + ) + + missing_or_invalid_score_ids = { + route_id for route_id, row in usage_index.items() if _score(row) is None + } + missing_or_invalid_chosen_mode_ids = { + route_id for route_id, row in usage_index.items() if _observed_mode(row) is None + } + unevaluable_usage_ids = missing_or_invalid_score_ids | missing_or_invalid_chosen_mode_ids + orphan_feedback_ids = set(feedback_index) - set(usage_index) + population_integrity_complete = bool( + not durable_lifecycle_present + or ( + usage_counts["rows"] > 0 + and usage_counts["rows_without_route_id"] == 0 + and usage_counts["duplicate_route_ids"] == 0 + and feedback_counts["rows_without_route_id"] == 0 + and feedback_counts["duplicate_route_ids"] == 0 + and not unevaluable_usage_ids + and not orphan_feedback_ids + ) + ) + + execution_states = { + route_id: _durable_execution_state(row) for route_id, row in usage_index.items() + } + row_status_counts = {"completed": 0, "failed": 0, "inflight": 0, "unknown": 0} + for state in execution_states.values(): + status = str(state["status"]) + if bool(state["status_valid"]): + row_status_counts[status] += 1 + else: + row_status_counts["unknown"] += 1 + status_success_mismatch_routes = sum( + 1 for state in execution_states.values() if not bool(state["status_success_consistent"]) + ) + successful_terminal_routes = sum( + 1 for state in execution_states.values() if bool(state["successful_terminal"]) + ) + unverified_successful_execution_routes = sum( + 1 + for state in execution_states.values() + if bool(state["successful_terminal"]) and not bool(state["execution_mode_verified"]) + ) + missing_executed_mode_routes = sum( + 1 + for state in execution_states.values() + if bool(state["successful_terminal"]) and not bool(state["executed_mode_present"]) + ) + invalid_executed_mode_routes = sum( + 1 + for state in execution_states.values() + if bool(state["successful_terminal"]) + and bool(state["executed_mode_present"]) + and state["executed_mode"] is None + ) + chosen_executed_mode_mismatch_routes = sum( + 1 for state in execution_states.values() if bool(state["mode_mismatch"]) + ) + feedback_route_status_mismatch_routes = sum( + 1 + for route_id in joined_ids + if str(feedback_index[route_id].get("route_status") or "").strip().lower() + and str(feedback_index[route_id].get("route_status") or "").strip().lower() + != str(execution_states[route_id]["status"]) + ) + execution_integrity_complete = bool( + not durable_lifecycle_present + or ( + row_status_counts["unknown"] == 0 + and status_success_mismatch_routes == 0 + and unverified_successful_execution_routes == 0 + and chosen_executed_mode_mismatch_routes == 0 + and feedback_route_status_mismatch_routes == 0 + ) + ) + + usage_scored = sum(1 for row in usage_index.values() if _score(row) is not None) + evaluable = 0 + matched = 0 + changed = 0 + explicit_allowed = 0 + assumed_allowed = 0 + approvals = [] + costs = [] + latencies = [] + propensity_results = [] + all_evaluable = 0 + all_changed = 0 + quality_observed = 0 + orphan_quality_outcomes = sum( + 1 + for route_id in orphan_feedback_ids + if _quality_outcome(feedback_index[route_id]) is not None + ) + raw_quality_outcomes = orphan_quality_outcomes + invalid_quality_outcomes = orphan_quality_outcomes + quality_on_non_successful_routes = 0 + quality_on_unverified_execution_routes = 0 + raw_observation_propensity_rows = 0 + versioned_observation_propensity_rows = 0 + positive_support_routes = 0 + target_probabilities: list[float] = [] + global_weights: list[float] = [] + per_action_work: Dict[str, Dict[str, Any]] = { + mode: { + "target_routes": 0, + "positive_support_routes": 0, + "observed_matches": 0, + "target_probabilities": [], + "weights": [], + } + for mode in AUTO_AGENT_MODE_ORDER + } + + for route_id in sorted(usage_index): + usage_row = usage_index[route_id] + score = _score(usage_row) + observed = _observed_mode(usage_row) + if score is None or observed is None: + continue + allowed, _was_explicit = _allowed_modes(usage_row) + candidate = select_profile_action(score, allowed, selected_profile) + all_evaluable += 1 + if candidate != observed: + all_changed += 1 + + policy = usage_row.get("auto_agent_policy") + result = propensity_readiness( + policy, + chosen_action=observed, + target_action=candidate, + expected_policy_id=expected_policy_id, + expected_policy_version=expected_policy_version, + expected_feature_schema=expected_feature_schema, + expected_context=_expected_context_for_route( + route_id, + expected_context=expected_context, + expected_context_by_route_id=expected_context_by_route_id, + ), + ) + propensity_results.append(result) + action_work = per_action_work[candidate] + action_work["target_routes"] += 1 + target_probability = result.get("target_probability") + if isinstance(target_probability, (int, float)) and float(target_probability) > 0.0: + cooked_target_probability = float(target_probability) + positive_support_routes += 1 + target_probabilities.append(cooked_target_probability) + action_work["positive_support_routes"] += 1 + action_work["target_probabilities"].append(cooked_target_probability) + + if result.get("ready") and observed == candidate: + chosen_probability = float(result.get("chosen_probability") or 0.0) + if chosen_probability > 0.0: + weight = 1.0 / chosen_probability + global_weights.append(weight) + action_work["weights"].append(weight) + action_work["observed_matches"] += 1 + + feedback_row = feedback_index.get(route_id) + quality = _quality_outcome(feedback_row) if feedback_row is not None else None + if quality is not None: + raw_quality_outcomes += 1 + execution_state = execution_states[route_id] + feedback_status = str(feedback_row.get("route_status") or "").strip().lower() + feedback_status_consistent = bool( + not feedback_status or feedback_status == str(execution_state["status"]) + ) + successful_route = bool(execution_state["successful_terminal"]) + verified_execution = bool(execution_state["execution_mode_verified"]) + quality_evidence_valid = bool( + not durable_lifecycle_present + or (successful_route and verified_execution and feedback_status_consistent) + ) + if quality_evidence_valid: + quality_observed += 1 + else: + invalid_quality_outcomes += 1 + if not successful_route: + quality_on_non_successful_routes += 1 + if successful_route and not verified_execution: + quality_on_unverified_execution_routes += 1 + if feedback_row is not None: + observation_probability = _finite_probability(feedback_row.get("observation_propensity")) + if observation_probability is not None and observation_probability > 0.0: + raw_observation_propensity_rows += 1 + observation_policy_id = str(feedback_row.get("observation_policy_id") or "").strip() + observation_policy_version = str( + feedback_row.get("observation_policy_version") or "" + ).strip() + outcome_definition_version = str( + feedback_row.get("outcome_definition_version") or "" + ).strip() + if observation_policy_id and observation_policy_version and outcome_definition_version: + versioned_observation_propensity_rows += 1 + + for route_id in joined_ids: + usage_row = usage_index[route_id] + feedback_row = feedback_index[route_id] + score = _score(usage_row) + observed = _observed_mode(usage_row) + if score is None or observed is None: + continue + allowed, was_explicit = _allowed_modes(usage_row) + if was_explicit: + explicit_allowed += 1 + else: + assumed_allowed += 1 + candidate = select_profile_action(score, allowed, selected_profile) + evaluable += 1 + + if candidate != observed: + changed += 1 + continue + matched += 1 + quality = _quality_outcome(feedback_row) + execution_state = execution_states[route_id] + feedback_status = str(feedback_row.get("route_status") or "").strip().lower() + quality_evidence_valid = bool( + not durable_lifecycle_present + or ( + bool(execution_state["successful_terminal"]) + and bool(execution_state["execution_mode_verified"]) + and (not feedback_status or feedback_status == str(execution_state["status"])) + ) + ) + if quality is not None and quality_evidence_valid: + approvals.append(quality) + economics = _actual_economics(usage_row, feedback_row) + cost = _finite_metric(economics.get("cost_units")) + latency = _finite_metric(economics.get("elapsed_ms")) + if cost is not None: + costs.append(cost) + if latency is not None: + latencies.append(latency) + + valid_propensity = sum(1 for result in propensity_results if result["ready"]) + approval_count = sum(approvals) + quality_count = len(approvals) + matched_observed = { + "routes": matched, + "quality_sample_count": quality_count, + "approved": approval_count, + "rejected": quality_count - approval_count, + "approval_rate": _ratio(approval_count, quality_count), + "approval_interval": _wilson_interval(approval_count, quality_count), + "cost_sample_count": len(costs), + "avg_cost_units": round(sum(costs) / len(costs), 6) if costs else None, + "total_cost_units": round(sum(costs), 6) if costs else None, + "max_cost_units": round(max(costs), 6) if costs else None, + "latency_sample_count": len(latencies), + "avg_elapsed_ms": round(sum(latencies) / len(latencies), 6) if latencies else None, + "max_elapsed_ms": round(max(latencies), 6) if latencies else None, + "interpretation": "observed outcomes where candidate and historical actions agree", + "causal": False, + } + + global_effective_sample_size = _effective_sample_size(global_weights) + per_action: Dict[str, Dict[str, Any]] = {} + for mode in AUTO_AGENT_MODE_ORDER: + work = per_action_work[mode] + if not int(work["target_routes"]): + continue + probabilities = list(work["target_probabilities"]) + per_action[mode] = { + "target_routes": int(work["target_routes"]), + "positive_support_routes": int(work["positive_support_routes"]), + "observed_matches": int(work["observed_matches"]), + "effective_sample_size": _effective_sample_size(work["weights"]), + "minimum_target_probability": round(min(probabilities), 9) if probabilities else None, + } + weakest_target_action = None + weakest_action_effective_sample_size = None + if per_action: + weakest_target_action, weakest_row = min( + per_action.items(), + key=lambda item: (float(item[1]["effective_sample_size"]), AUTO_AGENT_MODE_ORDER.index(item[0])), + ) + weakest_action_effective_sample_size = float(weakest_row["effective_sample_size"]) + + if durable_lifecycle_present: + lifecycle_report = ( + lifecycle_source.get("lifecycle") + if isinstance(lifecycle_source.get("lifecycle"), Mapping) + else lifecycle_source + ) + lifecycle_counts = ( + lifecycle_report.get("counts") if isinstance(lifecycle_report.get("counts"), Mapping) else {} + ) + analysis_window = ( + lifecycle_source.get("analysis_window") + if isinstance(lifecycle_source.get("analysis_window"), Mapping) + else {} + ) + durable_started = int(lifecycle_counts.get("started") or 0) + durable_completed = int(lifecycle_counts.get("completed") or 0) + durable_failed = int(lifecycle_counts.get("failed") or 0) + durable_inflight = int(lifecycle_counts.get("inflight") or 0) + row_lifecycle_counts_match = bool( + row_status_counts["unknown"] == 0 + and row_status_counts["completed"] == durable_completed + and row_status_counts["failed"] == durable_failed + and row_status_counts["inflight"] == durable_inflight + ) + lifecycle_reconciled = bool( + durable_started == usage_counts["unique_route_ids"] + and durable_completed + durable_failed + durable_inflight == durable_started + and durable_inflight == 0 + and not bool(analysis_window.get("truncated")) + and row_lifecycle_counts_match + ) + else: + durable_started = 0 + durable_completed = 0 + durable_failed = 0 + durable_inflight = 0 + row_lifecycle_counts_match = False + lifecycle_reconciled = False + + quality_missing = max(0, all_evaluable - quality_observed) + observation_propensity_logged = bool( + all_evaluable > 0 and versioned_observation_propensity_rows == all_evaluable + ) + outcome_evidence_integrity = bool( + not durable_lifecycle_present + or ( + invalid_quality_outcomes == 0 + and outcome_contract_integrity_complete + ) + ) + quality_observation_ready = bool( + all_evaluable > 0 + and quality_observed > 0 + and outcome_evidence_integrity + and (quality_missing == 0 or observation_propensity_logged) + ) + minimum_observed_target_probability = min(target_probabilities) if target_probabilities else None + logging_integrity_complete = bool( + all_evaluable > 0 + and valid_propensity == all_evaluable + and decision_record_fingerprint_complete + ) + target_probability_floor_met = bool( + logging_integrity_complete + and minimum_observed_target_probability is not None + and minimum_observed_target_probability >= min_target_probability + ) + global_overlap_met = global_effective_sample_size >= min_global_effective_sample_size + per_action_overlap_met = bool(per_action) and all( + float(row["effective_sample_size"]) >= min_per_action_effective_sample_size + for row in per_action.values() + ) + minimum_overlap_routes_met = valid_propensity >= min_overlap_routes + checks = { + "candidate_delta_present": all_changed > 0, + "population_integrity_complete": population_integrity_complete, + "execution_integrity_complete": execution_integrity_complete, + "logging_integrity_complete": logging_integrity_complete, + "minimum_overlap_routes_met": minimum_overlap_routes_met, + "target_probability_floor_met": target_probability_floor_met, + "global_overlap_ess_met": global_overlap_met, + "per_action_overlap_met": per_action_overlap_met, + "outcome_evidence_integrity": outcome_evidence_integrity, + "quality_observation_ready": quality_observation_ready, + "durable_lifecycle_present": durable_lifecycle_present, + "lifecycle_reconciled": lifecycle_reconciled, + } + ready_for_external_ope = all(checks.values()) + blocking_reasons: list[str] = [] + if all_evaluable == 0: + blocking_reasons.append("no_evaluable_routes") + elif valid_propensity == 0: + blocking_reasons.append("no_valid_randomized_overlap") + elif not logging_integrity_complete: + blocking_reasons.append("logging_integrity_incomplete") + if durable_lifecycle_present and not decision_record_fingerprint_complete: + blocking_reasons.append("decision_record_fingerprint_invalid") + if not checks["candidate_delta_present"]: + blocking_reasons.append("no_candidate_policy_delta") + if durable_lifecycle_present and not population_integrity_complete: + blocking_reasons.append("population_integrity_incomplete") + if usage_counts["duplicate_route_ids"] or feedback_counts["duplicate_route_ids"]: + blocking_reasons.append("duplicate_route_ids_in_durable_population") + if usage_counts["rows_without_route_id"] or feedback_counts["rows_without_route_id"]: + blocking_reasons.append("missing_route_ids_in_durable_population") + if unevaluable_usage_ids: + blocking_reasons.append("unevaluable_routes_in_durable_population") + if orphan_feedback_ids: + blocking_reasons.append("orphan_feedback_in_durable_population") + if durable_lifecycle_present and not execution_integrity_complete: + blocking_reasons.append("execution_integrity_incomplete") + if row_status_counts["unknown"] or status_success_mismatch_routes: + blocking_reasons.append("invalid_durable_route_state") + if unverified_successful_execution_routes: + blocking_reasons.append("unverified_successful_execution_mode") + if chosen_executed_mode_mismatch_routes: + blocking_reasons.append("chosen_executed_mode_mismatch") + if feedback_route_status_mismatch_routes: + blocking_reasons.append("feedback_route_status_mismatch") + if not minimum_overlap_routes_met: + blocking_reasons.append("insufficient_overlap_routes") + if not target_probability_floor_met: + blocking_reasons.append("insufficient_target_probability") + if not global_overlap_met: + blocking_reasons.append("insufficient_global_overlap_ess") + if not per_action_overlap_met: + blocking_reasons.append("insufficient_per_action_overlap") + if durable_lifecycle_present and not outcome_evidence_integrity: + blocking_reasons.append("outcome_evidence_integrity_failed") + if not outcome_contract_integrity_complete: + blocking_reasons.append("outcome_contract_not_precommitted") + if invalid_quality_outcomes: + blocking_reasons.append("quality_outcome_on_ineligible_route") + if quality_observed == 0: + blocking_reasons.append("no_observed_quality_outcomes") + if not quality_observation_ready: + blocking_reasons.append("unknown_quality_observation_process") + if not durable_lifecycle_present: + blocking_reasons.append("durable_lifecycle_required") + elif not lifecycle_reconciled: + blocking_reasons.append("lifecycle_not_reconciled") + blocking_reasons = list(dict.fromkeys(blocking_reasons)) + + evaluation_readiness = { + "schema_version": READINESS_SCHEMA_VERSION, + "thresholds": { + "minimum_valid_routes": min_overlap_routes, + "minimum_target_probability": min_target_probability, + "minimum_global_effective_sample_size": min_global_effective_sample_size, + "minimum_per_action_effective_sample_size": min_per_action_effective_sample_size, + "require_all_evaluable_logging_valid": True, + "require_complete_quality_observation_or_model": True, + "require_at_least_one_observed_quality_outcome": True, + "require_durable_population_and_execution_integrity": True, + "require_precommitted_outcome_contracts": True, + }, + "population_integrity": { + "durable_population": durable_lifecycle_present, + "usage_rows": usage_counts["rows"], + "usage_rows_with_route_id": usage_counts["rows_with_route_id"], + "usage_rows_without_route_id": usage_counts["rows_without_route_id"], + "unique_usage_route_ids": usage_counts["unique_route_ids"], + "duplicate_usage_route_ids": usage_counts["duplicate_route_ids"], + "feedback_rows": feedback_counts["rows"], + "feedback_rows_without_route_id": feedback_counts["rows_without_route_id"], + "duplicate_feedback_route_ids": feedback_counts["duplicate_route_ids"], + "orphan_feedback_routes": len(orphan_feedback_ids), + "evaluable_usage_routes": all_evaluable, + "unevaluable_usage_routes": len(unevaluable_usage_ids), + "missing_or_invalid_score_routes": len(missing_or_invalid_score_ids), + "missing_or_invalid_chosen_mode_routes": len(missing_or_invalid_chosen_mode_ids), + "complete": population_integrity_complete, + }, + "execution_integrity": { + "durable_population": durable_lifecycle_present, + "row_status_counts": dict(row_status_counts), + "row_lifecycle_counts_match": row_lifecycle_counts_match, + "successful_terminal_routes": successful_terminal_routes, + "failed_or_nonterminal_routes": row_status_counts["failed"] + row_status_counts["inflight"], + "status_success_mismatch_routes": status_success_mismatch_routes, + "missing_executed_mode_routes": missing_executed_mode_routes, + "invalid_executed_mode_routes": invalid_executed_mode_routes, + "unverified_successful_execution_routes": unverified_successful_execution_routes, + "chosen_executed_mode_mismatch_routes": chosen_executed_mode_mismatch_routes, + "feedback_route_status_mismatch_routes": feedback_route_status_mismatch_routes, + "complete": execution_integrity_complete, + }, + "logging_integrity": { + "evaluable_usage_routes": all_evaluable, + "checked_usage_routes": len(propensity_results), + "valid_routes": valid_propensity, + "invalid_routes": len(propensity_results) - valid_propensity, + "valid_rate": _ratio(valid_propensity, len(propensity_results)), + "all_evaluable_routes_valid": logging_integrity_complete, + "invalid_reason_counts": _reason_counts(propensity_results), + "decision_record_fingerprint_required": durable_lifecycle_present, + "decision_record_fingerprint_valid_routes": fingerprint_valid_routes, + "decision_record_fingerprint_invalid_routes": max( + 0, usage_counts["unique_route_ids"] - fingerprint_valid_routes + ), + "decision_record_fingerprint_invalid_reason_counts": fingerprint_invalid_reason_counts, + "decision_record_fingerprint_complete": decision_record_fingerprint_complete, + }, + "target_overlap": { + "positive_support_routes": positive_support_routes, + "unsupported_routes": max(0, all_evaluable - positive_support_routes), + "minimum_target_probability": ( + round(minimum_observed_target_probability, 9) + if minimum_observed_target_probability is not None + else None + ), + "nonzero_weight_routes": len(global_weights), + "effective_sample_size": global_effective_sample_size, + "max_importance_weight": round(max(global_weights), 6) if global_weights else None, + "per_action": per_action, + "weakest_target_action": weakest_target_action, + "weakest_action_effective_sample_size": weakest_action_effective_sample_size, + }, + "outcome_observation": { + "evaluable_routes": all_evaluable, + "raw_quality_outcome_routes": raw_quality_outcomes, + "quality_observed_routes": quality_observed, + "quality_missing_routes": quality_missing, + "quality_coverage_rate": _ratio(quality_observed, all_evaluable), + "has_observed_quality_outcomes": quality_observed > 0, + "invalid_quality_outcome_routes": invalid_quality_outcomes, + "orphan_quality_outcome_routes": orphan_quality_outcomes, + "quality_on_non_successful_routes": quality_on_non_successful_routes, + "quality_on_unverified_execution_routes": quality_on_unverified_execution_routes, + "outcome_contract_required": durable_lifecycle_present, + "precommitted_outcome_contract_routes": precommitted_outcome_contract_routes, + "missing_or_invalid_outcome_contract_routes": max( + 0, + usage_counts["unique_route_ids"] + - precommitted_outcome_contract_routes, + ), + "outcome_contract_integrity_complete": outcome_contract_integrity_complete, + "evidence_integrity_complete": outcome_evidence_integrity, + "observation_propensity_logged": observation_propensity_logged, + "raw_observation_propensity_routes": raw_observation_propensity_rows, + "versioned_observation_propensity_routes": versioned_observation_propensity_rows, + "required_observation_contract_fields": [ + "observation_policy_id", + "observation_policy_version", + "outcome_definition_version", + ], + "missingness_semantics": "unknown", + "ready": quality_observation_ready, + }, + "lifecycle_integrity": { + "durable_lifecycle_present": durable_lifecycle_present, + "durable_started_routes": durable_started, + "durable_terminal_routes": durable_completed + durable_failed, + "durable_inflight_routes": durable_inflight, + "replay_usage_routes": usage_counts["unique_route_ids"], + "replay_decision_coverage_rate": _ratio(usage_counts["unique_route_ids"], durable_started), + "reconciled": lifecycle_reconciled, + }, + "ready_for_external_ope": ready_for_external_ope, + "policy_value_estimated": False, + "causal_claim": False, + } + + propensity_summary = { + "checked_joined_evaluable_routes": evaluable, + "checked_evaluable_usage_routes": len(propensity_results), + "valid_routes": valid_propensity, + "invalid_routes": len(propensity_results) - valid_propensity, + "valid_rate": _ratio(valid_propensity, len(propensity_results)), + "required_policy_id": str(expected_policy_id or "") or None, + "required_policy_version": str(expected_policy_version or "") or None, + "required_feature_schema": str(expected_feature_schema or "") or None, + "required_decision_context": ( + "route_specific" + if expected_context_by_route_id is not None + else (None if expected_context is _MISSING else expected_context) + ), + "invalid_reason_counts": _reason_counts(propensity_results), + "minimum_routes_for_external_ope": min_overlap_routes, + "ready_as_external_ope_input": ready_for_external_ope, + "off_policy_estimate_computed": False, + "causal_claim": False, + } + + gate_status = "external_ope_required" if ready_for_external_ope else "blocked" + gate_reason = ( + "validated_off_policy_estimator_and_review_required" + if ready_for_external_ope + else (blocking_reasons[0] if blocking_reasons else "readiness_checks_failed") + ) + + warnings = [ + "Historical matched replay is associational and cannot identify outcomes for actions not taken.", + "The Wilson interval describes observed matched approvals; it is not a policy-value confidence interval.", + "The readiness certificate diagnoses overlap and observation integrity; it does not estimate policy value.", + "Missing quality feedback remains unknown unless an observation propensity or validated observation model is supplied.", + "Durable quality evidence counts only successfully completed routes whose executed and chosen modes agree.", + ] + if assumed_allowed: + warnings.append( + "Some rows lacked an explicit allowed_agent_modes set; canonical eligibility was assumed for replay only." + ) + + return { + "analysis_kind": "associational_matched_route_replay", + "profile": selected_profile.as_dict(), + "mode_order": list(AUTO_AGENT_MODE_ORDER), + "causal_interpretation": { + "causal": False, + "off_policy_estimate": False, + "counterfactual_actions_scored": False, + "label": "historical associational evidence", + }, + "support": { + "usage": usage_counts, + "feedback": feedback_counts, + "exact_joined_route_ids": len(joined_ids), + "exact_usage_join_coverage": _ratio(len(joined_ids), usage_counts["unique_route_ids"]), + "exact_feedback_join_coverage": _ratio(len(joined_ids), feedback_counts["unique_route_ids"]), + "usage_scored_routes": usage_scored, + "joined_evaluable_routes": evaluable, + "joined_explicit_allowed_modes": explicit_allowed, + "joined_assumed_canonical_allowed_modes": assumed_allowed, + }, + "candidate_action_agreement": { + "evaluable_routes": evaluable, + "matched_routes": matched, + "changed_routes": changed, + "agreement_rate": _ratio(matched, evaluable), + }, + "matched_observed": matched_observed, + "propensity_readiness": propensity_summary, + "evaluation_readiness": evaluation_readiness, + "promotion_gate": { + "status": gate_status, + "deployment": "shadow_only", + "gate_policy_version": READINESS_SCHEMA_VERSION, + "reason_code": gate_reason, + "blocking_reason_codes": blocking_reasons, + "checks": checks, + "passed_checks": sum(1 for passed in checks.values() if passed), + "total_checks": len(checks), + "automatic_promotion_allowed": False, + "requires_validated_external_ope": True, + "causal_claim": False, + }, + "warnings": warnings, + } + + +build_policy_evidence_report = analyze_route_policy + + +__all__ = [ + "AUTO_AGENT_MODE_ORDER", + "POLICY_PROFILES", + "READINESS_MIN_GLOBAL_EFFECTIVE_SAMPLE_SIZE", + "READINESS_MIN_PER_ACTION_EFFECTIVE_SAMPLE_SIZE", + "READINESS_MIN_TARGET_PROBABILITY", + "READINESS_SCHEMA_VERSION", + "RoutePolicyProfile", + "analyze_route_policy", + "build_policy_evidence_report", + "get_policy_profile", + "ordered_allowed_modes", + "profile_decision", + "propensity_readiness", + "select_profile_action", +] diff --git a/source/route_policy_ledger.py b/source/route_policy_ledger.py new file mode 100644 index 000000000..e568aaf56 --- /dev/null +++ b/source/route_policy_ledger.py @@ -0,0 +1,2511 @@ +"""Durable two-phase route-decision and feedback ledger. + +The ledger commits a decision before route execution begins, then records its +terminal outcome in a separate transaction. A process failure therefore +leaves durable, explicitly ``inflight`` evidence instead of silently dropping +the attempted route. SQLite WAL mode plus short per-operation connections and +``BEGIN IMMEDIATE`` transactions provide safe access from concurrent workers. + +Only a domain-separated SHA-256 digest of the session identifier is stored. +Missing feedback is reported as unknown; it is never interpreted as negative. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import math +import re +import sqlite3 +import time +import uuid +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple + + +LEDGER_SCHEMA_VERSION = 3 +DECISION_STATUSES: Tuple[str, ...] = ("inflight", "completed", "failed") +SUPPORT_SCHEMA_VERSION = "route-support-v1" +EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION = "route-execution-assignment-v1" +DECISION_FINGERPRINT_SCHEMA_VERSION = "route-decision-fingerprint-v1" +DECISION_TYPES: Tuple[str, ...] = ("deterministic", "randomized", "legacy_unknown") +OUTCOME_CONTRACT_SCHEMA_VERSION = "route-outcome-contract-v1" +OUTCOME_MATURITY_SCHEMA_VERSION = "route-outcome-maturity-v1" +OUTCOME_NAMES: Tuple[str, ...] = ( + "route_success", + "user_quality_rating", + "cost", + "latency", +) +_SESSION_HASH_DOMAIN = b"supermix-route-policy-ledger-v1\x00" +_CANDIDATE_HASH_DOMAIN = b"supermix-route-candidate-set-v1\x00" +_DISTRIBUTION_HASH_DOMAIN = b"supermix-route-distribution-v1\x00" +_DECISION_FINGERPRINT_DOMAIN = b"supermix-route-decision-record-v1\x00" +_OUTCOME_CONTRACT_HASH_DOMAIN = b"supermix-route-outcome-contract-v1\x00" +_EXECUTED_ASSIGNMENT_COMMITMENT_RE = re.compile( + rf"^{EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION}:[0-9a-f]{{64}}$" +) +_PROMPT_BEARING_CONTEXT_KEYS = frozenset( + { + "conversation", + "input", + "input_text", + "memory_context", + "message", + "messages", + "prompt", + "query", + "raw_prompt", + "raw_text", + "text", + "tool_context", + "user_prompt", + } +) + +_OUTCOME_CONTRACT_DEFAULTS: Dict[str, Dict[str, Any]] = { + "route_success": { + "outcome_definition_version": "route-success-v1", + "observation_policy_id": "route-completion", + "observation_policy_version": "1", + "value_type": "boolean", + "unit": "boolean", + "maturity_delay_seconds": 0.0, + }, + "user_quality_rating": { + "outcome_definition_version": "user-quality-rating-v1", + "observation_policy_id": "explicit-route-feedback", + "observation_policy_version": "1", + "value_type": "ordinal", + "unit": "signed_unit_interval", + "maturity_delay_seconds": 0.0, + }, + "cost": { + "outcome_definition_version": "route-cost-v1", + "observation_policy_id": "runtime-economics", + "observation_policy_version": "1", + "value_type": "number", + "unit": "cost_units", + "maturity_delay_seconds": 0.0, + }, + "latency": { + "outcome_definition_version": "route-latency-v1", + "observation_policy_id": "runtime-economics", + "observation_policy_version": "1", + "value_type": "number", + "unit": "milliseconds", + "maturity_delay_seconds": 0.0, + }, +} + + +class RoutePolicyLedgerError(RuntimeError): + """Base class for ledger failures with stable caller-facing semantics.""" + + +class DecisionNotFoundError(RoutePolicyLedgerError, KeyError): + """Raised when a route id does not exist in the ledger.""" + + +class LedgerConflictError(RoutePolicyLedgerError): + """Raised when an idempotency key or completed decision is reused differently.""" + + +def hash_session_identity(session_id: str) -> str: + """Return the stable, domain-separated session identity stored in SQLite.""" + + cooked = str(session_id or "").strip() + if not cooked: + raise ValueError("session_id is required") + return hashlib.sha256(_SESSION_HASH_DOMAIN + cooked.encode("utf-8")).hexdigest() + + +def _text(value: Any, name: str, *, limit: int = 240) -> str: + cooked = str(value or "").strip() + if not cooked: + raise ValueError(f"{name} is required") + if len(cooked) > limit: + raise ValueError(f"{name} must be at most {limit} characters") + return cooked + + +def _route_id(value: Any) -> str: + if value is None: + return str(uuid.uuid4()) + cooked = str(value).strip().lower() + if not cooked: + raise ValueError("route_id must be a UUID or omitted") + try: + uuid.UUID(cooked) + except (ValueError, AttributeError, TypeError) as exc: + raise ValueError("route_id must be a valid UUID") from exc + # Preserve a caller's valid hex-vs-hyphen representation so an already + # propagated route id remains byte-for-byte usable by feedback clients. + return cooked + + +def _existing_route_id(value: Any) -> str: + """Normalize a lookup key, treating invalid legacy ids as absent rows.""" + + try: + return _route_id(value) + except ValueError as exc: + raise DecisionNotFoundError(f"unknown route_id: {str(value or '').strip()}") from exc + + +def _json_mapping(value: Any, name: str) -> str: + if value is None: + value = {} + if not isinstance(value, Mapping): + raise ValueError(f"{name} must be a JSON object") + try: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be JSON serializable without non-finite numbers") from exc + + +def _eligible_modes(value: Any) -> List[str]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError("eligible_modes must be a non-empty sequence") + modes = [str(mode or "").strip() for mode in value] + if not modes or any(not mode for mode in modes): + raise ValueError("eligible_modes must contain non-empty mode names") + if len(set(modes)) != len(modes): + raise ValueError("eligible_modes must not contain duplicates") + return modes + + +def _probabilities(value: Any, eligible_modes: Sequence[str], chosen_mode: str) -> Dict[str, float]: + if value is None: + raise ValueError("action_probabilities are required for durable route decisions") + if not isinstance(value, Mapping): + raise ValueError("action_probabilities must be a JSON object") + probabilities: Dict[str, float] = {} + for raw_mode, raw_probability in value.items(): + mode = str(raw_mode or "").strip() + if isinstance(raw_probability, bool): + raise ValueError("action probabilities must be finite numbers between 0 and 1") + try: + probability = float(raw_probability) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("action probabilities must be finite numbers between 0 and 1") from exc + if not math.isfinite(probability) or probability < 0.0 or probability > 1.0: + raise ValueError("action probabilities must be finite numbers between 0 and 1") + probabilities[mode] = probability + if set(probabilities) != set(eligible_modes): + raise ValueError("action_probabilities must have exactly the eligible mode keys") + if not math.isclose(sum(probabilities.values()), 1.0, rel_tol=0.0, abs_tol=1e-6): + raise ValueError("action probabilities must sum to 1") + if probabilities.get(chosen_mode, 0.0) <= 0.0: + raise ValueError("chosen_mode must have positive logged probability") + return probabilities + + +def _legacy_probability_mapping(value: Any) -> Dict[str, float]: + """Decode a legacy probability object without trusting its JSON shape.""" + + try: + raw = json.loads(str(value)) + except (TypeError, ValueError, json.JSONDecodeError): + return {} + if not isinstance(raw, Mapping): + return {} + + probabilities: Dict[str, float] = {} + for raw_mode, raw_probability in raw.items(): + if isinstance(raw_probability, bool): + return {} + try: + probability = float(raw_probability) + except (TypeError, ValueError, OverflowError): + return {} + if not math.isfinite(probability): + return {} + probabilities[str(raw_mode)] = probability + return probabilities + + +def _canonical_json(value: Any, name: str) -> str: + try: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be JSON serializable without non-finite numbers") from exc + + +def _prompt_free_context(value: Any) -> Any: + """Remove raw-text fields from the read-only policy-evidence projection.""" + + if isinstance(value, Mapping): + sanitized: Dict[str, Any] = {} + for raw_key, raw_value in value.items(): + key = str(raw_key) + normalized = key.strip().lower() + if ( + normalized in _PROMPT_BEARING_CONTEXT_KEYS + or normalized.endswith("_prompt") + or normalized.endswith("_message") + ): + continue + sanitized[key] = _prompt_free_context(raw_value) + return sanitized + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [_prompt_free_context(item) for item in value] + return value + + +def _domain_hash(domain: bytes, value: Any, name: str) -> str: + return hashlib.sha256(domain + _canonical_json(value, name).encode("utf-8")).hexdigest() + + +def _decision_record_fingerprint( + *, + policy_name: str, + policy_version: str, + policy_schema_version: str, + decision_context: Mapping[str, Any], + eligible_modes: Sequence[str], + action_probabilities: Mapping[str, float], + chosen_mode: str, + candidate_set_hash: str, + distribution_hash: str, +) -> str: + """Hash the immutable, prompt-free route decision made before execution.""" + + payload = { + "schema_version": DECISION_FINGERPRINT_SCHEMA_VERSION, + "policy": { + "name": policy_name, + "version": policy_version, + "schema_version": policy_schema_version, + }, + "decision_context": dict(decision_context), + "eligible_modes": list(eligible_modes), + "action_probabilities": dict(action_probabilities), + "chosen_mode": chosen_mode, + "candidate_set_hash": candidate_set_hash, + "distribution_hash": distribution_hash, + } + return _domain_hash( + _DECISION_FINGERPRINT_DOMAIN, + payload, + "route decision record fingerprint", + ) + + +def _normalize_logging_support( + value: Optional[Mapping[str, Any]], + *, + eligible_modes: Sequence[str], + probabilities: Mapping[str, float], + chosen_mode: str, +) -> Dict[str, Any]: + """Validate and fingerprint the immutable post-filter logging envelope.""" + + positive_actions = [mode for mode in eligible_modes if float(probabilities.get(mode, 0.0)) > 0.0] + if value is None: + if len(positive_actions) != 1: + raise ValueError( + "randomized action_probabilities require logging_support with an assignment commitment" + ) + raw: Mapping[str, Any] = { + "schema_version": SUPPORT_SCHEMA_VERSION, + "decision_type": "deterministic", + "probability_stage": "post_filter", + "sampler": { + "name": "argmax", + "version": "1", + "exploration_rate": 0.0, + "assignment_unit": "route", + "assignment_commitment": None, + }, + "candidates": [{"action": mode} for mode in eligible_modes], + "exclusions": [], + } + elif isinstance(value, Mapping): + raw = value + else: + raise ValueError("logging_support must be a JSON object") + + schema_version = _text( + raw.get("schema_version") or SUPPORT_SCHEMA_VERSION, + "support schema_version", + limit=80, + ) + if schema_version != SUPPORT_SCHEMA_VERSION: + raise ValueError( + f"logging_support schema_version must be {SUPPORT_SCHEMA_VERSION}" + ) + if raw.get("shadow_only") is True or raw.get("rehearsal_only") is True: + raise ValueError("shadow-only or rehearsal-only support is never ledger eligible") + if "ledger_eligible" in raw and raw.get("ledger_eligible") is not True: + raise ValueError("logging_support explicitly marked ledger_eligible=false is rejected") + decision_type = str(raw.get("decision_type") or "").strip().lower() + if decision_type == "stochastic": + decision_type = "randomized" + if decision_type not in {"deterministic", "randomized"}: + raise ValueError("logging_support decision_type must be deterministic or randomized") + probability_stage = str(raw.get("probability_stage") or "").strip().lower() + if probability_stage != "post_filter": + raise ValueError("logging_support probability_stage must be post_filter") + + candidates_raw = raw.get("candidates") + if not isinstance(candidates_raw, Sequence) or isinstance(candidates_raw, (str, bytes)): + raise ValueError("logging_support candidates must be a sequence") + candidates: List[Dict[str, Any]] = [] + candidate_actions: List[str] = [] + for candidate in candidates_raw: + if not isinstance(candidate, Mapping): + raise ValueError("logging_support candidates must be JSON objects") + action = _text(candidate.get("action"), "candidate action", limit=80) + candidate_actions.append(action) + normalized_candidate = dict(candidate) + normalized_candidate["action"] = action + _canonical_json(normalized_candidate, "logging_support candidate") + candidates.append(normalized_candidate) + if candidate_actions != list(eligible_modes): + raise ValueError("logging_support candidate actions must exactly match eligible_modes order") + + exclusions_raw = raw.get("exclusions") or [] + if not isinstance(exclusions_raw, Sequence) or isinstance(exclusions_raw, (str, bytes)): + raise ValueError("logging_support exclusions must be a sequence") + exclusions: List[Dict[str, Any]] = [] + excluded_actions: set[str] = set() + for exclusion in exclusions_raw: + if not isinstance(exclusion, Mapping): + raise ValueError("logging_support exclusions must be JSON objects") + action = _text(exclusion.get("action"), "excluded action", limit=80) + if action in eligible_modes or action in excluded_actions: + raise ValueError("logging_support exclusions must be unique and outside eligible_modes") + reasons_raw = exclusion.get("reasons") + if not isinstance(reasons_raw, Sequence) or isinstance(reasons_raw, (str, bytes)): + raise ValueError("logging_support exclusion reasons must be a non-empty sequence") + reasons = [_text(reason, "exclusion reason", limit=120) for reason in reasons_raw] + if not reasons: + raise ValueError("logging_support exclusion reasons must be a non-empty sequence") + excluded_actions.add(action) + normalized_exclusion = dict(exclusion) + normalized_exclusion.update({"action": action, "reasons": reasons}) + _canonical_json(normalized_exclusion, "logging_support exclusion") + exclusions.append(normalized_exclusion) + + sampler_raw = raw.get("sampler") + if not isinstance(sampler_raw, Mapping): + raise ValueError("logging_support sampler must be a JSON object") + if sampler_raw.get("shadow_only") is True or sampler_raw.get("rehearsal_only") is True: + raise ValueError("shadow-only or rehearsal-only sampler support is never ledger eligible") + if "ledger_eligible" in sampler_raw and sampler_raw.get("ledger_eligible") is not True: + raise ValueError("sampler explicitly marked ledger_eligible=false is rejected") + sampler_name = _text(sampler_raw.get("name"), "sampler name", limit=80) + sampler_version = _text(sampler_raw.get("version"), "sampler version", limit=80) + assignment_unit = _text(sampler_raw.get("assignment_unit") or "route", "assignment unit", limit=80) + exploration_raw = sampler_raw.get("exploration_rate", 0.0) + if isinstance(exploration_raw, bool): + raise ValueError("sampler exploration_rate must be a finite number between 0 and 1") + try: + exploration_rate = float(exploration_raw) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("sampler exploration_rate must be a finite number between 0 and 1") from exc + if not math.isfinite(exploration_rate) or not 0.0 <= exploration_rate <= 1.0: + raise ValueError("sampler exploration_rate must be a finite number between 0 and 1") + commitment_raw = sampler_raw.get("assignment_commitment") + assignment_commitment = str(commitment_raw or "").strip() or None + if assignment_commitment is not None and len(assignment_commitment) > 240: + raise ValueError("sampler assignment_commitment must be at most 240 characters") + if assignment_commitment is not None and assignment_commitment.lower().startswith( + ("shadow:", "shadow-v1:", "route-study-shadow", "route-shadow:") + ): + raise ValueError("shadow assignment commitments cannot support executed ledger writes") + if assignment_commitment is not None and not _EXECUTED_ASSIGNMENT_COMMITMENT_RE.fullmatch( + assignment_commitment + ): + raise ValueError( + "sampler assignment_commitment must use the closed " + f"{EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION}: namespace" + ) + + if decision_type == "deterministic": + if len(positive_actions) != 1 or exploration_rate != 0.0: + raise ValueError("deterministic logging_support requires one positive action and zero exploration") + else: + if len(positive_actions) < 2: + raise ValueError("randomized logging_support requires at least two positive actions") + if exploration_rate <= 0.0: + raise ValueError("randomized logging_support requires positive exploration_rate") + if assignment_commitment is None: + raise ValueError("randomized logging_support requires an assignment commitment") + + sampler = { + "name": sampler_name, + "version": sampler_version, + "exploration_rate": exploration_rate, + "assignment_unit": assignment_unit, + "assignment_commitment": assignment_commitment, + } + envelope = { + "schema_version": schema_version, + "decision_type": decision_type, + "probability_stage": probability_stage, + "sampler": sampler, + "candidates": candidates, + "exclusions": exclusions, + "chosen_probability": float(probabilities[chosen_mode]), + } + candidate_payload = { + "schema_version": schema_version, + "candidates": candidates, + "exclusions": exclusions, + } + distribution_payload = { + "schema_version": schema_version, + "decision_type": decision_type, + "probability_stage": probability_stage, + "sampler": sampler, + "eligible_modes": list(eligible_modes), + "action_probabilities": dict(probabilities), + } + envelope["candidate_set_hash"] = _domain_hash( + _CANDIDATE_HASH_DOMAIN, candidate_payload, "candidate support envelope" + ) + envelope["distribution_hash"] = _domain_hash( + _DISTRIBUTION_HASH_DOMAIN, distribution_payload, "logging distribution" + ) + return envelope + + +def _utc(timestamp: Optional[float]) -> Optional[str]: + if timestamp is None: + return None + return datetime.fromtimestamp(float(timestamp), tz=timezone.utc).isoformat().replace("+00:00", "Z") + + +def build_logging_support_envelope( + logging_support: Optional[Mapping[str, Any]], + *, + eligible_modes: Sequence[str], + action_probabilities: Mapping[str, Any], + chosen_mode: str, +) -> Dict[str, Any]: + """Build the same canonical support envelope used by durable writes.""" + + modes = _eligible_modes(eligible_modes) + cooked_chosen = _text(chosen_mode, "chosen_mode", limit=80) + if cooked_chosen not in modes: + raise ValueError("chosen_mode must be present in eligible_modes") + probabilities = _probabilities(action_probabilities, modes, cooked_chosen) + return _normalize_logging_support( + logging_support, + eligible_modes=modes, + probabilities=probabilities, + chosen_mode=cooked_chosen, + ) + + +def build_route_outcome_contracts( + outcome_contracts: Optional[Mapping[str, Mapping[str, Any]]] = None, + *, + precommitted: bool = True, + commitment_source: Optional[str] = None, +) -> Dict[str, Dict[str, Any]]: + """Return the canonical, complete Route Outcome Contract v1 set. + + Callers may omit the set entirely to use conservative built-in definitions, + or supply exactly the four supported outcome names with per-outcome + definition, observation-policy, and maturity overrides. The builder never + accepts additional outcomes: changing an outcome universe requires a new + schema version rather than silently changing an existing experiment. + """ + + if not isinstance(precommitted, bool): + raise ValueError("precommitted must be a boolean") + if outcome_contracts is not None and not isinstance(outcome_contracts, Mapping): + raise ValueError("outcome_contracts must be a JSON object") + if outcome_contracts is not None and set(outcome_contracts) != set(OUTCOME_NAMES): + raise ValueError("outcome_contracts must contain exactly the four canonical outcome names") + + if commitment_source is None: + supplied_sources = { + str(contract.get("commitment_source") or "").strip() + for contract in (outcome_contracts or {}).values() + if isinstance(contract, Mapping) and contract.get("commitment_source") + } + if len(supplied_sources) == 1: + commitment_source = supplied_sources.pop() + elif supplied_sources: + raise ValueError("outcome contracts must use one commitment_source") + else: + commitment_source = ( + "legacy_posthoc" + if not precommitted + else ("caller" if outcome_contracts is not None else "safe_default") + ) + cooked_source = _text(commitment_source, "commitment_source", limit=80) + if not precommitted and cooked_source != "legacy_posthoc": + raise ValueError("non-precommitted outcome contracts must be legacy_posthoc") + if precommitted and cooked_source == "legacy_posthoc": + raise ValueError("precommitted outcome contracts cannot be legacy_posthoc") + + contracts: Dict[str, Dict[str, Any]] = {} + for outcome_name in OUTCOME_NAMES: + defaults = _OUTCOME_CONTRACT_DEFAULTS[outcome_name] + raw = (outcome_contracts or {}).get(outcome_name) or {} + if not isinstance(raw, Mapping): + raise ValueError(f"outcome contract {outcome_name} must be a JSON object") + + schema_version = str(raw.get("schema_version") or OUTCOME_CONTRACT_SCHEMA_VERSION).strip() + if schema_version != OUTCOME_CONTRACT_SCHEMA_VERSION: + raise ValueError(f"outcome contract {outcome_name} has an unsupported schema_version") + raw_name = str(raw.get("outcome_name") or outcome_name).strip() + if raw_name != outcome_name: + raise ValueError("outcome contract outcome_name must match its mapping key") + maturity_basis = str(raw.get("maturity_basis") or "decision_started_at").strip() + if maturity_basis != "decision_started_at": + raise ValueError("outcome contract maturity_basis must be decision_started_at") + + definition_version = _text( + raw.get("outcome_definition_version") + or defaults["outcome_definition_version"], + f"{outcome_name} outcome_definition_version", + limit=120, + ) + observation_policy_id = _text( + raw.get("observation_policy_id") or defaults["observation_policy_id"], + f"{outcome_name} observation_policy_id", + limit=120, + ) + observation_policy_version = _text( + raw.get("observation_policy_version") + or defaults["observation_policy_version"], + f"{outcome_name} observation_policy_version", + limit=120, + ) + value_type = str(raw.get("value_type") or defaults["value_type"]).strip() + unit = str(raw.get("unit") or defaults["unit"]).strip() + if value_type != defaults["value_type"] or unit != defaults["unit"]: + raise ValueError(f"outcome contract {outcome_name} cannot change value_type or unit") + if "precommitted" in raw and raw.get("precommitted") is not precommitted: + raise ValueError(f"outcome contract {outcome_name} precommitted flag conflicts with the write phase") + raw_source = str(raw.get("commitment_source") or cooked_source).strip() + if raw_source != cooked_source: + raise ValueError("outcome contracts must use one commitment_source") + + delay_raw = raw.get("maturity_delay_seconds", defaults["maturity_delay_seconds"]) + if isinstance(delay_raw, bool): + raise ValueError("maturity_delay_seconds must be a finite non-negative number") + try: + maturity_delay_seconds = float(delay_raw) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("maturity_delay_seconds must be a finite non-negative number") from exc + if not math.isfinite(maturity_delay_seconds) or maturity_delay_seconds < 0.0: + raise ValueError("maturity_delay_seconds must be a finite non-negative number") + + canonical: Dict[str, Any] = { + "schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + "outcome_name": outcome_name, + "outcome_definition_version": definition_version, + "observation_policy_id": observation_policy_id, + "observation_policy_version": observation_policy_version, + "value_type": value_type, + "unit": unit, + "maturity_delay_seconds": maturity_delay_seconds, + "maturity_basis": "decision_started_at", + "precommitted": precommitted, + "commitment_source": cooked_source, + } + contract_hash = _domain_hash( + _OUTCOME_CONTRACT_HASH_DOMAIN, + canonical, + f"{outcome_name} outcome contract", + ) + supplied_hash = raw.get("contract_hash") + if supplied_hash is not None and str(supplied_hash).strip() != contract_hash: + raise ValueError(f"outcome contract {outcome_name} contract_hash does not match its canonical fields") + canonical["contract_hash"] = contract_hash + contracts[outcome_name] = canonical + return contracts + + +def _finite_actual_metric(actual_economics: Mapping[str, Any], key: str) -> Optional[float]: + """Read current flattened economics and the older ``actual`` nesting.""" + + raw = actual_economics.get(key) + if raw is None: + nested = actual_economics.get("actual") + if isinstance(nested, Mapping): + raw = nested.get(key) + if isinstance(raw, bool) or raw is None: + return None + try: + value = float(raw) + except (TypeError, ValueError, OverflowError): + return None + return value if math.isfinite(value) and value >= 0.0 else None + + +def _quality_observation( + feedback: Mapping[str, Any], + *, + revision: int, +) -> Tuple[str, Optional[Any], Dict[str, Any]]: + """Project only an explicit quality signal; never infer one from absence.""" + + intent = str(feedback.get("feedback_intent") or feedback.get("intent") or "").strip().lower() + metadata: Dict[str, Any] = { + "feedback_revision": int(revision), + "feedback_intent": intent or None, + "quality_signal_source": None, + "raw_rating": None, + } + non_quality_intents = { + "cost", + "latency", + "lower_cost", + "faster", + "too_costly", + "too_slow", + } + if intent in non_quality_intents or str(feedback.get("observation_status") or "").strip().lower() == "not_observed": + return "not_observed", None, metadata + + axes = feedback.get("feedback_axes") + if isinstance(axes, Mapping): + raw_quality = axes.get("quality") + if not isinstance(raw_quality, bool) and raw_quality is not None: + try: + quality = float(raw_quality) + except (TypeError, ValueError, OverflowError): + quality = math.nan + if math.isfinite(quality) and -1.0 <= quality <= 1.0: + metadata["quality_signal_source"] = "feedback_axes.quality" + return "observed", quality, metadata + + rating = str(feedback.get("rating") or "").strip().lower() + scores = {"up": 1.0, "down": -1.0} + if rating in scores: + metadata["quality_signal_source"] = "rating" + metadata["raw_rating"] = rating + return "observed", scores[rating], metadata + return "not_observed", None, metadata + + +def _verify_decision_record_fingerprint( + *, + policy_name: str, + policy_version: str, + policy_schema_version: str, + decision_context: Any, + eligible_modes: Any, + action_probabilities: Mapping[str, float], + chosen_mode: str, + logging_support: Mapping[str, Any], + support_candidate_set_hash: Optional[str], + support_distribution_hash: Optional[str], +) -> Tuple[Optional[str], bool, str]: + """Fail closed when a durable decision no longer matches its fingerprint.""" + + raw_fingerprint = logging_support.get("decision_record_fingerprint") + stored_fingerprint = ( + str(raw_fingerprint).strip() + if isinstance(raw_fingerprint, str) and raw_fingerprint.strip() + else None + ) + if stored_fingerprint is None: + reason = ( + "legacy_unverifiable" + if logging_support.get("migration_source") == "ledger_schema_v1" + else "missing_unverifiable" + ) + return None, False, reason + + fingerprint_schema = logging_support.get( + "decision_record_fingerprint_schema_version" + ) + if fingerprint_schema != DECISION_FINGERPRINT_SCHEMA_VERSION: + return stored_fingerprint, False, "unsupported_schema" + if len(stored_fingerprint) != 64 or any( + character not in "0123456789abcdef" for character in stored_fingerprint + ): + return stored_fingerprint, False, "malformed_fingerprint" + if not isinstance(decision_context, Mapping): + return stored_fingerprint, False, "invalid_decision_context" + if not isinstance(eligible_modes, Sequence) or isinstance( + eligible_modes, (str, bytes) + ): + return stored_fingerprint, False, "invalid_eligible_modes" + + envelope_candidate_hash = logging_support.get("candidate_set_hash") + envelope_distribution_hash = logging_support.get("distribution_hash") + if ( + not isinstance(envelope_candidate_hash, str) + or not isinstance(envelope_distribution_hash, str) + or envelope_candidate_hash != support_candidate_set_hash + or envelope_distribution_hash != support_distribution_hash + ): + return stored_fingerprint, False, "support_projection_mismatch" + + try: + expected_fingerprint = _decision_record_fingerprint( + policy_name=policy_name, + policy_version=policy_version, + policy_schema_version=policy_schema_version, + decision_context=decision_context, + eligible_modes=eligible_modes, + action_probabilities=action_probabilities, + chosen_mode=chosen_mode, + candidate_set_hash=envelope_candidate_hash, + distribution_hash=envelope_distribution_hash, + ) + except (TypeError, ValueError): + return stored_fingerprint, False, "invalid_canonical_record" + if not hmac.compare_digest(stored_fingerprint, expected_fingerprint): + return stored_fingerprint, False, "fingerprint_mismatch" + return stored_fingerprint, True, "verified" + + +class RoutePolicyLedger: + """SQLite-backed route ledger using one durable transaction per phase.""" + + def __init__(self, db_path: Any, *, timeout_seconds: float = 30.0) -> None: + self.db_path = Path(db_path) + if str(self.db_path) == ":memory:": + raise ValueError("RoutePolicyLedger requires a durable filesystem path") + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self.timeout_seconds = max(1.0, float(timeout_seconds)) + self._initialize() + + def _connect(self) -> sqlite3.Connection: + connection = sqlite3.connect( + str(self.db_path), + timeout=self.timeout_seconds, + isolation_level=None, + ) + connection.row_factory = sqlite3.Row + connection.execute(f"PRAGMA busy_timeout = {int(self.timeout_seconds * 1000)}") + connection.execute("PRAGMA foreign_keys = ON") + connection.execute("PRAGMA synchronous = FULL") + return connection + + def _initialize(self) -> None: + connection = self._connect() + try: + connection.execute("PRAGMA journal_mode = WAL") + current_version = int(connection.execute("PRAGMA user_version").fetchone()[0]) + if current_version not in (0, 1, 2, LEDGER_SCHEMA_VERSION): + raise RoutePolicyLedgerError( + f"unsupported route ledger schema {current_version}; expected 1, 2, or {LEDGER_SCHEMA_VERSION}" + ) + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS ledger_metadata ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS session_counters ( + session_hash TEXT PRIMARY KEY, + next_sequence INTEGER NOT NULL CHECK (next_sequence >= 1) + ); + + CREATE TABLE IF NOT EXISTS route_decisions ( + route_id TEXT PRIMARY KEY, + session_hash TEXT NOT NULL, + session_sequence INTEGER NOT NULL CHECK (session_sequence >= 1), + ledger_schema_version INTEGER NOT NULL, + policy_name TEXT NOT NULL, + policy_version TEXT NOT NULL, + policy_schema_version TEXT NOT NULL, + decision_context_json TEXT NOT NULL, + eligible_modes_json TEXT NOT NULL, + action_probabilities_json TEXT NOT NULL, + chosen_mode TEXT NOT NULL, + executed_mode TEXT, + estimated_economics_json TEXT NOT NULL, + actual_economics_json TEXT, + status TEXT NOT NULL CHECK (status IN ('inflight', 'completed', 'failed')), + success INTEGER CHECK (success IS NULL OR success IN (0, 1)), + error_category TEXT, + error_message TEXT, + started_at REAL NOT NULL, + completed_at REAL, + UNIQUE (session_hash, session_sequence) + ); + + CREATE INDEX IF NOT EXISTS idx_route_decisions_session_sequence + ON route_decisions(session_hash, session_sequence DESC); + CREATE INDEX IF NOT EXISTS idx_route_decisions_status + ON route_decisions(status, started_at); + + CREATE TABLE IF NOT EXISTS route_feedback_revisions ( + route_id TEXT NOT NULL REFERENCES route_decisions(route_id) ON DELETE CASCADE, + revision INTEGER NOT NULL CHECK (revision >= 1), + idempotency_key TEXT NOT NULL, + feedback_json TEXT NOT NULL, + recorded_at REAL NOT NULL, + PRIMARY KEY (route_id, revision), + UNIQUE (route_id, idempotency_key) + ); + + CREATE INDEX IF NOT EXISTS idx_route_feedback_latest + ON route_feedback_revisions(route_id, revision DESC); + + CREATE TABLE IF NOT EXISTS route_decision_support ( + route_id TEXT PRIMARY KEY + REFERENCES route_decisions(route_id) ON DELETE CASCADE, + support_schema_version TEXT NOT NULL, + decision_type TEXT NOT NULL + CHECK (decision_type IN ('deterministic', 'randomized', 'legacy_unknown')), + probability_stage TEXT NOT NULL, + sampler_name TEXT NOT NULL, + sampler_version TEXT NOT NULL, + exploration_rate REAL NOT NULL + CHECK (exploration_rate >= 0.0 AND exploration_rate <= 1.0), + chosen_probability REAL + CHECK (chosen_probability IS NULL OR (chosen_probability > 0.0 AND chosen_probability <= 1.0)), + candidate_set_hash TEXT NOT NULL, + distribution_hash TEXT NOT NULL, + logging_envelope_json TEXT NOT NULL + ); + + CREATE INDEX IF NOT EXISTS idx_route_support_decision_type + ON route_decision_support(decision_type, probability_stage); + + CREATE TABLE IF NOT EXISTS route_outcome_contracts ( + route_id TEXT NOT NULL + REFERENCES route_decisions(route_id) ON DELETE CASCADE, + outcome_name TEXT NOT NULL + CHECK (outcome_name IN ('route_success', 'user_quality_rating', 'cost', 'latency')), + contract_schema_version TEXT NOT NULL, + outcome_definition_version TEXT NOT NULL, + observation_policy_id TEXT NOT NULL, + observation_policy_version TEXT NOT NULL, + value_type TEXT NOT NULL, + unit TEXT NOT NULL, + maturity_delay_seconds REAL NOT NULL + CHECK (maturity_delay_seconds >= 0.0), + maturity_basis TEXT NOT NULL + CHECK (maturity_basis = 'decision_started_at'), + precommitted INTEGER NOT NULL + CHECK (precommitted IN (0, 1)), + commitment_source TEXT NOT NULL, + contract_hash TEXT NOT NULL, + contract_json TEXT NOT NULL, + committed_at REAL NOT NULL, + PRIMARY KEY (route_id, outcome_name) + ); + + CREATE INDEX IF NOT EXISTS idx_route_outcome_contracts_name + ON route_outcome_contracts(outcome_name, precommitted); + + CREATE TABLE IF NOT EXISTS route_outcome_observation_events ( + event_id INTEGER PRIMARY KEY AUTOINCREMENT, + route_id TEXT NOT NULL, + outcome_name TEXT NOT NULL, + event_key TEXT NOT NULL, + observation_status TEXT NOT NULL + CHECK (observation_status IN ('observed', 'not_observed')), + value_json TEXT, + event_source TEXT NOT NULL, + metadata_json TEXT NOT NULL, + observed_at REAL NOT NULL, + recorded_at REAL NOT NULL, + FOREIGN KEY (route_id, outcome_name) + REFERENCES route_outcome_contracts(route_id, outcome_name) ON DELETE CASCADE, + UNIQUE (route_id, outcome_name, event_key), + CHECK ( + (observation_status = 'observed' AND value_json IS NOT NULL) + OR (observation_status = 'not_observed' AND value_json IS NULL) + ) + ); + + CREATE INDEX IF NOT EXISTS idx_route_outcome_events_route_time + ON route_outcome_observation_events(route_id, recorded_at, event_id); + CREATE INDEX IF NOT EXISTS idx_route_outcome_events_name_status + ON route_outcome_observation_events(outcome_name, observation_status, recorded_at); + + CREATE TRIGGER IF NOT EXISTS route_outcome_contracts_no_update + BEFORE UPDATE ON route_outcome_contracts + BEGIN + SELECT RAISE(ABORT, 'route outcome contracts are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS route_outcome_contracts_no_delete + BEFORE DELETE ON route_outcome_contracts + BEGIN + SELECT RAISE(ABORT, 'route outcome contracts are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS route_outcome_events_no_update + BEFORE UPDATE ON route_outcome_observation_events + BEGIN + SELECT RAISE(ABORT, 'route outcome observation events are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS route_outcome_events_no_delete + BEFORE DELETE ON route_outcome_observation_events + BEGIN + SELECT RAISE(ABORT, 'route outcome observation events are append-only'); + END; + """ + ) + self._backfill_support_rows(connection) + self._backfill_outcome_rows(connection) + connection.execute( + "INSERT OR REPLACE INTO ledger_metadata(key, value) VALUES ('schema_version', ?)", + (str(LEDGER_SCHEMA_VERSION),), + ) + connection.execute(f"PRAGMA user_version = {LEDGER_SCHEMA_VERSION}") + finally: + connection.close() + + @staticmethod + def _backfill_support_rows(connection: sqlite3.Connection) -> None: + """Idempotently migrate v1 decisions into an explicitly non-causal envelope.""" + + rows = connection.execute( + """ + SELECT d.route_id, d.eligible_modes_json, d.action_probabilities_json, d.chosen_mode + FROM route_decisions d + LEFT JOIN route_decision_support s ON s.route_id = d.route_id + WHERE s.route_id IS NULL + """ + ).fetchall() + for row in rows: + route_id = str(row["route_id"]) + chosen_mode = str(row["chosen_mode"]) + try: + modes = [str(mode) for mode in json.loads(str(row["eligible_modes_json"]))] + except (TypeError, ValueError, json.JSONDecodeError): + modes = [] + + probabilities = _legacy_probability_mapping(row["action_probabilities_json"]) + valid_vector = ( + bool(modes) + and bool(probabilities) + and set(probabilities) == set(modes) + and all(0.0 <= value <= 1.0 for value in probabilities.values()) + and math.isclose(sum(probabilities.values()), 1.0, rel_tol=0.0, abs_tol=1e-6) + and probabilities.get(chosen_mode, 0.0) > 0.0 + ) + positive_actions = ( + [mode for mode in modes if probabilities.get(mode, 0.0) > 0.0] + if valid_vector + else [] + ) + + deterministic = valid_vector and len(positive_actions) == 1 + decision_type = "deterministic" if deterministic else "legacy_unknown" + chosen_probability = probabilities.get(chosen_mode) if valid_vector else None + candidates = [{"action": mode, "legacy_migrated": True} for mode in modes] + sampler = { + "name": "legacy_v1", + "version": "1", + "exploration_rate": 0.0, + "assignment_unit": "route", + "assignment_commitment": None, + } + envelope = { + "schema_version": SUPPORT_SCHEMA_VERSION, + "decision_type": decision_type, + "probability_stage": "post_filter", + "sampler": sampler, + "candidates": candidates, + "exclusions": [], + "chosen_probability": chosen_probability, + "migration_source": "ledger_schema_v1", + "decision_record_fingerprint_schema_version": None, + "decision_record_fingerprint": None, + } + candidate_hash = _domain_hash( + _CANDIDATE_HASH_DOMAIN, + { + "schema_version": SUPPORT_SCHEMA_VERSION, + "candidates": candidates, + "exclusions": [], + }, + "legacy candidate support envelope", + ) + distribution_hash = _domain_hash( + _DISTRIBUTION_HASH_DOMAIN, + { + "schema_version": SUPPORT_SCHEMA_VERSION, + "decision_type": decision_type, + "probability_stage": "post_filter", + "sampler": sampler, + "eligible_modes": modes, + "action_probabilities": probabilities, + }, + "legacy logging distribution", + ) + envelope.update( + { + "candidate_set_hash": candidate_hash, + "distribution_hash": distribution_hash, + } + ) + connection.execute( + """ + INSERT OR IGNORE INTO route_decision_support( + route_id, support_schema_version, decision_type, probability_stage, + sampler_name, sampler_version, exploration_rate, chosen_probability, + candidate_set_hash, distribution_hash, logging_envelope_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + route_id, + SUPPORT_SCHEMA_VERSION, + decision_type, + "post_filter", + sampler["name"], + sampler["version"], + 0.0, + chosen_probability, + candidate_hash, + distribution_hash, + _canonical_json(envelope, "legacy logging envelope"), + ), + ) + + @staticmethod + def _insert_outcome_contract_rows( + connection: sqlite3.Connection, + route_id: str, + contracts: Mapping[str, Mapping[str, Any]], + *, + committed_at: float, + ) -> None: + for outcome_name in OUTCOME_NAMES: + contract = dict(contracts[outcome_name]) + connection.execute( + """ + INSERT OR IGNORE INTO route_outcome_contracts( + route_id, outcome_name, contract_schema_version, + outcome_definition_version, observation_policy_id, + observation_policy_version, value_type, unit, + maturity_delay_seconds, maturity_basis, precommitted, + commitment_source, contract_hash, contract_json, committed_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + route_id, + outcome_name, + contract["schema_version"], + contract["outcome_definition_version"], + contract["observation_policy_id"], + contract["observation_policy_version"], + contract["value_type"], + contract["unit"], + contract["maturity_delay_seconds"], + contract["maturity_basis"], + 1 if contract["precommitted"] else 0, + contract["commitment_source"], + contract["contract_hash"], + _canonical_json(contract, f"{outcome_name} outcome contract"), + float(committed_at), + ), + ) + + @staticmethod + def _insert_outcome_event( + connection: sqlite3.Connection, + *, + route_id: str, + outcome_name: str, + event_key: str, + observation_status: str, + value: Optional[Any], + event_source: str, + metadata: Optional[Mapping[str, Any]], + observed_at: float, + recorded_at: float, + ) -> None: + if outcome_name not in OUTCOME_NAMES: + raise ValueError("unsupported outcome_name") + if observation_status not in {"observed", "not_observed"}: + raise ValueError("observation_status must be observed or not_observed") + if observation_status == "observed" and value is None: + raise ValueError("observed outcome events require a value") + if observation_status == "not_observed" and value is not None: + raise ValueError("not_observed outcome events cannot carry a value") + value_json = ( + _canonical_json(value, f"{outcome_name} observed value") + if value is not None + else None + ) + metadata_json = _json_mapping(metadata, "outcome event metadata") + connection.execute( + """ + INSERT OR IGNORE INTO route_outcome_observation_events( + route_id, outcome_name, event_key, observation_status, + value_json, event_source, metadata_json, observed_at, recorded_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + route_id, + outcome_name, + _text(event_key, "event_key", limit=240), + observation_status, + value_json, + _text(event_source, "event_source", limit=80), + metadata_json, + float(observed_at), + float(recorded_at), + ), + ) + + @classmethod + def _backfill_outcome_rows(cls, connection: sqlite3.Connection) -> None: + """Idempotently mark old lifecycle data as descriptive, post-hoc evidence.""" + + migrated_at = float(time.time()) + decisions = connection.execute( + """ + SELECT d.* + FROM route_decisions d + LEFT JOIN route_outcome_contracts c ON c.route_id = d.route_id + GROUP BY d.route_id + HAVING COUNT(c.outcome_name) < 4 + """ + ).fetchall() + for row in decisions: + route_id = str(row["route_id"]) + contracts = build_route_outcome_contracts( + precommitted=False, + commitment_source="legacy_posthoc", + ) + cls._insert_outcome_contract_rows( + connection, + route_id, + contracts, + committed_at=migrated_at, + ) + + # Events are separately idempotent so a crash between contract and + # event backfill can be completed on the next open. + legacy_rows = connection.execute( + """ + SELECT d.* + FROM route_decisions d + JOIN route_outcome_contracts c ON c.route_id = d.route_id + WHERE c.outcome_name = 'route_success' + AND c.precommitted = 0 + AND d.status != 'inflight' + """ + ).fetchall() + for row in legacy_rows: + route_id = str(row["route_id"]) + observed_at = float(row["completed_at"] or row["started_at"]) + legacy_version = int(row["ledger_schema_version"] or 1) + metadata = { + "migration_source": f"ledger_schema_v{legacy_version}", + "descriptive_only": True, + } + success_value = ( + bool(row["success"]) + if row["success"] is not None + else None + ) + cls._insert_outcome_event( + connection, + route_id=route_id, + outcome_name="route_success", + event_key="legacy_completion:route_success", + observation_status="observed" if success_value is not None else "not_observed", + value=success_value, + event_source="legacy_posthoc", + metadata=metadata, + observed_at=observed_at, + recorded_at=migrated_at, + ) + try: + actual = json.loads(str(row["actual_economics_json"] or "{}")) + except (TypeError, ValueError, json.JSONDecodeError): + actual = {} + if not isinstance(actual, Mapping): + actual = {} + for outcome_name, metric_name in (("cost", "cost_units"), ("latency", "elapsed_ms")): + metric = _finite_actual_metric(actual, metric_name) + cls._insert_outcome_event( + connection, + route_id=route_id, + outcome_name=outcome_name, + event_key=f"legacy_completion:{outcome_name}", + observation_status="observed" if metric is not None else "not_observed", + value=metric, + event_source="legacy_posthoc", + metadata=metadata, + observed_at=observed_at, + recorded_at=migrated_at, + ) + + feedback_rows = connection.execute( + """ + SELECT f.route_id, f.revision, f.feedback_json, f.recorded_at, + d.ledger_schema_version + FROM route_feedback_revisions f + JOIN route_decisions d ON d.route_id = f.route_id + JOIN route_outcome_contracts c + ON c.route_id = f.route_id + AND c.outcome_name = 'user_quality_rating' + AND c.precommitted = 0 + ORDER BY f.route_id, f.revision + """ + ).fetchall() + for row in feedback_rows: + try: + feedback = json.loads(str(row["feedback_json"])) + except (TypeError, ValueError, json.JSONDecodeError): + feedback = {} + if not isinstance(feedback, Mapping): + feedback = {} + revision = int(row["revision"]) + status, value, metadata = _quality_observation(feedback, revision=revision) + metadata.update( + { + "migration_source": f"ledger_schema_v{int(row['ledger_schema_version'] or 1)}", + "descriptive_only": True, + } + ) + cls._insert_outcome_event( + connection, + route_id=str(row["route_id"]), + outcome_name="user_quality_rating", + event_key=f"legacy_feedback_revision:{revision}", + observation_status=status, + value=value, + event_source="legacy_posthoc", + metadata=metadata, + observed_at=float(row["recorded_at"]), + recorded_at=migrated_at, + ) + + @contextmanager + def _write_transaction(self) -> Iterator[sqlite3.Connection]: + connection = self._connect() + try: + connection.execute("BEGIN IMMEDIATE") + yield connection + connection.commit() + except BaseException: + connection.rollback() + raise + finally: + connection.close() + + @contextmanager + def _read_transaction(self) -> Iterator[sqlite3.Connection]: + """Keep a multi-query projection on one SQLite snapshot.""" + + connection = self._connect() + try: + connection.execute("BEGIN") + yield connection + connection.commit() + except BaseException: + connection.rollback() + raise + finally: + connection.close() + + def begin_decision( + self, + *, + session_id: str, + policy_name: str, + policy_version: str, + policy_schema_version: str, + decision_context: Mapping[str, Any], + eligible_modes: Sequence[str], + chosen_mode: str, + action_probabilities: Optional[Mapping[str, Any]] = None, + logging_support: Optional[Mapping[str, Any]] = None, + estimated_economics: Optional[Mapping[str, Any]] = None, + outcome_contracts: Optional[Mapping[str, Mapping[str, Any]]] = None, + route_id: Optional[Any] = None, + ) -> Dict[str, Any]: + """Commit an inflight decision before executing its chosen route.""" + + cooked_route_id = _route_id(route_id) + session_hash = hash_session_identity(session_id) + cooked_policy_name = _text(policy_name, "policy_name", limit=120) + cooked_policy_version = _text(policy_version, "policy_version", limit=120) + cooked_policy_schema = _text(policy_schema_version, "policy_schema_version", limit=120) + context_json = _json_mapping(decision_context, "decision_context") + modes = _eligible_modes(eligible_modes) + cooked_chosen = _text(chosen_mode, "chosen_mode", limit=80) + if cooked_chosen not in modes: + raise ValueError("chosen_mode must be present in eligible_modes") + probabilities = _probabilities(action_probabilities, modes, cooked_chosen) + support = _normalize_logging_support( + logging_support, + eligible_modes=modes, + probabilities=probabilities, + chosen_mode=cooked_chosen, + ) + support["decision_record_fingerprint_schema_version"] = ( + DECISION_FINGERPRINT_SCHEMA_VERSION + ) + support["decision_record_fingerprint"] = _decision_record_fingerprint( + policy_name=cooked_policy_name, + policy_version=cooked_policy_version, + policy_schema_version=cooked_policy_schema, + decision_context=json.loads(context_json), + eligible_modes=modes, + action_probabilities=probabilities, + chosen_mode=cooked_chosen, + candidate_set_hash=support["candidate_set_hash"], + distribution_hash=support["distribution_hash"], + ) + probabilities_json = json.dumps(probabilities, sort_keys=True, separators=(",", ":"), allow_nan=False) + support_json = _canonical_json(support, "logging_support") + modes_json = json.dumps(modes, separators=(",", ":"), ensure_ascii=False) + estimates_json = _json_mapping(estimated_economics, "estimated_economics") + contracts = build_route_outcome_contracts( + outcome_contracts, + precommitted=True, + ) + started_at = float(time.time()) + + try: + with self._write_transaction() as connection: + connection.execute( + "INSERT OR IGNORE INTO session_counters(session_hash, next_sequence) VALUES (?, 1)", + (session_hash,), + ) + sequence = int( + connection.execute( + "SELECT next_sequence FROM session_counters WHERE session_hash = ?", + (session_hash,), + ).fetchone()[0] + ) + connection.execute( + "UPDATE session_counters SET next_sequence = ? WHERE session_hash = ?", + (sequence + 1, session_hash), + ) + connection.execute( + """ + INSERT INTO route_decisions( + route_id, session_hash, session_sequence, ledger_schema_version, + policy_name, policy_version, policy_schema_version, + decision_context_json, eligible_modes_json, action_probabilities_json, + chosen_mode, estimated_economics_json, status, started_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'inflight', ?) + """, + ( + cooked_route_id, + session_hash, + sequence, + LEDGER_SCHEMA_VERSION, + cooked_policy_name, + cooked_policy_version, + cooked_policy_schema, + context_json, + modes_json, + probabilities_json, + cooked_chosen, + estimates_json, + started_at, + ), + ) + sampler = support["sampler"] + connection.execute( + """ + INSERT INTO route_decision_support( + route_id, support_schema_version, decision_type, probability_stage, + sampler_name, sampler_version, exploration_rate, chosen_probability, + candidate_set_hash, distribution_hash, logging_envelope_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + cooked_route_id, + support["schema_version"], + support["decision_type"], + support["probability_stage"], + sampler["name"], + sampler["version"], + sampler["exploration_rate"], + support["chosen_probability"], + support["candidate_set_hash"], + support["distribution_hash"], + support_json, + ), + ) + self._insert_outcome_contract_rows( + connection, + cooked_route_id, + contracts, + committed_at=started_at, + ) + except sqlite3.IntegrityError as exc: + raise LedgerConflictError(f"route_id already exists: {cooked_route_id}") from exc + return self.get_decision(cooked_route_id) + + def complete_decision( + self, + route_id: Any, + *, + success: bool, + executed_mode: Optional[str] = None, + actual_economics: Optional[Mapping[str, Any]] = None, + error_category: Optional[str] = None, + error_message: Optional[str] = None, + ) -> Dict[str, Any]: + """Complete an inflight decision, idempotently for an identical retry.""" + + cooked_route_id = _existing_route_id(route_id) + if not isinstance(success, bool): + raise ValueError("success must be a boolean") + cooked_executed = str(executed_mode or "").strip() or None + if cooked_executed is not None and len(cooked_executed) > 80: + raise ValueError("executed_mode must be at most 80 characters") + actual_json = _json_mapping(actual_economics, "actual_economics") + actual_mapping = json.loads(actual_json) + cooked_error_category = str(error_category or "").strip() or None + cooked_error_message = str(error_message or "").strip() or None + if success and (cooked_error_category or cooked_error_message): + raise ValueError("successful decisions cannot have error metadata") + if not success and not cooked_error_category: + raise ValueError("failed decisions require error_category") + if cooked_error_category is not None and len(cooked_error_category) > 120: + raise ValueError("error_category must be at most 120 characters") + if cooked_error_message is not None and len(cooked_error_message) > 2000: + raise ValueError("error_message must be at most 2000 characters") + terminal_status = "completed" if success else "failed" + completed_at = float(time.time()) + + with self._write_transaction() as connection: + existing = connection.execute( + "SELECT * FROM route_decisions WHERE route_id = ?", (cooked_route_id,) + ).fetchone() + if existing is None: + raise DecisionNotFoundError(f"unknown route_id: {cooked_route_id}") + if str(existing["status"]) != "inflight": + same_completion = ( + str(existing["status"]) == terminal_status + and bool(existing["success"]) is success + and existing["executed_mode"] == cooked_executed + and str(existing["actual_economics_json"] or "{}") == actual_json + and existing["error_category"] == cooked_error_category + and existing["error_message"] == cooked_error_message + ) + if not same_completion: + raise LedgerConflictError(f"route decision already completed differently: {cooked_route_id}") + else: + connection.execute( + """ + UPDATE route_decisions + SET status = ?, success = ?, executed_mode = ?, actual_economics_json = ?, + error_category = ?, error_message = ?, completed_at = ? + WHERE route_id = ? AND status = 'inflight' + """, + ( + terminal_status, + 1 if success else 0, + cooked_executed, + actual_json, + cooked_error_category, + cooked_error_message, + completed_at, + cooked_route_id, + ), + ) + completion_metadata = { + "executed_mode": cooked_executed, + "terminal_status": terminal_status, + "error_category": cooked_error_category, + } + self._insert_outcome_event( + connection, + route_id=cooked_route_id, + outcome_name="route_success", + event_key="completion", + observation_status="observed", + value=success, + event_source="route_completion", + metadata=completion_metadata, + observed_at=completed_at, + recorded_at=completed_at, + ) + for outcome_name, metric_name in (("cost", "cost_units"), ("latency", "elapsed_ms")): + metric = _finite_actual_metric(actual_mapping, metric_name) + self._insert_outcome_event( + connection, + route_id=cooked_route_id, + outcome_name=outcome_name, + event_key="completion", + observation_status="observed" if metric is not None else "not_observed", + value=metric, + event_source="route_completion", + metadata=completion_metadata, + observed_at=completed_at, + recorded_at=completed_at, + ) + return self.get_decision(cooked_route_id) + + def record_feedback( + self, + route_id: Any, + feedback: Mapping[str, Any], + *, + idempotency_key: Optional[str] = None, + ) -> Dict[str, Any]: + """Append a feedback revision or return the identical prior revision.""" + + cooked_route_id = _existing_route_id(route_id) + feedback_json = _json_mapping(feedback, "feedback") + explicit_key = idempotency_key is not None + cooked_key = _text(idempotency_key, "idempotency_key", limit=240) if explicit_key else None + recorded_at = float(time.time()) + idempotent = False + + with self._write_transaction() as connection: + exists = connection.execute( + "SELECT 1 FROM route_decisions WHERE route_id = ?", (cooked_route_id,) + ).fetchone() + if exists is None: + raise DecisionNotFoundError(f"unknown route_id: {cooked_route_id}") + if explicit_key: + prior = connection.execute( + """ + SELECT revision, idempotency_key, feedback_json, recorded_at + FROM route_feedback_revisions + WHERE route_id = ? AND idempotency_key = ? + """, + (cooked_route_id, cooked_key), + ).fetchone() + else: + # Content-derived compatibility is deliberately limited to an + # adjacent identical retry. Matching an older revision would + # suppress a legitimate up -> down -> up sequence. + prior = connection.execute( + """ + SELECT revision, idempotency_key, feedback_json, recorded_at + FROM route_feedback_revisions + WHERE route_id = ? ORDER BY revision DESC LIMIT 1 + """, + (cooked_route_id,), + ).fetchone() + if prior is not None and str(prior["feedback_json"]) != feedback_json: + prior = None + if prior is not None: + if explicit_key and str(prior["feedback_json"]) != feedback_json: + raise LedgerConflictError( + f"feedback idempotency key reused with different content for {cooked_route_id}" + ) + revision = int(prior["revision"]) + cooked_key = str(prior["idempotency_key"]) + recorded_at = float(prior["recorded_at"]) + idempotent = True + else: + prior_revision = connection.execute( + "SELECT COALESCE(MAX(revision), 0) FROM route_feedback_revisions WHERE route_id = ?", + (cooked_route_id,), + ).fetchone()[0] + revision = int(prior_revision) + 1 + if not explicit_key: + content_hash = hashlib.sha256(feedback_json.encode("utf-8")).hexdigest() + cooked_key = f"content-sha256:{content_hash}:revision:{revision}" + connection.execute( + """ + INSERT INTO route_feedback_revisions( + route_id, revision, idempotency_key, feedback_json, recorded_at + ) VALUES (?, ?, ?, ?, ?) + """, + (cooked_route_id, revision, cooked_key, feedback_json, recorded_at), + ) + feedback_mapping = json.loads(feedback_json) + observation_status, observation_value, observation_metadata = _quality_observation( + feedback_mapping, + revision=revision, + ) + self._insert_outcome_event( + connection, + route_id=cooked_route_id, + outcome_name="user_quality_rating", + event_key=f"feedback_revision:{revision}", + observation_status=observation_status, + value=observation_value, + event_source="feedback_revision", + metadata=observation_metadata, + observed_at=recorded_at, + recorded_at=recorded_at, + ) + return { + "route_id": cooked_route_id, + "revision": revision, + "idempotency_key": cooked_key, + "feedback": json.loads(feedback_json), + "recorded_at": recorded_at, + "recorded_at_utc": _utc(recorded_at), + "idempotent": idempotent, + } + + def feedback_history(self, route_id: Any) -> List[Dict[str, Any]]: + cooked_route_id = _existing_route_id(route_id) + connection = self._connect() + try: + exists = connection.execute( + "SELECT 1 FROM route_decisions WHERE route_id = ?", (cooked_route_id,) + ).fetchone() + if exists is None: + raise DecisionNotFoundError(f"unknown route_id: {cooked_route_id}") + rows = connection.execute( + """ + SELECT revision, idempotency_key, feedback_json, recorded_at + FROM route_feedback_revisions + WHERE route_id = ? ORDER BY revision ASC + """, + (cooked_route_id,), + ).fetchall() + return [self._feedback_row(cooked_route_id, row) for row in rows] + finally: + connection.close() + + @staticmethod + def _feedback_row(route_id: str, row: sqlite3.Row) -> Dict[str, Any]: + timestamp = float(row["recorded_at"]) + return { + "route_id": route_id, + "revision": int(row["revision"]), + "idempotency_key": str(row["idempotency_key"]), + "feedback": json.loads(str(row["feedback_json"])), + "recorded_at": timestamp, + "recorded_at_utc": _utc(timestamp), + } + + @staticmethod + def _outcome_contract_row(row: sqlite3.Row) -> Dict[str, Any]: + canonical: Dict[str, Any] = { + "schema_version": str(row["contract_schema_version"]), + "outcome_name": str(row["outcome_name"]), + "outcome_definition_version": str(row["outcome_definition_version"]), + "observation_policy_id": str(row["observation_policy_id"]), + "observation_policy_version": str(row["observation_policy_version"]), + "value_type": str(row["value_type"]), + "unit": str(row["unit"]), + "maturity_delay_seconds": float(row["maturity_delay_seconds"]), + "maturity_basis": str(row["maturity_basis"]), + "precommitted": bool(row["precommitted"]), + "commitment_source": str(row["commitment_source"]), + } + expected_hash = _domain_hash( + _OUTCOME_CONTRACT_HASH_DOMAIN, + canonical, + f"{canonical['outcome_name']} outcome contract", + ) + stored_hash = str(row["contract_hash"]) + if len(stored_hash) != 64 or any(character not in "0123456789abcdef" for character in stored_hash): + hash_valid = False + hash_reason = "malformed_hash" + elif not hmac.compare_digest(stored_hash, expected_hash): + hash_valid = False + hash_reason = "hash_mismatch" + else: + try: + raw = json.loads(str(row["contract_json"])) + except (TypeError, ValueError, json.JSONDecodeError): + raw = None + projected = dict(canonical) + projected["contract_hash"] = stored_hash + try: + projection_matches = bool( + isinstance(raw, Mapping) + and _canonical_json(dict(raw), "stored outcome contract") + == _canonical_json(projected, "projected outcome contract") + ) + except ValueError: + projection_matches = False + if not projection_matches: + hash_valid = False + hash_reason = "contract_projection_mismatch" + else: + hash_valid = True + hash_reason = "verified" + contract = dict(canonical) + contract["contract_hash"] = stored_hash + try: + committed_at = float(row["committed_at"]) + except (TypeError, ValueError, OverflowError): + committed_at = math.nan + contract.update( + { + "contract_hash_valid": hash_valid, + "contract_hash_reason": hash_reason, + "committed_at": committed_at, + "committed_at_utc": _utc(committed_at) if math.isfinite(committed_at) else None, + } + ) + return contract + + @staticmethod + def _outcome_event_row(row: sqlite3.Row) -> Dict[str, Any]: + observed_at = float(row["observed_at"]) + recorded_at = float(row["recorded_at"]) + return { + "event_id": int(row["event_id"]), + "route_id": str(row["route_id"]), + "outcome_name": str(row["outcome_name"]), + "event_key": str(row["event_key"]), + "observation_status": str(row["observation_status"]), + "value": ( + json.loads(str(row["value_json"])) + if row["value_json"] is not None + else None + ), + "event_source": str(row["event_source"]), + "metadata": json.loads(str(row["metadata_json"])), + "observed_at": observed_at, + "observed_at_utc": _utc(observed_at), + "recorded_at": recorded_at, + "recorded_at_utc": _utc(recorded_at), + } + + def _decision_from_row( + self, + connection: sqlite3.Connection, + row: sqlite3.Row, + *, + as_of: Optional[float] = None, + ) -> Dict[str, Any]: + route_id = str(row["route_id"]) + support_row = connection.execute( + "SELECT * FROM route_decision_support WHERE route_id = ?", (route_id,) + ).fetchone() + contract_cutoff = "AND committed_at <= ?" if as_of is not None else "" + contract_parameters: Tuple[Any, ...] = ( + (route_id, float(as_of)) if as_of is not None else (route_id,) + ) + contract_rows = connection.execute( + f""" + SELECT * FROM route_outcome_contracts + WHERE route_id = ? {contract_cutoff} ORDER BY outcome_name + """, + contract_parameters, + ).fetchall() + outcome_contracts = { + str(contract_row["outcome_name"]): self._outcome_contract_row(contract_row) + for contract_row in contract_rows + } + event_parameters: Tuple[Any, ...] = ( + (route_id, float(as_of)) if as_of is not None else (route_id,) + ) + event_cutoff = "AND recorded_at <= ?" if as_of is not None else "" + event_rows = connection.execute( + f""" + SELECT * FROM route_outcome_observation_events + WHERE route_id = ? {event_cutoff} + ORDER BY recorded_at, event_id + """, + event_parameters, + ).fetchall() + outcome_events = [self._outcome_event_row(event_row) for event_row in event_rows] + feedback_cutoff = "AND recorded_at <= ?" if as_of is not None else "" + feedback_parameters: Tuple[Any, ...] = ( + (route_id, float(as_of), route_id) + if as_of is not None + else (route_id, route_id) + ) + feedback_projection = connection.execute( + f""" + SELECT + f.revision, + f.idempotency_key, + f.feedback_json, + f.recorded_at, + summary.revision_count + FROM ( + SELECT COUNT(*) AS revision_count, MAX(revision) AS latest_revision + FROM route_feedback_revisions + WHERE route_id = ? {feedback_cutoff} + ) AS summary + LEFT JOIN route_feedback_revisions AS f + ON f.route_id = ? AND f.revision = summary.latest_revision + """, + feedback_parameters, + ).fetchone() + revision_count = int(feedback_projection["revision_count"]) + latest_feedback = ( + self._feedback_row(route_id, feedback_projection) + if feedback_projection["revision"] is not None + else None + ) + raw_logging_support = ( + json.loads(str(support_row["logging_envelope_json"])) + if support_row is not None + else None + ) + logging_support = ( + dict(raw_logging_support) + if isinstance(raw_logging_support, Mapping) + else { + "schema_version": None, + "decision_type": "legacy_unknown", + "probability_stage": None, + "sampler": {}, + "candidates": [], + "exclusions": [], + "chosen_probability": None, + "candidate_set_hash": None, + "distribution_hash": None, + "decision_record_fingerprint_schema_version": None, + "decision_record_fingerprint": None, + } + ) + decision_context = json.loads(str(row["decision_context_json"])) + eligible_modes = json.loads(str(row["eligible_modes_json"])) + action_probabilities = _legacy_probability_mapping(row["action_probabilities_json"]) + support_candidate_set_hash = ( + str(support_row["candidate_set_hash"]) + if support_row is not None + else None + ) + support_distribution_hash = ( + str(support_row["distribution_hash"]) + if support_row is not None + else None + ) + ( + decision_record_fingerprint, + decision_record_fingerprint_valid, + decision_record_fingerprint_reason, + ) = _verify_decision_record_fingerprint( + policy_name=str(row["policy_name"]), + policy_version=str(row["policy_version"]), + policy_schema_version=str(row["policy_schema_version"]), + decision_context=decision_context, + eligible_modes=eligible_modes, + action_probabilities=action_probabilities, + chosen_mode=str(row["chosen_mode"]), + logging_support=logging_support, + support_candidate_set_hash=support_candidate_set_hash, + support_distribution_hash=support_distribution_hash, + ) + started_at = float(row["started_at"]) + completed_at = float(row["completed_at"]) if row["completed_at"] is not None else None + historical_inflight = ( + as_of is not None + and completed_at is not None + and completed_at > float(as_of) + ) + projected_completed_at = None if historical_inflight else completed_at + projected_status = "inflight" if historical_inflight else str(row["status"]) + projected_success = ( + None + if historical_inflight or row["success"] is None + else bool(row["success"]) + ) + for contract in outcome_contracts.values(): + if not bool(contract.get("precommitted")): + contract["commitment_timing_valid"] = False + contract["commitment_timing_reason"] = "not_precommitted" + continue + try: + committed_at = float(contract.get("committed_at")) + except (TypeError, ValueError, OverflowError): + committed_at = math.nan + timing_valid = math.isfinite(committed_at) and committed_at <= started_at + contract["commitment_timing_valid"] = timing_valid + contract["commitment_timing_reason"] = ( + "verified" if timing_valid else "committed_after_decision_start" + ) + contract_sources = { + str(contract.get("commitment_source") or "") + for contract in outcome_contracts.values() + } + complete_contract_set = set(outcome_contracts) == set(OUTCOME_NAMES) + return { + "route_id": route_id, + "session_hash": str(row["session_hash"]), + "session_sequence": int(row["session_sequence"]), + "ledger_schema_version": int(row["ledger_schema_version"]), + "policy_name": str(row["policy_name"]), + "policy_version": str(row["policy_version"]), + "policy_schema_version": str(row["policy_schema_version"]), + "decision_context": decision_context, + "eligible_modes": eligible_modes, + "action_probabilities": action_probabilities, + "logging_support": logging_support, + "decision_type": logging_support.get("decision_type"), + "probability_stage": logging_support.get("probability_stage"), + "chosen_probability": logging_support.get("chosen_probability"), + "candidate_set_hash": logging_support.get("candidate_set_hash"), + "distribution_hash": logging_support.get("distribution_hash"), + "decision_record_fingerprint": decision_record_fingerprint, + "decision_record_fingerprint_valid": decision_record_fingerprint_valid, + "decision_record_fingerprint_reason": decision_record_fingerprint_reason, + "chosen_mode": str(row["chosen_mode"]), + "executed_mode": ( + str(row["executed_mode"]) + if not historical_inflight and row["executed_mode"] is not None + else None + ), + "estimated_economics": json.loads(str(row["estimated_economics_json"])), + "actual_economics": ( + json.loads(str(row["actual_economics_json"] or "{}")) + if not historical_inflight + else {} + ), + "outcome_contracts": outcome_contracts, + "outcome_events": outcome_events, + "outcome_contracts_precommitted_at_begin": ( + complete_contract_set + and all( + bool(contract.get("precommitted")) + and bool(contract.get("contract_hash_valid")) + and bool(contract.get("commitment_timing_valid")) + for contract in outcome_contracts.values() + ) + ), + "outcome_contracts_defaulted_at_begin": ( + complete_contract_set + and contract_sources == {"safe_default"} + and all( + bool(contract.get("contract_hash_valid")) + and bool(contract.get("commitment_timing_valid")) + for contract in outcome_contracts.values() + ) + ), + "outcome_contract_commitment_source": ( + next(iter(contract_sources)) if len(contract_sources) == 1 else "mixed" + ), + "status": projected_status, + "success": projected_success, + "error_category": ( + str(row["error_category"]) + if not historical_inflight and row["error_category"] is not None + else None + ), + "error_message": ( + str(row["error_message"]) + if not historical_inflight and row["error_message"] is not None + else None + ), + "started_at": started_at, + "started_at_utc": _utc(started_at), + "completed_at": projected_completed_at, + "completed_at_utc": _utc(projected_completed_at), + "duration_ms": ( + round((projected_completed_at - started_at) * 1000.0, 3) + if projected_completed_at is not None + else None + ), + "feedback_status": "known" if revision_count else "unknown", + "feedback_revision_count": revision_count, + "latest_feedback": latest_feedback, + } + + def get_decision(self, route_id: Any) -> Dict[str, Any]: + cooked_route_id = _existing_route_id(route_id) + with self._read_transaction() as connection: + row = connection.execute( + "SELECT * FROM route_decisions WHERE route_id = ?", (cooked_route_id,) + ).fetchone() + if row is None: + raise DecisionNotFoundError(f"unknown route_id: {cooked_route_id}") + return self._decision_from_row(connection, row) + + def list_decisions( + self, + *, + session_id: Optional[str] = None, + policy_name: Optional[str] = None, + policy_version: Optional[str] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + cooked_limit = max(1, min(1000, int(limit))) + session_hash = hash_session_identity(session_id) if session_id is not None else None + clauses: List[str] = [] + parameters: List[Any] = [] + if session_hash is not None: + clauses.append("session_hash = ?") + parameters.append(session_hash) + if policy_name is not None: + clauses.append("policy_name = ?") + parameters.append(_text(policy_name, "policy_name", limit=120)) + if policy_version is not None: + clauses.append("policy_version = ?") + parameters.append(_text(policy_version, "policy_version", limit=120)) + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + with self._read_transaction() as connection: + rows = connection.execute( + f""" + SELECT * FROM route_decisions + {where} + ORDER BY started_at DESC, route_id DESC LIMIT ? + """, + (*parameters, cooked_limit), + ).fetchall() + return [self._decision_from_row(connection, row) for row in rows] + + @staticmethod + def _outcome_contract_maturity( + decisions: Sequence[Mapping[str, Any]], + *, + as_of: float, + ) -> Dict[str, Any]: + by_outcome: Dict[str, Dict[str, Any]] = {} + for outcome_name in OUTCOME_NAMES: + contract_count = 0 + invalid_contract_count = 0 + late_commitment_count = 0 + precommitted_count = 0 + legacy_posthoc_count = 0 + mature_contract_count = 0 + pending_contract_count = 0 + recorded_event_count = 0 + observed_event_count = 0 + not_observed_event_count = 0 + routes_with_recorded_event: set[str] = set() + for decision in decisions: + route_id = str(decision.get("route_id") or "") + contracts = decision.get("outcome_contracts") + contract = contracts.get(outcome_name) if isinstance(contracts, Mapping) else None + if isinstance(contract, Mapping): + contract_count += 1 + if not bool(contract.get("contract_hash_valid")): + invalid_contract_count += 1 + continue + if bool(contract.get("precommitted")) and not bool( + contract.get("commitment_timing_valid") + ): + late_commitment_count += 1 + continue + if bool(contract.get("precommitted")): + precommitted_count += 1 + if contract.get("commitment_source") == "legacy_posthoc": + legacy_posthoc_count += 1 + maturity_at = float(decision.get("started_at") or 0.0) + float( + contract.get("maturity_delay_seconds") or 0.0 + ) + if maturity_at <= as_of: + mature_contract_count += 1 + else: + pending_contract_count += 1 + events = decision.get("outcome_events") + if not isinstance(events, Sequence) or isinstance(events, (str, bytes)): + continue + for event in events: + if not isinstance(event, Mapping) or event.get("outcome_name") != outcome_name: + continue + recorded_event_count += 1 + routes_with_recorded_event.add(route_id) + if event.get("observation_status") == "observed": + observed_event_count += 1 + elif event.get("observation_status") == "not_observed": + not_observed_event_count += 1 + by_outcome[outcome_name] = { + "contract_count": contract_count, + "invalid_contract_count": invalid_contract_count, + "late_commitment_count": late_commitment_count, + "precommitted_count": precommitted_count, + "legacy_posthoc_count": legacy_posthoc_count, + "mature_contract_count": mature_contract_count, + "pending_contract_count": pending_contract_count, + "recorded_event_count": recorded_event_count, + "observed_event_count": observed_event_count, + "not_observed_event_count": not_observed_event_count, + "routes_with_recorded_event": len(routes_with_recorded_event), + } + + complete_sets = [ + decision + for decision in decisions + if isinstance(decision.get("outcome_contracts"), Mapping) + and set(decision["outcome_contracts"]) == set(OUTCOME_NAMES) + and all( + bool(contract.get("contract_hash_valid")) + for contract in decision["outcome_contracts"].values() + if isinstance(contract, Mapping) + ) + and all( + isinstance(contract, Mapping) + for contract in decision["outcome_contracts"].values() + ) + ] + return { + "schema_version": OUTCOME_MATURITY_SCHEMA_VERSION, + "as_of": as_of, + "as_of_utc": _utc(as_of), + "basis": "decision_started_at_plus_maturity_delay_seconds", + "descriptive_only": True, + "policy_value_estimate": None, + "causal_identification": "not_performed", + "missingness_identification": "not_performed", + "absence_semantics": "no_event_recorded_is_not_classified", + "included_routes": len(decisions), + "complete_contract_sets": len(complete_sets), + "precommitted_routes": sum( + 1 + for decision in complete_sets + if bool(decision.get("outcome_contracts_precommitted_at_begin")) + ), + "legacy_posthoc_routes": sum( + 1 + for decision in complete_sets + if decision.get("outcome_contract_commitment_source") == "legacy_posthoc" + ), + "by_outcome": by_outcome, + } + + def policy_evidence_snapshot( + self, + *, + session_id: Optional[str] = None, + policy_name: Optional[str] = None, + policy_version: Optional[str] = None, + limit: int = 1000, + as_of: Optional[float] = None, + ) -> Dict[str, Any]: + """Return a prompt-free, ledger-native policy-evidence projection. + + The projection keeps lifecycle and missing-feedback state separate and + exposes the exact post-filter support envelope used at decision time. + It is intentionally read-only and contains no policy-value estimate. + """ + + if isinstance(as_of, bool): + raise ValueError("as_of must be a finite timestamp") + try: + as_of_value = float(time.time() if as_of is None else as_of) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("as_of must be a finite timestamp") from exc + if not math.isfinite(as_of_value): + raise ValueError("as_of must be a finite timestamp") + cooked_limit = max(1, min(1000, int(limit))) + session_hash = hash_session_identity(session_id) if session_id is not None else None + clauses = ["started_at <= ?"] + parameters: List[Any] = [as_of_value] + if session_hash is not None: + clauses.append("session_hash = ?") + parameters.append(session_hash) + cooked_policy_name = ( + _text(policy_name, "policy_name", limit=120) + if policy_name is not None + else None + ) + cooked_policy_version = ( + _text(policy_version, "policy_version", limit=120) + if policy_version is not None + else None + ) + if cooked_policy_name is not None: + clauses.append("policy_name = ?") + parameters.append(cooked_policy_name) + if cooked_policy_version is not None: + clauses.append("policy_version = ?") + parameters.append(cooked_policy_version) + where = f"WHERE {' AND '.join(clauses)}" + + # Lifecycle counts, decisions, contracts, events, and latest feedback + # are all read from one SQLite snapshot at one explicit cutoff. + with self._read_transaction() as connection: + status = connection.execute( + f""" + SELECT + COUNT(*) AS started, + COALESCE(SUM(CASE WHEN status = 'completed' AND completed_at <= ? THEN 1 ELSE 0 END), 0) AS completed, + COALESCE(SUM(CASE WHEN status = 'failed' AND completed_at <= ? THEN 1 ELSE 0 END), 0) AS failed + FROM route_decisions {where} + """, + (as_of_value, as_of_value, *parameters), + ).fetchone() + started = int(status["started"]) + completed = int(status["completed"]) + failed = int(status["failed"]) + terminal = completed + failed + feedback_clauses = [ + clause.replace("session_hash", "d.session_hash") + .replace("policy_name", "d.policy_name") + .replace("policy_version", "d.policy_version") + .replace("started_at", "d.started_at") + for clause in clauses + ] + feedback_where = f"WHERE {' AND '.join(feedback_clauses)}" + feedback = connection.execute( + f""" + SELECT + COUNT(f.revision) AS revision_count, + COUNT(DISTINCT f.route_id) AS known, + COUNT(DISTINCT CASE WHEN d.completed_at <= ? THEN f.route_id END) AS terminal_known + FROM route_decisions d + LEFT JOIN route_feedback_revisions f + ON f.route_id = d.route_id AND f.recorded_at <= ? + {feedback_where} + """, + (as_of_value, as_of_value, *parameters), + ).fetchone() + known = int(feedback["known"]) + terminal_known = int(feedback["terminal_known"]) + report = { + "ledger_schema_version": LEDGER_SCHEMA_VERSION, + "session_hash": session_hash, + "policy_filter": { + "policy_name": policy_name, + "policy_version": policy_version, + }, + "counts": { + "started": started, + "completed": completed, + "failed": failed, + "inflight": started - terminal, + }, + "feedback_coverage": { + "known": known, + "unknown": started - known, + "coverage_rate": round(known / started, 6) if started else 0.0, + "revision_count": int(feedback["revision_count"]), + "terminal_known": terminal_known, + "terminal_unknown": terminal - terminal_known, + "terminal_coverage_rate": ( + round(terminal_known / terminal, 6) if terminal else 0.0 + ), + "missing_feedback_semantics": "unknown", + }, + } + rows = connection.execute( + f""" + SELECT * FROM route_decisions + {where} + ORDER BY started_at DESC, route_id DESC LIMIT ? + """, + (*parameters, cooked_limit), + ).fetchall() + decisions = [ + self._decision_from_row(connection, row, as_of=as_of_value) + for row in rows + ] + decisions.sort(key=lambda row: (int(row.get("session_sequence") or 0), str(row.get("route_id") or ""))) + usage_rows: List[Dict[str, Any]] = [] + feedback_rows: List[Dict[str, Any]] = [] + expected_context_by_route_id: Dict[str, Any] = {} + + for decision in decisions: + route_id = str(decision["route_id"]) + context = _prompt_free_context(decision.get("decision_context") or {}) + if not isinstance(context, dict): + context = {} + support = dict(decision.get("logging_support") or {}) + eligible_modes = list(decision.get("eligible_modes") or []) + allowed_modes_raw = context.get("allowed_agent_modes") + allowed_modes = ( + [str(mode) for mode in allowed_modes_raw] + if isinstance(allowed_modes_raw, (list, tuple)) + else list(eligible_modes) + ) + auto_policy = { + "policy_id": decision.get("policy_name"), + "policy_version": decision.get("policy_version"), + "feature_schema_version": decision.get("policy_schema_version"), + "decision_type": support.get("decision_type"), + "probability_stage": support.get("probability_stage"), + "decision_context": context, + "action_mode": context.get("action_mode"), + "budget_profile": context.get("budget_profile"), + "score": context.get("score"), + "selected_agent_mode": decision.get("chosen_mode"), + "allowed_agent_modes": allowed_modes, + "eligible_actions": eligible_modes, + "action_probabilities": dict(decision.get("action_probabilities") or {}), + "post_filter_action_probabilities": dict(decision.get("action_probabilities") or {}), + "logging_propensity": support.get("chosen_probability"), + "candidate_set_hash": support.get("candidate_set_hash"), + "distribution_hash": support.get("distribution_hash"), + "decision_record_fingerprint": decision.get( + "decision_record_fingerprint" + ), + "decision_record_fingerprint_valid": bool( + decision.get("decision_record_fingerprint_valid") + ), + "decision_record_fingerprint_reason": decision.get( + "decision_record_fingerprint_reason" + ), + "outcome_contracts_precommitted_at_begin": bool( + decision.get("outcome_contracts_precommitted_at_begin") + ), + "logging_support": support, + } + usage_rows.append( + { + "route_id": route_id, + "session_hash": decision.get("session_hash"), + "session_sequence": decision.get("session_sequence"), + "selected_agent_mode": decision.get("chosen_mode"), + "executed_agent_mode": decision.get("executed_mode"), + "status": decision.get("status"), + "success": decision.get("success"), + "decision_record_fingerprint": decision.get( + "decision_record_fingerprint" + ), + "decision_record_fingerprint_valid": bool( + decision.get("decision_record_fingerprint_valid") + ), + "decision_record_fingerprint_reason": decision.get( + "decision_record_fingerprint_reason" + ), + "outcome_contracts": dict(decision.get("outcome_contracts") or {}), + "outcome_events": list(decision.get("outcome_events") or []), + "outcome_contracts_precommitted_at_begin": bool( + decision.get("outcome_contracts_precommitted_at_begin") + ), + "outcome_contracts_defaulted_at_begin": bool( + decision.get("outcome_contracts_defaulted_at_begin") + ), + "auto_agent_policy": auto_policy, + "route_economics": { + "estimate": dict(decision.get("estimated_economics") or {}), + "actual": dict(decision.get("actual_economics") or {}), + }, + } + ) + expected_context_by_route_id[route_id] = context + + latest = decision.get("latest_feedback") + if isinstance(latest, Mapping) and isinstance(latest.get("feedback"), Mapping): + feedback = dict(latest["feedback"]) + feedback.update( + { + "route_id": route_id, + "feedback_revision": latest.get("revision"), + "recorded_at": latest.get("recorded_at"), + "route_status": decision.get("status"), + "route_economics": { + "actual": dict(decision.get("actual_economics") or {}) + }, + } + ) + feedback_rows.append(feedback) + + started = int((report.get("counts") or {}).get("started") or 0) + return { + "ledger_schema_version": LEDGER_SCHEMA_VERSION, + "support_schema_version": SUPPORT_SCHEMA_VERSION, + "outcome_contract_schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + "as_of": as_of_value, + "as_of_utc": _utc(as_of_value), + "policy_filter": { + "policy_name": policy_name, + "policy_version": policy_version, + }, + "analysis_window": { + "limit": cooked_limit, + "included_decisions": len(decisions), + "total_decisions": started, + "truncated": started > len(decisions), + "order": "session_sequence_ascending", + "as_of": as_of_value, + "as_of_utc": _utc(as_of_value), + }, + "usage_rows": usage_rows, + "feedback_rows": feedback_rows, + "expected_context_by_route_id": expected_context_by_route_id, + "lifecycle": report, + "outcome_contract_maturity": self._outcome_contract_maturity( + decisions, + as_of=as_of_value, + ), + } + + def report( + self, + *, + session_id: Optional[str] = None, + policy_name: Optional[str] = None, + policy_version: Optional[str] = None, + ) -> Dict[str, Any]: + session_hash = hash_session_identity(session_id) if session_id is not None else None + clauses: List[str] = [] + parameters_list: List[Any] = [] + if session_hash is not None: + clauses.append("session_hash = ?") + parameters_list.append(session_hash) + if policy_name is not None: + clauses.append("policy_name = ?") + parameters_list.append(_text(policy_name, "policy_name", limit=120)) + if policy_version is not None: + clauses.append("policy_version = ?") + parameters_list.append(_text(policy_version, "policy_version", limit=120)) + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + parameters: Tuple[Any, ...] = tuple(parameters_list) + connection = self._connect() + try: + status = connection.execute( + f""" + SELECT + COUNT(*) AS started, + COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0) AS completed, + COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS failed, + COALESCE(SUM(CASE WHEN status = 'inflight' THEN 1 ELSE 0 END), 0) AS inflight + FROM route_decisions {where} + """, + parameters, + ).fetchone() + feedback_clauses = [clause.replace("session_hash", "d.session_hash").replace("policy_name", "d.policy_name").replace("policy_version", "d.policy_version") for clause in clauses] + feedback_where = f"WHERE {' AND '.join(feedback_clauses)}" if feedback_clauses else "" + feedback = connection.execute( + f""" + SELECT + COUNT(f.revision) AS revision_count, + COUNT(DISTINCT f.route_id) AS known, + COUNT(DISTINCT CASE WHEN d.status != 'inflight' THEN f.route_id END) AS terminal_known + FROM route_decisions d + LEFT JOIN route_feedback_revisions f ON f.route_id = d.route_id + {feedback_where} + """, + parameters, + ).fetchone() + finally: + connection.close() + + started = int(status["started"]) + completed = int(status["completed"]) + failed = int(status["failed"]) + inflight = int(status["inflight"]) + terminal = completed + failed + known = int(feedback["known"]) + terminal_known = int(feedback["terminal_known"]) + return { + "ledger_schema_version": LEDGER_SCHEMA_VERSION, + "session_hash": session_hash, + "policy_filter": { + "policy_name": policy_name, + "policy_version": policy_version, + }, + "counts": { + "started": started, + "completed": completed, + "failed": failed, + "inflight": inflight, + }, + "feedback_coverage": { + "known": known, + "unknown": started - known, + "coverage_rate": round(known / started, 6) if started else 0.0, + "revision_count": int(feedback["revision_count"]), + "terminal_known": terminal_known, + "terminal_unknown": terminal - terminal_known, + "terminal_coverage_rate": round(terminal_known / terminal, 6) if terminal else 0.0, + "missing_feedback_semantics": "unknown", + }, + } + + def snapshot(self, *, session_id: Optional[str] = None, limit: int = 50) -> Dict[str, Any]: + report = self.report(session_id=session_id) + report["recent_decisions"] = self.list_decisions(session_id=session_id, limit=limit) + return report + + +__all__ = [ + "DECISION_TYPES", + "DECISION_STATUSES", + "DECISION_FINGERPRINT_SCHEMA_VERSION", + "LEDGER_SCHEMA_VERSION", + "OUTCOME_CONTRACT_SCHEMA_VERSION", + "OUTCOME_MATURITY_SCHEMA_VERSION", + "OUTCOME_NAMES", + "SUPPORT_SCHEMA_VERSION", + "DecisionNotFoundError", + "LedgerConflictError", + "RoutePolicyLedger", + "RoutePolicyLedgerError", + "build_logging_support_envelope", + "build_route_outcome_contracts", + "hash_session_identity", +] diff --git a/source/route_policy_protocol.py b/source/route_policy_protocol.py new file mode 100644 index 000000000..dd552d530 --- /dev/null +++ b/source/route_policy_protocol.py @@ -0,0 +1,1541 @@ +"""Fail-closed stateful experiment protocol preflight for route rehearsals. + +The adjacent-route explorer describes one prompt-specific support stratum. +This module wraps one or more of those immutable plans in a prompt-free, +reviewable protocol *draft*. It freezes a target-policy class, population and +resource declarations, Route Outcome Contracts, and a stateful design screen. + +It deliberately does not seal a protocol, generate or reveal a random seed, +assign a cluster, execute a route, write evidence, estimate policy value, or +authorize promotion. Declarations are kept distinct from independent +validation so that state carryover, interference, and missingness cannot be +silently assumed away. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple + +try: + from .route_policy_explorer import ( + ACTIVATION_BLOCKERS, + STUDY_ID, + STUDY_PLAN_SCHEMA_VERSION, + STUDY_VERSION, + validate_adjacent_route_study, + ) + from .route_policy_lab import AUTO_AGENT_MODE_ORDER, get_policy_profile + from .route_policy_ledger import ( + OUTCOME_CONTRACT_SCHEMA_VERSION, + build_route_outcome_contracts, + ) +except ImportError: # pragma: no cover - direct source/ execution compatibility + from route_policy_explorer import ( + ACTIVATION_BLOCKERS, + STUDY_ID, + STUDY_PLAN_SCHEMA_VERSION, + STUDY_VERSION, + validate_adjacent_route_study, + ) + from route_policy_lab import AUTO_AGENT_MODE_ORDER, get_policy_profile + from route_policy_ledger import ( + OUTCOME_CONTRACT_SCHEMA_VERSION, + build_route_outcome_contracts, + ) + + +PROTOCOL_SCHEMA_VERSION = "route-study-protocol-preflight-v1" +PROTOCOL_VERSION = "1.0.0" +PROTOCOL_LABEL = "Stateful Route Experiment Preflight v1" +PROTOCOL_BUILD_INPUT_SCHEMA_VERSION = "route-study-protocol-build-input-v1" +REVIEW_BUNDLE_SCHEMA_VERSION = "route-study-review-bundle-v1" +REVIEW_BUNDLE_VERSION = "1.0.0" +REVIEW_BUNDLE_LABEL = "Route Protocol Review Bundle v1" +TARGET_POLICY_CLASS_SCHEMA_VERSION = "route-target-policy-class-v1" +POPULATION_SCHEMA_VERSION = "route-study-population-v1" +STATEFUL_DESIGN_SCHEMA_VERSION = "route-stateful-design-screen-v1" +STOPPING_SCHEMA_VERSION = "route-study-stopping-v1" +RANDOMNESS_SCHEMA_VERSION = "route-study-randomness-v1" + +DESIGN_MODES: Tuple[str, ...] = ( + "sticky_session_cluster", + "clustered_switchback", +) +CARRYOVER_SCOPES: Tuple[str, ...] = ( + "unknown", + "none_declared", + "within_session", + "cross_session", +) +INTERFERENCE_SCOPES: Tuple[str, ...] = ( + "unknown", + "none_declared", + "shared_resource", + "cross_cluster", +) +TEMPORAL_VARIATION_SCOPES: Tuple[str, ...] = ( + "unknown", + "stable_declared", + "nonstationary", +) + +DEFAULT_PLANNED_CLUSTERS = 200 +DEFAULT_MAX_ROUTES_PER_CLUSTER = 20 +DEFAULT_ANALYSIS_EVERY_CLUSTERS = 50 +MAX_PLANNED_CLUSTERS = 1_000_000 +MAX_ROUTES_PER_CLUSTER = 10_000 + +_PROTOCOL_HASH_DOMAIN = b"supermix.route-study.protocol-preflight.v1\x00" +_REVIEW_BUNDLE_HASH_DOMAIN = b"supermix.route-study.review-bundle.v1\x00" +_TARGET_CLASS_HASH_DOMAIN = b"supermix.route-study.target-class.v1\x00" +_OUTCOME_SET_HASH_DOMAIN = b"supermix.route-study.outcome-set.v1\x00" +_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,159}$") +_SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + +_TOP_LEVEL_KEYS = {"schema_version", "protocol", "charter", "protocol_hash"} +_PROTOCOL_KEYS = {"version", "label", "state", "activation_available"} +_CHARTER_KEYS = { + "source_studies", + "target_policy_class", + "population", + "stateful_design", + "outcome_observation", + "stopping_and_resources", + "randomness", + "external_evaluation", + "blocker_register", + "prompt_free_contract", + "causal_boundaries", +} +PROTOCOL_OPTION_KEYS = frozenset( + { + "target_policy_profile", + "design_mode", + "carryover_scope", + "interference_scope", + "temporal_variation", + "population_rule_id", + "population_rule_version", + "cluster_key_schema_version", + "planned_clusters", + "max_routes_per_cluster", + "analysis_every_clusters", + "block_length_routes", + "washout_routes", + "seed_commitment", + "external_estimator_id", + "external_reviewer_id", + } +) +PROTOCOL_BUILD_INPUT_KEYS = frozenset({"study_plans", *PROTOCOL_OPTION_KEYS}) + +_SOURCE_STUDIES_KEYS = { + "study_schema_version", + "study_id", + "study_version", + "support_stratum_count", + "common_source_contract", + "support_strata", + "source_plans_embedded", +} +_COMMON_SOURCE_CONTRACT_KEYS = { + "policy_id", + "policy_version", + "feature_schema_version", + "support_schema_version", + "outcome_contract_schema_version", +} +_SUPPORT_STRATUM_KEYS = { + "study_design_hash", + "candidate_set_hash", + "distribution_hash", + "baseline_action", + "eligible_actions", + "rehearsed_action_probabilities", + "executed_logging_propensities", +} +_POPULATION_KEYS = { + "schema_version", + "population_rule_id", + "population_rule_version", + "cluster_key_schema_version", + "cluster_identifier_exported", + "study_scoped_pseudonym_required", + "start_after_protocol_seal_required", + "one_study_policy_per_cluster", + "planned_clusters", + "max_routes_per_cluster", + "admitted_support_strata", + "population_precommitted", + "cluster_map_validated", +} +_STOPPING_KEYS = { + "schema_version", + "planned_cluster_ceiling", + "planned_route_ceiling", + "analysis_every_clusters", + "fixed_analysis_schedule", + "outcome_dependent_stopping_allowed", + "automatic_promotion_allowed", + "resource_guardrail_action", + "rules_precommitted", +} +_RANDOMNESS_KEYS = { + "schema_version", + "planned_derivation", + "seed_commitment", + "seed_material_included", + "seed_commitment_sealed", + "caller_selected_nonce_allowed", + "assignment_receipt_required_before_first_route", + "assignment_implementation_available", + "assignment_performed", + "nonce_grinding_resistant", +} +_OUTCOME_KEYS = { + "schema_version", + "contracts", + "observation_process_validated", + "missingness_identified", + "mnar_correction_available", + "rating_nonresponse_treated_as_outcome", + "contract_set_hash", +} +_EXTERNAL_KEYS = { + "estimator_id", + "reviewer_id", + "validated", + "causal_estimate_available", + "winner_available", +} +_PROMPT_FREE_KEYS = { + "prompt_free", + "raw_prompt_included", + "raw_session_id_included", + "free_form_text_fields_allowed", + "canonical_json", +} +_CAUSAL_BOUNDARY_KEYS = { + "deployment", + "protocol_sealed", + "activation_available", + "assignment_available", + "assignment_performed", + "execution_enabled", + "io_performed", + "ledger_write_performed", + "model_inference_performed", + "off_policy_estimate_computed", + "causal_identification_performed", + "power_analysis_performed", + "automatic_promotion_allowed", + "always_valid_inference_available", + "declarations_are_not_validation", + "activation_blockers", + "interpretation", +} +_DESIGN_KEYS = { + "schema_version", + "selected_design_mode", + "selected_design_status", + "assignment_unit", + "sticky_within_live_session", + "carryover_scope", + "interference_scope", + "temporal_variation", + "block_length_routes", + "washout_routes", + "candidate_screen", + "selected_design_blocking_reasons", + "carryover_validated", + "interference_validated", + "mixing_validated", + "causal_design_certified", +} +_DESIGN_CANDIDATE_KEYS = { + "design_mode", + "status", + "blocking_reasons", + "assignment_implementation_available", + "selected", +} +_BLOCKER_KEYS = {"code", "status", "activation_blocking"} +_BUNDLE_TOP_LEVEL_KEYS = { + "schema_version", + "bundle", + "protocol_builder", + "source_study_plans", + "protocol", + "bundle_hash", +} +_BUNDLE_META_KEYS = { + "version", + "label", + "state", + "verification_level", + "authenticity_proof_available", + "trusted_timestamp_available", + "activation_available", +} + + +def _canonical_json(value: Any) -> str: + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError("protocol must be canonical JSON without non-finite values") from exc + + +def _domain_hash(domain: bytes, value: Any) -> str: + return hashlib.sha256(domain + _canonical_json(value).encode("utf-8")).hexdigest() + + +def _identifier(value: Any, name: str) -> str: + cooked = str(value or "").strip() + if not _IDENTIFIER_RE.fullmatch(cooked): + raise ValueError(f"{name} must be a versioned prompt-free identifier") + return cooked + + +def _enum(value: Any, allowed: Sequence[str], name: str) -> str: + cooked = str(value or "").strip() + if cooked not in allowed: + raise ValueError(f"{name} must be one of: {', '.join(allowed)}") + return cooked + + +def _bounded_int(value: Any, name: str, *, minimum: int, maximum: int) -> int: + if isinstance(value, bool): + raise ValueError(f"{name} must be an integer between {minimum} and {maximum}") + try: + cooked = int(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"{name} must be an integer between {minimum} and {maximum}") from exc + if isinstance(value, float) and (not math.isfinite(value) or value != cooked): + raise ValueError(f"{name} must be an integer between {minimum} and {maximum}") + if not minimum <= cooked <= maximum: + raise ValueError(f"{name} must be an integer between {minimum} and {maximum}") + return cooked + + +def _seed_commitment(value: Any) -> Optional[str]: + if value is None or str(value).strip() == "": + return None + cooked = str(value).strip() + if not _SHA256_RE.fullmatch(cooked): + raise ValueError("seed_commitment must be a lowercase SHA-256 digest or omitted") + return cooked + + +def _normalize_study_plans(value: Any) -> list[Dict[str, Any]]: + if isinstance(value, Mapping): + raw_plans = [value] + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + raw_plans = list(value) + else: + raise ValueError("study_plans must be an eligible plan or a non-empty sequence") + if not raw_plans: + raise ValueError("study_plans must contain at least one eligible rehearsal plan") + if len(raw_plans) > 1_000: + raise ValueError("study_plans supports at most 1000 prompt-free support strata") + + validated = [validate_adjacent_route_study(plan) for plan in raw_plans] + hashes = [row["design_hash"] for row in validated] + if len(set(hashes)) != len(hashes): + raise ValueError("study_plans must not contain duplicate design hashes") + validated.sort(key=lambda row: row["design_hash"]) + + common_keys = ( + "policy_id", + "policy_version", + "feature_schema_version", + "support_schema_version", + "outcome_contract_schema_version", + ) + first_contract = validated[0]["source_contract"] + for row in validated[1:]: + if any(row["source_contract"][key] != first_contract[key] for key in common_keys): + raise ValueError("all study support strata must share one source policy/schema cohort") + return validated + + +def _target_policy_class(profile_name: str) -> Dict[str, Any]: + profile = get_policy_profile(profile_name) + payload = { + "schema_version": TARGET_POLICY_CLASS_SCHEMA_VERSION, + "profile_name": profile.name, + "thresholds": dict(profile.thresholds), + "supported_actions": list(AUTO_AGENT_MODE_ORDER), + "frozen_in_draft": True, + "externally_validated": False, + "optimality_claim": False, + } + return {**payload, "class_hash": _domain_hash(_TARGET_CLASS_HASH_DOMAIN, payload)} + + +def _design_reasons( + mode: str, + *, + carryover_scope: str, + interference_scope: str, + temporal_variation: str, + washout_routes: int, +) -> list[str]: + reasons: list[str] = [] + if carryover_scope == "unknown": + reasons.append("carryover_scope_unknown") + if interference_scope == "unknown": + reasons.append("interference_scope_unknown") + if temporal_variation == "unknown": + reasons.append("temporal_variation_unknown") + + if mode == "sticky_session_cluster": + if carryover_scope == "cross_session": + reasons.append("session_cluster_does_not_cover_cross_session_state") + if interference_scope in {"shared_resource", "cross_cluster"}: + reasons.append("session_cluster_does_not_block_declared_interference") + if temporal_variation == "nonstationary": + reasons.append("temporal_blocking_not_declared") + elif mode == "clustered_switchback": + if carryover_scope in {"within_session", "cross_session"} and washout_routes <= 0: + reasons.append("positive_washout_required_for_declared_carryover") + if carryover_scope == "cross_session": + reasons.append("finite_cross_session_carryover_not_validated") + if interference_scope in {"shared_resource", "cross_cluster"}: + reasons.append("interference_graph_and_exposure_mapping_not_validated") + if temporal_variation == "nonstationary": + reasons.append("block_balance_and_mixing_not_validated") + return reasons + + +def _design_screen( + selected_mode: str, + *, + carryover_scope: str, + interference_scope: str, + temporal_variation: str, + block_length_routes: int, + washout_routes: int, +) -> Dict[str, Any]: + candidates = [ + { + "design_mode": "route_randomization", + "status": "screened_out", + "blocking_reasons": [ + "stateful_product_campaigns_cannot_randomize_independent_routes_in_v1" + ], + "assignment_implementation_available": False, + } + ] + selected_reasons: list[str] = [] + for mode in DESIGN_MODES: + reasons = _design_reasons( + mode, + carryover_scope=carryover_scope, + interference_scope=interference_scope, + temporal_variation=temporal_variation, + washout_routes=washout_routes, + ) + if mode == selected_mode: + selected_reasons = list(reasons) + candidates.append( + { + "design_mode": mode, + "selected": mode == selected_mode, + "status": ( + "assumptions_declared_unvalidated" if not reasons else "blocked_for_review" + ), + "blocking_reasons": reasons, + "assignment_implementation_available": False, + } + ) + + has_unknown = any( + value == "unknown" + for value in (carryover_scope, interference_scope, temporal_variation) + ) + selected_status = ( + "declaration_incomplete" + if has_unknown + else ( + "assumptions_declared_unvalidated" + if not selected_reasons + else "incompatible_or_unvalidated_assumptions" + ) + ) + return { + "schema_version": STATEFUL_DESIGN_SCHEMA_VERSION, + "selected_design_mode": selected_mode, + "selected_design_status": selected_status, + "assignment_unit": ( + "session_hash" if selected_mode == "sticky_session_cluster" else "session_hash_x_time_block" + ), + "sticky_within_live_session": True, + "carryover_scope": carryover_scope, + "interference_scope": interference_scope, + "temporal_variation": temporal_variation, + "block_length_routes": block_length_routes if selected_mode == "clustered_switchback" else None, + "washout_routes": washout_routes if selected_mode == "clustered_switchback" else None, + "candidate_screen": candidates, + "selected_design_blocking_reasons": selected_reasons, + "carryover_validated": False, + "interference_validated": False, + "mixing_validated": False, + "causal_design_certified": False, + } + + +def _blocker_register( + *, + carryover_scope: str, + interference_scope: str, + seed_commitment: Optional[str], +) -> list[Dict[str, Any]]: + status_by_code = { + "target_policy_class_not_precommitted": "drafted_unsealed", + "outcome_observation_process_not_validated": "drafted_unvalidated", + "population_definition_not_precommitted": "drafted_unsealed", + "session_carryover_not_addressed": ( + "unresolved" if carryover_scope == "unknown" else "declared_unvalidated" + ), + "interference_not_addressed": ( + "unresolved" if interference_scope == "unknown" else "declared_unvalidated" + ), + "stopping_rules_not_precommitted": "drafted_unsealed", + "preassignment_seed_commitment_not_sealed": ( + "unresolved" if seed_commitment is None else "drafted_unsealed" + ), + "external_ope_not_validated": "unresolved", + } + return [ + { + "code": code, + "status": status_by_code[code], + "activation_blocking": True, + } + for code in ACTIVATION_BLOCKERS + ] + + +def _normalize_protocol_options( + *, + target_policy_profile: str = "balanced", + design_mode: str = "sticky_session_cluster", + carryover_scope: str = "unknown", + interference_scope: str = "unknown", + temporal_variation: str = "unknown", + population_rule_id: str = "interactive-auto-route-opt-in", + population_rule_version: str = "1", + cluster_key_schema_version: str = "session-hash-v1", + planned_clusters: int = DEFAULT_PLANNED_CLUSTERS, + max_routes_per_cluster: int = DEFAULT_MAX_ROUTES_PER_CLUSTER, + analysis_every_clusters: int = DEFAULT_ANALYSIS_EVERY_CLUSTERS, + block_length_routes: int = 20, + washout_routes: int = 0, + seed_commitment: Optional[str] = None, + external_estimator_id: Optional[str] = None, + external_reviewer_id: Optional[str] = None, +) -> Dict[str, Any]: + """Return the complete canonical builder option set used by review bundles.""" + + profile = get_policy_profile(str(target_policy_profile or "").strip()) + cooked_mode = _enum(design_mode, DESIGN_MODES, "design_mode") + cooked_carryover = _enum(carryover_scope, CARRYOVER_SCOPES, "carryover_scope") + cooked_interference = _enum( + interference_scope, INTERFERENCE_SCOPES, "interference_scope" + ) + cooked_temporal = _enum( + temporal_variation, TEMPORAL_VARIATION_SCOPES, "temporal_variation" + ) + cooked_clusters = _bounded_int( + planned_clusters, "planned_clusters", minimum=2, maximum=MAX_PLANNED_CLUSTERS + ) + cooked_max_routes = _bounded_int( + max_routes_per_cluster, + "max_routes_per_cluster", + minimum=1, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + cooked_analysis_every = _bounded_int( + analysis_every_clusters, + "analysis_every_clusters", + minimum=1, + maximum=cooked_clusters, + ) + cooked_block = _bounded_int( + block_length_routes, + "block_length_routes", + minimum=2, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + cooked_washout = _bounded_int( + washout_routes, + "washout_routes", + minimum=0, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + if cooked_washout >= cooked_block: + raise ValueError("washout_routes must be smaller than block_length_routes") + estimator_id = ( + _identifier(external_estimator_id, "external_estimator_id") + if external_estimator_id is not None and str(external_estimator_id).strip() + else None + ) + reviewer_id = ( + _identifier(external_reviewer_id, "external_reviewer_id") + if external_reviewer_id is not None and str(external_reviewer_id).strip() + else None + ) + return { + "target_policy_profile": profile.name, + "design_mode": cooked_mode, + "carryover_scope": cooked_carryover, + "interference_scope": cooked_interference, + "temporal_variation": cooked_temporal, + "population_rule_id": _identifier(population_rule_id, "population_rule_id"), + "population_rule_version": _identifier( + population_rule_version, "population_rule_version" + ), + "cluster_key_schema_version": _identifier( + cluster_key_schema_version, "cluster_key_schema_version" + ), + "planned_clusters": cooked_clusters, + "max_routes_per_cluster": cooked_max_routes, + "analysis_every_clusters": cooked_analysis_every, + "block_length_routes": cooked_block, + "washout_routes": cooked_washout, + "seed_commitment": _seed_commitment(seed_commitment), + "external_estimator_id": estimator_id, + "external_reviewer_id": reviewer_id, + } + + +def _split_protocol_build_input(payload: Any) -> Tuple[Any, Dict[str, Any]]: + if not isinstance(payload, Mapping): + raise ValueError("protocol build input must be a JSON object") + unknown = set(payload) - set(PROTOCOL_BUILD_INPUT_KEYS) + if unknown: + raise ValueError( + "protocol build input contains unsupported or non-prompt-free fields: " + + ", ".join(sorted(map(str, unknown))) + ) + if "study_plans" not in payload: + raise ValueError("protocol build input is missing required field: study_plans") + options = _normalize_protocol_options( + **{key: payload[key] for key in PROTOCOL_OPTION_KEYS if key in payload} + ) + return payload["study_plans"], options + + +def build_route_study_protocol_from_input(payload: Any) -> Dict[str, Any]: + """Build from the one closed, prompt-free input schema shared by web and CLI.""" + + study_plans, options = _split_protocol_build_input(payload) + return build_route_study_protocol(study_plans, **options) + + +def build_route_study_protocol( + study_plans: Any, + *, + target_policy_profile: str = "balanced", + design_mode: str = "sticky_session_cluster", + carryover_scope: str = "unknown", + interference_scope: str = "unknown", + temporal_variation: str = "unknown", + population_rule_id: str = "interactive-auto-route-opt-in", + population_rule_version: str = "1", + cluster_key_schema_version: str = "session-hash-v1", + planned_clusters: int = DEFAULT_PLANNED_CLUSTERS, + max_routes_per_cluster: int = DEFAULT_MAX_ROUTES_PER_CLUSTER, + analysis_every_clusters: int = DEFAULT_ANALYSIS_EVERY_CLUSTERS, + block_length_routes: int = 20, + washout_routes: int = 0, + seed_commitment: Optional[str] = None, + external_estimator_id: Optional[str] = None, + external_reviewer_id: Optional[str] = None, +) -> Dict[str, Any]: + """Build a canonical, non-sealed protocol preflight around study strata.""" + + validated = _normalize_study_plans(study_plans) + options = _normalize_protocol_options( + target_policy_profile=target_policy_profile, + design_mode=design_mode, + carryover_scope=carryover_scope, + interference_scope=interference_scope, + temporal_variation=temporal_variation, + population_rule_id=population_rule_id, + population_rule_version=population_rule_version, + cluster_key_schema_version=cluster_key_schema_version, + planned_clusters=planned_clusters, + max_routes_per_cluster=max_routes_per_cluster, + analysis_every_clusters=analysis_every_clusters, + block_length_routes=block_length_routes, + washout_routes=washout_routes, + seed_commitment=seed_commitment, + external_estimator_id=external_estimator_id, + external_reviewer_id=external_reviewer_id, + ) + + first_contract = validated[0]["source_contract"] + common_source_contract = { + key: first_contract[key] + for key in ( + "policy_id", + "policy_version", + "feature_schema_version", + "support_schema_version", + "outcome_contract_schema_version", + ) + } + strata = [ + { + "study_design_hash": row["design_hash"], + "candidate_set_hash": row["source_contract"]["candidate_set_hash"], + "distribution_hash": row["source_contract"]["distribution_hash"], + "baseline_action": row["enrollment"]["baseline_action"], + "eligible_actions": list(row["probability_design"]["eligible_actions"]), + "rehearsed_action_probabilities": dict( + row["probability_design"]["action_probabilities"] + ), + "executed_logging_propensities": False, + } + for row in validated + ] + + target_class = _target_policy_class(options["target_policy_profile"]) + outcome_contracts = build_route_outcome_contracts( + precommitted=True, + commitment_source="protocol_draft", + ) + outcome_payload = { + "schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + "contracts": outcome_contracts, + "observation_process_validated": False, + "missingness_identified": False, + "mnar_correction_available": False, + "rating_nonresponse_treated_as_outcome": False, + } + outcome_payload["contract_set_hash"] = _domain_hash( + _OUTCOME_SET_HASH_DOMAIN, outcome_payload + ) + + design = _design_screen( + options["design_mode"], + carryover_scope=options["carryover_scope"], + interference_scope=options["interference_scope"], + temporal_variation=options["temporal_variation"], + block_length_routes=options["block_length_routes"], + washout_routes=options["washout_routes"], + ) + population = { + "schema_version": POPULATION_SCHEMA_VERSION, + "population_rule_id": options["population_rule_id"], + "population_rule_version": options["population_rule_version"], + "cluster_key_schema_version": options["cluster_key_schema_version"], + "cluster_identifier_exported": False, + "study_scoped_pseudonym_required": True, + "start_after_protocol_seal_required": True, + "one_study_policy_per_cluster": True, + "planned_clusters": options["planned_clusters"], + "max_routes_per_cluster": options["max_routes_per_cluster"], + "admitted_support_strata": [row["study_design_hash"] for row in strata], + "population_precommitted": False, + "cluster_map_validated": False, + } + route_ceiling = options["planned_clusters"] * options["max_routes_per_cluster"] + stopping = { + "schema_version": STOPPING_SCHEMA_VERSION, + "planned_cluster_ceiling": options["planned_clusters"], + "planned_route_ceiling": route_ceiling, + "analysis_every_clusters": options["analysis_every_clusters"], + "fixed_analysis_schedule": True, + "outcome_dependent_stopping_allowed": False, + "automatic_promotion_allowed": False, + "resource_guardrail_action": "pause_and_independent_review", + "rules_precommitted": False, + } + randomness = { + "schema_version": RANDOMNESS_SCHEMA_VERSION, + "planned_derivation": "hmac_sha256_seed_design_hash_cluster_hash_block_id", + "seed_commitment": options["seed_commitment"], + "seed_material_included": False, + "seed_commitment_sealed": False, + "caller_selected_nonce_allowed": False, + "assignment_receipt_required_before_first_route": True, + "assignment_implementation_available": False, + "assignment_performed": False, + "nonce_grinding_resistant": False, + } + blocker_register = _blocker_register( + carryover_scope=options["carryover_scope"], + interference_scope=options["interference_scope"], + seed_commitment=options["seed_commitment"], + ) + charter = { + "source_studies": { + "study_schema_version": STUDY_PLAN_SCHEMA_VERSION, + "study_id": STUDY_ID, + "study_version": STUDY_VERSION, + "support_stratum_count": len(strata), + "common_source_contract": common_source_contract, + "support_strata": strata, + "source_plans_embedded": False, + }, + "target_policy_class": target_class, + "population": population, + "stateful_design": design, + "outcome_observation": outcome_payload, + "stopping_and_resources": stopping, + "randomness": randomness, + "external_evaluation": { + "estimator_id": options["external_estimator_id"], + "reviewer_id": options["external_reviewer_id"], + "validated": False, + "causal_estimate_available": False, + "winner_available": False, + }, + "blocker_register": blocker_register, + "prompt_free_contract": { + "prompt_free": True, + "raw_prompt_included": False, + "raw_session_id_included": False, + "free_form_text_fields_allowed": False, + "canonical_json": "sorted_keys_compact_utf8_no_nan", + }, + "causal_boundaries": { + "deployment": "shadow_only", + "protocol_sealed": False, + "activation_available": False, + "assignment_available": False, + "assignment_performed": False, + "execution_enabled": False, + "io_performed": False, + "ledger_write_performed": False, + "model_inference_performed": False, + "off_policy_estimate_computed": False, + "causal_identification_performed": False, + "power_analysis_performed": False, + "automatic_promotion_allowed": False, + "always_valid_inference_available": False, + "declarations_are_not_validation": True, + "activation_blockers": list(ACTIVATION_BLOCKERS), + "interpretation": "prompt-free protocol draft for independent review only", + }, + } + payload = { + "schema_version": PROTOCOL_SCHEMA_VERSION, + "protocol": { + "version": PROTOCOL_VERSION, + "label": PROTOCOL_LABEL, + "state": "draft_for_independent_review", + "activation_available": False, + }, + "charter": charter, + } + return {**payload, "protocol_hash": _domain_hash(_PROTOCOL_HASH_DOMAIN, payload)} + + +def _require_exact_keys(value: Any, expected: set[str], name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or set(value) != expected: + raise ValueError(f"{name} does not match the v1 schema") + return value + + +def _sha256_digest(value: Any, name: str) -> str: + cooked = str(value or "").strip() + if not _SHA256_RE.fullmatch(cooked): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + return cooked + + +def _optional_identifier(value: Any, name: str) -> Optional[str]: + if value is None: + return None + return _identifier(value, name) + + +def _validate_source_studies(value: Any) -> Tuple[Mapping[str, Any], list[str]]: + source = _require_exact_keys(value, _SOURCE_STUDIES_KEYS, "protocol source studies") + if source.get("study_schema_version") != STUDY_PLAN_SCHEMA_VERSION: + raise ValueError("protocol source study schema_version is unsupported") + if source.get("study_id") != STUDY_ID or source.get("study_version") != STUDY_VERSION: + raise ValueError("protocol source study identity is unsupported") + if source.get("source_plans_embedded") is not False: + raise ValueError("protocol draft cannot claim embedded source plans") + + common = _require_exact_keys( + source.get("common_source_contract"), + _COMMON_SOURCE_CONTRACT_KEYS, + "protocol common source contract", + ) + for key in ( + "policy_id", + "policy_version", + "feature_schema_version", + "support_schema_version", + "outcome_contract_schema_version", + ): + _identifier(common.get(key), f"common source contract {key}") + if common.get("outcome_contract_schema_version") != OUTCOME_CONTRACT_SCHEMA_VERSION: + raise ValueError("protocol source outcome contract schema is unsupported") + + raw_strata = source.get("support_strata") + if ( + not isinstance(raw_strata, Sequence) + or isinstance(raw_strata, (str, bytes)) + or not raw_strata + or len(raw_strata) > 1_000 + ): + raise ValueError("protocol must bind between 1 and 1000 support strata") + if source.get("support_stratum_count") != len(raw_strata): + raise ValueError("protocol support stratum count does not match its inventory") + + design_hashes: list[str] = [] + for index, raw in enumerate(raw_strata): + row = _require_exact_keys( + raw, _SUPPORT_STRATUM_KEYS, f"protocol support stratum {index}" + ) + design_hash = _sha256_digest(row.get("study_design_hash"), "study_design_hash") + _sha256_digest(row.get("candidate_set_hash"), "candidate_set_hash") + _sha256_digest(row.get("distribution_hash"), "distribution_hash") + eligible_raw = row.get("eligible_actions") + if not isinstance(eligible_raw, list) or not eligible_raw: + raise ValueError("protocol support stratum eligible_actions must be a non-empty list") + eligible = [str(action) for action in eligible_raw] + if len(set(eligible)) != len(eligible) or any( + action not in AUTO_AGENT_MODE_ORDER for action in eligible + ): + raise ValueError("protocol support stratum eligible_actions are unsupported") + canonical_eligible = [action for action in AUTO_AGENT_MODE_ORDER if action in eligible] + if eligible != canonical_eligible: + raise ValueError("protocol support stratum eligible_actions are not canonical") + if row.get("baseline_action") not in eligible: + raise ValueError("protocol support stratum baseline is outside eligible_actions") + probabilities = row.get("rehearsed_action_probabilities") + if not isinstance(probabilities, Mapping) or set(probabilities) != set(eligible): + raise ValueError("protocol support stratum probability vector is incomplete") + total = 0.0 + for action in eligible: + value = probabilities.get(action) + if isinstance(value, bool): + raise ValueError("protocol support probabilities must be finite positive numbers") + try: + probability = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError( + "protocol support probabilities must be finite positive numbers" + ) from exc + if not math.isfinite(probability) or probability <= 0.0 or probability > 1.0: + raise ValueError("protocol support probabilities must be finite positive numbers") + total += probability + if not math.isclose(total, 1.0, rel_tol=0.0, abs_tol=1e-9): + raise ValueError("protocol support probabilities must sum to one") + if row.get("executed_logging_propensities") is not False: + raise ValueError("protocol strata cannot claim executed logging propensities") + design_hashes.append(design_hash) + if len(set(design_hashes)) != len(design_hashes) or design_hashes != sorted(design_hashes): + raise ValueError("protocol support strata must be unique and canonically ordered") + return source, design_hashes + + +def _validate_design_screen(value: Any) -> Mapping[str, Any]: + design = _require_exact_keys(value, _DESIGN_KEYS, "protocol stateful design screen") + if design.get("schema_version") != STATEFUL_DESIGN_SCHEMA_VERSION: + raise ValueError("protocol stateful design schema_version is unsupported") + selected_mode = _enum( + design.get("selected_design_mode"), DESIGN_MODES, "selected_design_mode" + ) + carryover = _enum(design.get("carryover_scope"), CARRYOVER_SCOPES, "carryover_scope") + interference = _enum( + design.get("interference_scope"), INTERFERENCE_SCOPES, "interference_scope" + ) + temporal = _enum( + design.get("temporal_variation"), TEMPORAL_VARIATION_SCOPES, "temporal_variation" + ) + if selected_mode == "clustered_switchback": + block = _bounded_int( + design.get("block_length_routes"), + "block_length_routes", + minimum=2, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + washout = _bounded_int( + design.get("washout_routes"), + "washout_routes", + minimum=0, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + if washout >= block: + raise ValueError("protocol switchback washout must be smaller than block length") + expected = _design_screen( + selected_mode, + carryover_scope=carryover, + interference_scope=interference, + temporal_variation=temporal, + block_length_routes=block, + washout_routes=washout, + ) + if dict(design) != expected: + raise ValueError("protocol stateful design is not a canonical builder result") + return design + + if design.get("block_length_routes") is not None or design.get("washout_routes") is not None: + raise ValueError("sticky-session design cannot expose switchback controls") + candidates = design.get("candidate_screen") + if not isinstance(candidates, list) or len(candidates) != 3: + raise ValueError("protocol design candidate screen is incomplete") + expected_route = { + "design_mode": "route_randomization", + "status": "screened_out", + "blocking_reasons": [ + "stateful_product_campaigns_cannot_randomize_independent_routes_in_v1" + ], + "assignment_implementation_available": False, + } + if candidates[0] != expected_route: + raise ValueError("protocol must keep route-level randomization screened out") + for index, mode in enumerate(DESIGN_MODES, start=1): + row = _require_exact_keys( + candidates[index], _DESIGN_CANDIDATE_KEYS, f"protocol design candidate {mode}" + ) + if row.get("design_mode") != mode or row.get("selected") is not (mode == selected_mode): + raise ValueError("protocol design candidate selection is inconsistent") + possible_reasons = [ + _design_reasons( + mode, + carryover_scope=carryover, + interference_scope=interference, + temporal_variation=temporal, + washout_routes=washout, + ) + for washout in (0, 1) + ] + if row.get("blocking_reasons") not in possible_reasons: + raise ValueError("protocol design candidate reasons are not canonical") + expected_status = ( + "assumptions_declared_unvalidated" + if not row.get("blocking_reasons") + else "blocked_for_review" + ) + if row.get("status") != expected_status or row.get( + "assignment_implementation_available" + ) is not False: + raise ValueError("protocol design candidate status is not canonical") + selected_reasons = _design_reasons( + selected_mode, + carryover_scope=carryover, + interference_scope=interference, + temporal_variation=temporal, + washout_routes=0, + ) + has_unknown = "unknown" in (carryover, interference, temporal) + selected_status = ( + "declaration_incomplete" + if has_unknown + else ( + "assumptions_declared_unvalidated" + if not selected_reasons + else "incompatible_or_unvalidated_assumptions" + ) + ) + expected_fixed = { + "selected_design_status": selected_status, + "assignment_unit": "session_hash", + "sticky_within_live_session": True, + "selected_design_blocking_reasons": selected_reasons, + "carryover_validated": False, + "interference_validated": False, + "mixing_validated": False, + "causal_design_certified": False, + } + if any(design.get(key) != expected for key, expected in expected_fixed.items()): + raise ValueError("protocol sticky-session design claims unsupported semantics") + return design + + +def _validate_protocol_semantics(charter: Mapping[str, Any]) -> Tuple[Mapping[str, Any], int]: + source, design_hashes = _validate_source_studies(charter.get("source_studies")) + + target_class = charter.get("target_policy_class") + if not isinstance(target_class, Mapping): + raise ValueError("protocol target-policy class is missing") + expected_target = _target_policy_class(str(target_class.get("profile_name") or "")) + if dict(target_class) != expected_target: + raise ValueError("protocol target-policy class does not match a frozen profile") + + design = _validate_design_screen(charter.get("stateful_design")) + population = _require_exact_keys( + charter.get("population"), _POPULATION_KEYS, "protocol population" + ) + if population.get("schema_version") != POPULATION_SCHEMA_VERSION: + raise ValueError("protocol population schema_version is unsupported") + _identifier(population.get("population_rule_id"), "population_rule_id") + _identifier(population.get("population_rule_version"), "population_rule_version") + _identifier(population.get("cluster_key_schema_version"), "cluster_key_schema_version") + planned_clusters = _bounded_int( + population.get("planned_clusters"), + "planned_clusters", + minimum=2, + maximum=MAX_PLANNED_CLUSTERS, + ) + max_routes = _bounded_int( + population.get("max_routes_per_cluster"), + "max_routes_per_cluster", + minimum=1, + maximum=MAX_ROUTES_PER_CLUSTER, + ) + expected_population_flags = { + "cluster_identifier_exported": False, + "study_scoped_pseudonym_required": True, + "start_after_protocol_seal_required": True, + "one_study_policy_per_cluster": True, + "population_precommitted": False, + "cluster_map_validated": False, + } + if any(population.get(key) is not expected for key, expected in expected_population_flags.items()): + raise ValueError("protocol population must remain draft-only and unvalidated") + if population.get("admitted_support_strata") != design_hashes: + raise ValueError("protocol population does not bind the canonical support strata") + + stopping = _require_exact_keys( + charter.get("stopping_and_resources"), _STOPPING_KEYS, "protocol stopping rules" + ) + if stopping.get("schema_version") != STOPPING_SCHEMA_VERSION: + raise ValueError("protocol stopping schema_version is unsupported") + analysis_every = _bounded_int( + stopping.get("analysis_every_clusters"), + "analysis_every_clusters", + minimum=1, + maximum=planned_clusters, + ) + expected_stopping = { + "schema_version": STOPPING_SCHEMA_VERSION, + "planned_cluster_ceiling": planned_clusters, + "planned_route_ceiling": planned_clusters * max_routes, + "analysis_every_clusters": analysis_every, + "fixed_analysis_schedule": True, + "outcome_dependent_stopping_allowed": False, + "automatic_promotion_allowed": False, + "resource_guardrail_action": "pause_and_independent_review", + "rules_precommitted": False, + } + if dict(stopping) != expected_stopping: + raise ValueError("protocol stopping rules are not canonical or fail-closed") + + randomness = _require_exact_keys( + charter.get("randomness"), _RANDOMNESS_KEYS, "protocol randomness contract" + ) + expected_randomness = { + "schema_version": RANDOMNESS_SCHEMA_VERSION, + "planned_derivation": "hmac_sha256_seed_design_hash_cluster_hash_block_id", + "seed_commitment": _seed_commitment(randomness.get("seed_commitment")), + "seed_material_included": False, + "seed_commitment_sealed": False, + "caller_selected_nonce_allowed": False, + "assignment_receipt_required_before_first_route": True, + "assignment_implementation_available": False, + "assignment_performed": False, + "nonce_grinding_resistant": False, + } + if dict(randomness) != expected_randomness: + raise ValueError("protocol randomness contract must remain non-assigning and unsealed") + + outcome = _require_exact_keys( + charter.get("outcome_observation"), _OUTCOME_KEYS, "protocol outcome observation" + ) + expected_outcome = { + "schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + "contracts": build_route_outcome_contracts( + precommitted=True, commitment_source="protocol_draft" + ), + "observation_process_validated": False, + "missingness_identified": False, + "mnar_correction_available": False, + "rating_nonresponse_treated_as_outcome": False, + } + expected_outcome["contract_set_hash"] = _domain_hash( + _OUTCOME_SET_HASH_DOMAIN, expected_outcome + ) + if dict(outcome) != expected_outcome: + raise ValueError("protocol outcome contracts are not the frozen v1 definitions") + + external = _require_exact_keys( + charter.get("external_evaluation"), _EXTERNAL_KEYS, "protocol external evaluation" + ) + expected_external = { + "estimator_id": _optional_identifier(external.get("estimator_id"), "estimator_id"), + "reviewer_id": _optional_identifier(external.get("reviewer_id"), "reviewer_id"), + "validated": False, + "causal_estimate_available": False, + "winner_available": False, + } + if dict(external) != expected_external: + raise ValueError("protocol external evaluation must remain unvalidated") + + blocker_register = charter.get("blocker_register") + if not isinstance(blocker_register, list) or any( + not isinstance(row, Mapping) or set(row) != _BLOCKER_KEYS for row in blocker_register + ): + raise ValueError("protocol blocker register does not match the v1 schema") + expected_blockers = _blocker_register( + carryover_scope=str(design["carryover_scope"]), + interference_scope=str(design["interference_scope"]), + seed_commitment=randomness["seed_commitment"], + ) + if blocker_register != expected_blockers: + raise ValueError("protocol blocker register is not canonical") + + prompt_free = _require_exact_keys( + charter.get("prompt_free_contract"), _PROMPT_FREE_KEYS, "protocol prompt-free contract" + ) + expected_prompt_free = { + "prompt_free": True, + "raw_prompt_included": False, + "raw_session_id_included": False, + "free_form_text_fields_allowed": False, + "canonical_json": "sorted_keys_compact_utf8_no_nan", + } + if dict(prompt_free) != expected_prompt_free: + raise ValueError("protocol prompt-free contract is not canonical") + + boundaries = _require_exact_keys( + charter.get("causal_boundaries"), _CAUSAL_BOUNDARY_KEYS, "protocol causal boundaries" + ) + expected_boundaries = { + "deployment": "shadow_only", + "protocol_sealed": False, + "activation_available": False, + "assignment_available": False, + "assignment_performed": False, + "execution_enabled": False, + "io_performed": False, + "ledger_write_performed": False, + "model_inference_performed": False, + "off_policy_estimate_computed": False, + "causal_identification_performed": False, + "power_analysis_performed": False, + "automatic_promotion_allowed": False, + "always_valid_inference_available": False, + "declarations_are_not_validation": True, + "activation_blockers": list(ACTIVATION_BLOCKERS), + "interpretation": "prompt-free protocol draft for independent review only", + } + if dict(boundaries) != expected_boundaries: + raise ValueError("protocol causal boundaries must remain exactly fail-closed") + return design, len(design_hashes) + + +def audit_route_study_protocol(protocol: Any) -> Dict[str, Any]: + """Verify the canonical hash and fail-closed boundaries of a draft.""" + + if not isinstance(protocol, Mapping): + raise ValueError("protocol must be a JSON object") + if set(protocol) != _TOP_LEVEL_KEYS: + raise ValueError("protocol top-level fields do not match the v1 schema") + if protocol.get("schema_version") != PROTOCOL_SCHEMA_VERSION: + raise ValueError("unsupported route study protocol schema_version") + protocol_meta = protocol.get("protocol") + charter = protocol.get("charter") + if not isinstance(protocol_meta, Mapping) or set(protocol_meta) != _PROTOCOL_KEYS: + raise ValueError("protocol metadata does not match the v1 schema") + if not isinstance(charter, Mapping) or set(charter) != _CHARTER_KEYS: + raise ValueError("protocol charter does not match the v1 schema") + expected_protocol_meta = { + "version": PROTOCOL_VERSION, + "label": PROTOCOL_LABEL, + "state": "draft_for_independent_review", + "activation_available": False, + } + if dict(protocol_meta) != expected_protocol_meta: + raise ValueError("protocol metadata is not the canonical fail-closed v1 draft") + semantic_design, semantic_stratum_count = _validate_protocol_semantics(charter) + boundaries = charter.get("causal_boundaries") + if not isinstance(boundaries, Mapping): + raise ValueError("protocol causal boundaries are missing") + forbidden_true = ( + "protocol_sealed", + "activation_available", + "assignment_available", + "assignment_performed", + "execution_enabled", + "io_performed", + "ledger_write_performed", + "model_inference_performed", + "off_policy_estimate_computed", + "causal_identification_performed", + "power_analysis_performed", + "automatic_promotion_allowed", + "always_valid_inference_available", + ) + if any(boundaries.get(key) is not False for key in forbidden_true): + raise ValueError("protocol causal boundaries must remain fail-closed") + if boundaries.get("activation_blockers") != list(ACTIVATION_BLOCKERS): + raise ValueError("protocol must retain every activation blocker") + + blocker_register = charter.get("blocker_register") + if not isinstance(blocker_register, Sequence) or isinstance( + blocker_register, (str, bytes) + ): + raise ValueError("protocol blocker register is missing") + blocker_codes = [ + str(row.get("code") or "") if isinstance(row, Mapping) else "" + for row in blocker_register + ] + if blocker_codes != list(ACTIVATION_BLOCKERS) or any( + not isinstance(row, Mapping) or row.get("activation_blocking") is not True + for row in blocker_register + ): + raise ValueError("protocol blocker register must retain every blocking gate") + + target_class = charter.get("target_policy_class") + if not isinstance(target_class, Mapping): + raise ValueError("protocol target-policy class is missing") + expected_target = _target_policy_class(str(target_class.get("profile_name") or "")) + if dict(target_class) != expected_target: + raise ValueError("protocol target-policy class does not match a frozen profile") + + design = charter.get("stateful_design") + if not isinstance(design, Mapping): + raise ValueError("protocol stateful design screen is missing") + if design.get("selected_design_mode") not in DESIGN_MODES: + raise ValueError("protocol selected design mode is unsupported") + if ( + design.get("assignment_implementation_available") is True + or design.get("causal_design_certified") is not False + or design.get("carryover_validated") is not False + or design.get("interference_validated") is not False + or design.get("mixing_validated") is not False + ): + raise ValueError("protocol design declarations cannot claim validation") + + randomness = charter.get("randomness") + if not isinstance(randomness, Mapping) or any( + randomness.get(key) is not False + for key in ( + "seed_material_included", + "seed_commitment_sealed", + "caller_selected_nonce_allowed", + "assignment_implementation_available", + "assignment_performed", + "nonce_grinding_resistant", + ) + ): + raise ValueError("protocol randomness contract must remain non-assigning and unsealed") + external = charter.get("external_evaluation") + if not isinstance(external, Mapping) or any( + external.get(key) is not False + for key in ("validated", "causal_estimate_available", "winner_available") + ): + raise ValueError("protocol external evaluation must remain unvalidated") + outcome = charter.get("outcome_observation") + if not isinstance(outcome, Mapping) or any( + outcome.get(key) is not False + for key in ( + "observation_process_validated", + "missingness_identified", + "mnar_correction_available", + "rating_nonresponse_treated_as_outcome", + ) + ): + raise ValueError("protocol outcome observation must remain unvalidated") + + source_studies = charter.get("source_studies") + strata = source_studies.get("support_strata") if isinstance(source_studies, Mapping) else None + if not isinstance(strata, Sequence) or isinstance(strata, (str, bytes)) or not strata: + raise ValueError("protocol must bind at least one support stratum") + if any( + not isinstance(row, Mapping) or row.get("executed_logging_propensities") is not False + for row in strata + ): + raise ValueError("protocol strata cannot claim executed logging propensities") + + supplied_hash = str(protocol.get("protocol_hash") or "").strip() + payload = { + "schema_version": protocol["schema_version"], + "protocol": dict(protocol_meta), + "charter": dict(charter), + } + expected_hash = _domain_hash(_PROTOCOL_HASH_DOMAIN, payload) + if supplied_hash != expected_hash: + raise ValueError("protocol_hash does not match the canonical protocol draft") + return { + "ok": True, + "schema_version": PROTOCOL_SCHEMA_VERSION, + "protocol_hash": supplied_hash, + "state": "draft_for_independent_review", + "support_stratum_count": semantic_stratum_count, + "selected_design_mode": str(semantic_design["selected_design_mode"]), + "selected_design_status": str(semantic_design["selected_design_status"]), + "verification_level": "structural_without_source_plans", + "source_plan_reconstruction_performed": False, + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "activation_available": False, + "activation_blockers": list(ACTIVATION_BLOCKERS), + } + + +def build_route_study_review_bundle( + study_plans: Any, + **protocol_options: Any, +) -> Dict[str, Any]: + """Create a portable, prompt-free bundle that supports full reconstruction. + + The bundle is still an unsigned, untimestamped review artifact. It binds + canonical source plans to the exact protocol builder options and result, but + it does not prove authorship, registration time, causal validity, or safety + to activate. + """ + + unknown = set(protocol_options) - set(PROTOCOL_OPTION_KEYS) + if unknown: + raise ValueError( + "review bundle contains unsupported protocol options: " + + ", ".join(sorted(map(str, unknown))) + ) + validated = _normalize_study_plans(study_plans) + canonical_plans = [row["plan"] for row in validated] + options = _normalize_protocol_options(**protocol_options) + protocol = build_route_study_protocol(canonical_plans, **options) + payload = { + "schema_version": REVIEW_BUNDLE_SCHEMA_VERSION, + "bundle": { + "version": REVIEW_BUNDLE_VERSION, + "label": REVIEW_BUNDLE_LABEL, + "state": "ready_for_semantic_review", + "verification_level": "full_source_bound_reconstruction", + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "activation_available": False, + }, + "protocol_builder": { + "schema_version": PROTOCOL_BUILD_INPUT_SCHEMA_VERSION, + "options": options, + }, + "source_study_plans": canonical_plans, + "protocol": protocol, + } + return {**payload, "bundle_hash": _domain_hash(_REVIEW_BUNDLE_HASH_DOMAIN, payload)} + + +def build_route_study_review_bundle_from_input(payload: Any) -> Dict[str, Any]: + """Build a review bundle from the shared closed prompt-free input schema.""" + + study_plans, options = _split_protocol_build_input(payload) + return build_route_study_review_bundle(study_plans, **options) + + +def audit_route_study_review_bundle(bundle: Any) -> Dict[str, Any]: + """Reconstruct a bundled protocol from its complete canonical source plans.""" + + if not isinstance(bundle, Mapping): + raise ValueError("review bundle must be a JSON object") + if set(bundle) != _BUNDLE_TOP_LEVEL_KEYS: + raise ValueError("review bundle top-level fields do not match the v1 schema") + if bundle.get("schema_version") != REVIEW_BUNDLE_SCHEMA_VERSION: + raise ValueError("unsupported route study review bundle schema_version") + meta = _require_exact_keys(bundle.get("bundle"), _BUNDLE_META_KEYS, "review bundle metadata") + expected_meta = { + "version": REVIEW_BUNDLE_VERSION, + "label": REVIEW_BUNDLE_LABEL, + "state": "ready_for_semantic_review", + "verification_level": "full_source_bound_reconstruction", + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "activation_available": False, + } + if dict(meta) != expected_meta: + raise ValueError("review bundle metadata is not the canonical fail-closed v1 contract") + + builder = _require_exact_keys( + bundle.get("protocol_builder"), + {"schema_version", "options"}, + "review bundle protocol builder", + ) + if builder.get("schema_version") != PROTOCOL_BUILD_INPUT_SCHEMA_VERSION: + raise ValueError("review bundle protocol builder schema_version is unsupported") + raw_options = builder.get("options") + if not isinstance(raw_options, Mapping) or set(raw_options) != set(PROTOCOL_OPTION_KEYS): + raise ValueError("review bundle protocol options are incomplete or unsupported") + canonical_options = _normalize_protocol_options(**dict(raw_options)) + if dict(raw_options) != canonical_options: + raise ValueError("review bundle protocol options are not canonical") + + raw_plans = bundle.get("source_study_plans") + if not isinstance(raw_plans, list): + raise ValueError("review bundle source_study_plans must be a canonical list") + validated_plans = _normalize_study_plans(raw_plans) + canonical_plans = [row["plan"] for row in validated_plans] + if raw_plans != canonical_plans: + raise ValueError("review bundle source plans are not canonical or canonically ordered") + + supplied_protocol = bundle.get("protocol") + structural = audit_route_study_protocol(supplied_protocol) + reconstructed = build_route_study_protocol(canonical_plans, **canonical_options) + if _canonical_json(supplied_protocol) != _canonical_json(reconstructed): + raise ValueError("review bundle protocol does not reconstruct from its source plans") + + supplied_hash = _sha256_digest(bundle.get("bundle_hash"), "bundle_hash") + payload = { + "schema_version": bundle["schema_version"], + "bundle": dict(meta), + "protocol_builder": { + "schema_version": builder["schema_version"], + "options": dict(raw_options), + }, + "source_study_plans": raw_plans, + "protocol": supplied_protocol, + } + expected_hash = _domain_hash(_REVIEW_BUNDLE_HASH_DOMAIN, payload) + if supplied_hash != expected_hash: + raise ValueError("bundle_hash does not match the canonical review bundle") + return { + "ok": True, + "schema_version": REVIEW_BUNDLE_SCHEMA_VERSION, + "bundle_hash": supplied_hash, + "protocol_hash": structural["protocol_hash"], + "state": "ready_for_semantic_review", + "support_stratum_count": len(canonical_plans), + "verification_level": "full_source_bound_reconstruction", + "source_plan_reconstruction_performed": True, + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "activation_available": False, + "activation_blockers": list(ACTIVATION_BLOCKERS), + } + + +__all__ = [ + "CARRYOVER_SCOPES", + "DEFAULT_ANALYSIS_EVERY_CLUSTERS", + "DEFAULT_MAX_ROUTES_PER_CLUSTER", + "DEFAULT_PLANNED_CLUSTERS", + "DESIGN_MODES", + "INTERFERENCE_SCOPES", + "PROTOCOL_BUILD_INPUT_KEYS", + "PROTOCOL_BUILD_INPUT_SCHEMA_VERSION", + "PROTOCOL_LABEL", + "PROTOCOL_OPTION_KEYS", + "PROTOCOL_SCHEMA_VERSION", + "PROTOCOL_VERSION", + "REVIEW_BUNDLE_LABEL", + "REVIEW_BUNDLE_SCHEMA_VERSION", + "REVIEW_BUNDLE_VERSION", + "TEMPORAL_VARIATION_SCOPES", + "audit_route_study_review_bundle", + "audit_route_study_protocol", + "build_route_study_protocol_from_input", + "build_route_study_review_bundle", + "build_route_study_review_bundle_from_input", + "build_route_study_protocol", +] diff --git a/source/route_policy_protocol_cli.py b/source/route_policy_protocol_cli.py new file mode 100644 index 000000000..b315df2fc --- /dev/null +++ b/source/route_policy_protocol_cli.py @@ -0,0 +1,163 @@ +"""Export or audit a prompt-free stateful route experiment preflight. + +It can also create a portable review bundle that carries the complete canonical +source plans needed for semantic reconstruction. Neither surface signs, seals, +assigns, touches the ledger, runs inference, or enables promotion. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Dict, Optional, Sequence + +try: + from .route_policy_protocol import ( + PROTOCOL_BUILD_INPUT_KEYS, + audit_route_study_review_bundle, + audit_route_study_protocol, + build_route_study_protocol_from_input, + build_route_study_review_bundle_from_input, + ) + from .route_policy_study_cli import _example_payload, build_plan +except ImportError: # pragma: no cover - direct ``python source/...`` use + from route_policy_protocol import ( + PROTOCOL_BUILD_INPUT_KEYS, + audit_route_study_review_bundle, + audit_route_study_protocol, + build_route_study_protocol_from_input, + build_route_study_review_bundle_from_input, + ) + from route_policy_study_cli import _example_payload, build_plan + + +_INPUT_KEYS = set(PROTOCOL_BUILD_INPUT_KEYS) + + +def _example_protocol_input() -> Dict[str, Any]: + return { + "study_plans": [build_plan(_example_payload())], + "target_policy_profile": "balanced", + "design_mode": "sticky_session_cluster", + "carryover_scope": "unknown", + "interference_scope": "unknown", + "temporal_variation": "unknown", + "population_rule_id": "interactive-auto-route-opt-in", + "population_rule_version": "1", + "cluster_key_schema_version": "session-hash-v1", + "planned_clusters": 200, + "max_routes_per_cluster": 20, + "analysis_every_clusters": 50, + "block_length_routes": 20, + "washout_routes": 0, + } + + +def _example_bundle_input() -> Dict[str, Any]: + first_payload = _example_payload() + second_payload = _example_payload() + second_payload["source_contract"] = { + **second_payload["source_contract"], + "candidate_set_hash": "c" * 64, + "distribution_hash": "d" * 64, + } + payload = _example_protocol_input() + payload["study_plans"] = [build_plan(first_payload), build_plan(second_payload)] + return payload + + +def _read_json(path: str) -> Any: + raw = sys.stdin.read() if path == "-" else Path(path).read_text(encoding="utf-8-sig") + try: + return json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"input is not valid JSON: {exc.msg}") from exc + + +def _load_build_input(path: str) -> Dict[str, Any]: + payload = _read_json(path) + if not isinstance(payload, dict): + raise ValueError("input JSON must be an object") + unknown = set(payload) - _INPUT_KEYS + if unknown: + raise ValueError( + "input contains unsupported or non-prompt-free fields: " + + ", ".join(sorted(map(str, unknown))) + ) + if "study_plans" not in payload: + raise ValueError("input is missing required field: study_plans") + return payload + + +def build_protocol(payload: Dict[str, Any]) -> Dict[str, Any]: + return build_route_study_protocol_from_input(payload) + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export or audit a prompt-free, fail-closed route experiment preflight." + ) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--input", help="protocol-input JSON path, or - for stdin") + source.add_argument( + "--example", + action="store_true", + help="render a synthetic protocol draft; it is not live study evidence", + ) + source.add_argument("--audit", help="audit an existing protocol draft JSON path, or -") + source.add_argument( + "--bundle-input", + help="prompt-free protocol-input JSON path to bundle, or - for stdin", + ) + source.add_argument( + "--example-bundle", + action="store_true", + help="render a synthetic two-stratum semantic-review bundle", + ) + source.add_argument( + "--audit-bundle", + help="fully reconstruct and audit an existing review bundle JSON path, or -", + ) + parser.add_argument("--output", help="optional JSON output path; stdout is the default") + parser.add_argument("--compact", action="store_true", help="emit compact canonical JSON") + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + try: + if args.audit_bundle: + result = audit_route_study_review_bundle(_read_json(str(args.audit_bundle))) + elif args.audit: + result = audit_route_study_protocol(_read_json(str(args.audit))) + elif args.example_bundle: + result = build_route_study_review_bundle_from_input(_example_bundle_input()) + elif args.bundle_input: + result = build_route_study_review_bundle_from_input( + _load_build_input(str(args.bundle_input)) + ) + else: + payload = _example_protocol_input() if args.example else _load_build_input(str(args.input)) + result = build_protocol(payload) + rendered = json.dumps( + result, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":") if args.compact else None, + indent=None if args.compact else 2, + ) + if args.output: + Path(args.output).write_text(rendered + "\n", encoding="utf-8") + else: + print(rendered) + except (OSError, TypeError, ValueError) as exc: + print(f"route protocol error: {exc}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/source/route_policy_shadow_cli.py b/source/route_policy_shadow_cli.py new file mode 100644 index 000000000..b09910b33 --- /dev/null +++ b/source/route_policy_shadow_cli.py @@ -0,0 +1,352 @@ +"""Operate the local, shadow-only whole-policy commitment registry. + +All mutations are explicit local CLI actions. The command never executes a +route, writes the executed-decision ledger, enables a policy, or prints private +seed material. ``seal`` writes the private seed capsule to a separate file +before committing the public campaign package to SQLite so an interrupted +operation can be recovered safely. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import sqlite3 +import stat +import sys +from pathlib import Path +from typing import Any, Dict, Iterable, Optional, Sequence, Tuple + +try: + from .route_policy_shadow_registry import ( + MAX_SHADOW_VERIFY_BATCH, + RouteShadowAssignmentRegistry, + RouteShadowRegistryError, + audit_shadow_seed_capsule, + build_shadow_design_binding, + create_shadow_campaign_artifacts, + generate_shadow_seed, + ) +except ImportError: # pragma: no cover - direct ``python source/...`` use + from route_policy_shadow_registry import ( + MAX_SHADOW_VERIFY_BATCH, + RouteShadowAssignmentRegistry, + RouteShadowRegistryError, + audit_shadow_seed_capsule, + build_shadow_design_binding, + create_shadow_campaign_artifacts, + generate_shadow_seed, + ) + + +_SEED_CAPSULE_KEYS = { + "schema_version", + "campaign_id", + "design_binding_hash", + "seed_commitment", + "seed_material_base64url", +} + + +def _reject_duplicate_keys(pairs: Iterable[Tuple[str, Any]]) -> Dict[str, Any]: + value: Dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError(f"input JSON contains duplicate object key: {key}") + value[key] = item + return value + + +def _reject_non_finite(value: str) -> None: + raise ValueError(f"input JSON contains non-finite number: {value}") + + +def _read_json(path: str, *, name: str) -> Any: + raw = sys.stdin.read() if path == "-" else Path(path).read_text(encoding="utf-8-sig") + try: + return json.loads( + raw, + object_pairs_hook=_reject_duplicate_keys, + parse_constant=_reject_non_finite, + ) + except json.JSONDecodeError as exc: + raise ValueError(f"{name} is not valid JSON: {exc.msg}") from exc + + +def _read_json_object(path: str, *, name: str) -> Dict[str, Any]: + value = _read_json(path, name=name) + if not isinstance(value, dict): + raise ValueError(f"{name} must be a JSON object") + return value + + +def _canonical_json(value: Any) -> str: + return json.dumps( + value, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ) + + +def _render(value: Any, *, compact: bool) -> str: + return json.dumps( + value, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":") if compact else None, + indent=None if compact else 2, + ) + + +def _decode_capsule_seed(capsule: Dict[str, Any]) -> bytes: + if set(capsule) != _SEED_CAPSULE_KEYS: + raise ValueError("private seed capsule has an unsupported field set") + encoded = capsule.get("seed_material_base64url") + if not isinstance(encoded, str) or not encoded: + raise ValueError("private seed capsule does not contain canonical seed material") + padding = "=" * (-len(encoded) % 4) + try: + seed = base64.b64decode(encoded + padding, altchars=b"-_", validate=True) + except (ValueError, TypeError) as exc: + raise ValueError("private seed capsule does not contain canonical seed material") from exc + canonical = base64.urlsafe_b64encode(seed).decode("ascii").rstrip("=") + if len(seed) != 32 or canonical != encoded: + raise ValueError("private seed capsule does not contain canonical seed material") + return seed + + +def _load_recovery_capsule(path: Path, review_bundle: Dict[str, Any]) -> Tuple[Dict[str, Any], bytes]: + capsule = _read_json_object(str(path), name="private seed capsule") + seed = _decode_capsule_seed(capsule) + artifacts = create_shadow_campaign_artifacts(review_bundle, seed) + audit_shadow_seed_capsule(artifacts["public_package"], capsule) + expected = artifacts["private_seed_capsule"] + if _canonical_json(capsule) != _canonical_json(expected): + raise ValueError("private seed capsule does not match the review bundle") + return capsule, seed + + +def _write_private_capsule(path: Path, capsule: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + descriptor = os.open(str(path), flags, stat.S_IRUSR | stat.S_IWUSR) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as handle: + descriptor = -1 + handle.write(_render(capsule, compact=False) + "\n") + handle.flush() + os.fsync(handle.fileno()) + try: + os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + # Windows ACLs, rather than POSIX mode bits, are authoritative. + pass + except BaseException: + if descriptor >= 0: + os.close(descriptor) + try: + path.unlink() + except OSError: + pass + raise + + +def _same_path(first: Path, second: Path) -> bool: + return os.path.normcase(str(first.expanduser().resolve())) == os.path.normcase( + str(second.expanduser().resolve()) + ) + + +def _seal(args: argparse.Namespace) -> Dict[str, Any]: + bundle = _read_json_object(str(args.bundle), name="review bundle") + registry_path = Path(args.registry).expanduser() + capsule_path = Path(args.seed_output).expanduser() + if _same_path(registry_path, capsule_path): + raise ValueError("registry and private seed capsule paths must be different") + + capsule_written = False + capsule_recovered = False + if capsule_path.exists(): + _capsule, seed = _load_recovery_capsule(capsule_path, bundle) + capsule_recovered = True + else: + if registry_path.is_file(): + origin_hash = build_shadow_design_binding(bundle)["origin_review_bundle_hash"] + existing = RouteShadowAssignmentRegistry( + registry_path, read_only=True + ).snapshot() + if any( + row["origin_review_bundle_hash"] == origin_hash + for row in existing["campaigns"] + ): + raise ValueError( + "origin review bundle is already sealed; recover its original seed capsule" + ) + seed = generate_shadow_seed() + artifacts = create_shadow_campaign_artifacts(bundle, seed) + try: + _write_private_capsule(capsule_path, artifacts["private_seed_capsule"]) + capsule_written = True + except FileExistsError: + _capsule, seed = _load_recovery_capsule(capsule_path, bundle) + capsule_recovered = True + + result = RouteShadowAssignmentRegistry(registry_path).seal_campaign(bundle, seed) + + # Deliberately reconstruct the result instead of mutating the registry return + # so private seed material cannot reach stdout through a future extra field. + return { + "ok": bool(result["ok"]), + "created": bool(result["created"]), + "public_package": result["public_package"], + "private_seed_capsule_written": capsule_written, + "private_seed_capsule_recovered": capsule_recovered, + "private_seed_material_persisted_in_registry": False, + "private_seed_material_returned": False, + "execution_enabled": False, + "activation_available": False, + } + + +def _commit(args: argparse.Namespace) -> Dict[str, Any]: + capsule = _read_json_object(str(args.seed_input), name="private seed capsule") + cluster = _read_json_object(str(args.cluster_input), name="cluster input") + if set(cluster) != {"cluster_identifier"}: + raise ValueError("cluster input must contain exactly one field: cluster_identifier") + identifier = cluster["cluster_identifier"] + if not isinstance(identifier, str): + raise ValueError("cluster_identifier must be a JSON string") + return RouteShadowAssignmentRegistry(args.registry).append_assignment_commitment( + campaign_id=args.campaign, + seed_capsule=capsule, + cluster_identifier=identifier, + ) + + +def _close(args: argparse.Namespace) -> Dict[str, Any]: + return RouteShadowAssignmentRegistry(args.registry).close_campaign(args.campaign) + + +def _reveal(args: argparse.Namespace) -> Dict[str, Any]: + capsule = _read_json_object(str(args.seed_input), name="private seed capsule") + artifact = RouteShadowAssignmentRegistry(args.registry).reveal_seed( + campaign_id=args.campaign, + seed_capsule=capsule, + ) + reveal = artifact["reveal"] + # The registry retains the public opening needed for reconstruction, but + # command output remains a non-secret receipt by contract. + return { + "ok": True, + "campaign_id": reveal["campaign_id"], + "seal_hash": reveal["seal_hash"], + "seed_commitment": reveal["seed_commitment"], + "seed_material_fingerprint": reveal["seed_material_fingerprint"], + "seed_material_revealed_in_registry": True, + "seed_material_returned": False, + "state": reveal["state"], + "seed_reveal_hash": artifact["seed_reveal_hash"], + "execution_enabled": False, + "activation_available": False, + } + + +def _verify(args: argparse.Namespace) -> Dict[str, Any]: + return RouteShadowAssignmentRegistry(args.registry).verify_assignment_reveals( + args.campaign, + batch_size=args.batch_size, + ) + + +def _status(args: argparse.Namespace) -> Dict[str, Any]: + registry_path = Path(args.registry).expanduser() + if not registry_path.is_file(): + raise ValueError("shadow registry does not exist") + snapshot = RouteShadowAssignmentRegistry( + registry_path, read_only=True + ).snapshot(args.campaign) + return {**snapshot, "read_only": True} + + +def _add_common(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--registry", required=True, help="local shadow-registry SQLite path") + parser.add_argument("--compact", action="store_true", help="emit compact canonical JSON") + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Operate a local shadow-only whole-policy commitment/reveal registry; " + "no command executes or promotes a route policy." + ) + ) + commands = parser.add_subparsers(dest="command", required=True) + + seal = commands.add_parser("seal", help="seal a source-bound review bundle") + _add_common(seal) + seal.add_argument("--bundle", required=True, help="review-bundle JSON path") + seal.add_argument( + "--seed-output", + required=True, + help="separate private seed-capsule JSON path (created exclusively)", + ) + seal.set_defaults(handler=_seal) + + commit = commands.add_parser("commit", help="append one opaque cluster commitment") + _add_common(commit) + commit.add_argument("--campaign", required=True, help="sealed shadow campaign identifier") + commit.add_argument("--seed-input", required=True, help="private seed-capsule JSON path") + commit.add_argument( + "--cluster-input", + required=True, + help='JSON path containing only {"cluster_identifier": "..."}', + ) + commit.set_defaults(handler=_commit) + + close = commands.add_parser("close", help="freeze the commitment population") + _add_common(close) + close.add_argument("--campaign", required=True, help="sealed shadow campaign identifier") + close.set_defaults(handler=_close) + + reveal = commands.add_parser("reveal", help="reveal the seed after commitment closure") + _add_common(reveal) + reveal.add_argument("--campaign", required=True, help="sealed shadow campaign identifier") + reveal.add_argument("--seed-input", required=True, help="private seed-capsule JSON path") + reveal.set_defaults(handler=_reveal) + + verify = commands.add_parser("verify", help="reconstruct a bounded batch of assignments") + _add_common(verify) + verify.add_argument("--campaign", required=True, help="sealed shadow campaign identifier") + verify.add_argument( + "--batch-size", + type=int, + default=MAX_SHADOW_VERIFY_BATCH, + help=f"maximum assignments to reconstruct (1-{MAX_SHADOW_VERIFY_BATCH})", + ) + verify.set_defaults(handler=_verify) + + status = commands.add_parser("status", help="read and verify registry state") + _add_common(status) + status.add_argument("--campaign", help="optional sealed shadow campaign identifier") + status.set_defaults(handler=_status) + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + try: + result = args.handler(args) + print(_render(result, compact=bool(args.compact))) + except (KeyError, OSError, sqlite3.Error, TypeError, ValueError, RouteShadowRegistryError) as exc: + print(f"route shadow registry error: {exc}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/source/route_policy_shadow_registry.py b/source/route_policy_shadow_registry.py new file mode 100644 index 000000000..0182c060f --- /dev/null +++ b/source/route_policy_shadow_registry.py @@ -0,0 +1,2268 @@ +"""Append-only, shadow-only whole-policy assignment commitments. + +The registry is deliberately separate from the executed route-decision ledger. +It can freeze a reviewed campaign, commit opaque cluster assignments, close +enrollment, reveal the seed, and reconstruct every committed whole-policy arm. +It cannot execute a route, emit a logging propensity, estimate a causal effect, +or promote a policy. +""" + +from __future__ import annotations + +import base64 +import hashlib +import hmac +import json +import re +import secrets +import sqlite3 +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple + +try: + from .route_policy_protocol import ( + PROTOCOL_OPTION_KEYS, + audit_route_study_review_bundle, + build_route_study_review_bundle, + ) +except ImportError: # pragma: no cover - direct ``python source/...`` use + from route_policy_protocol import ( + PROTOCOL_OPTION_KEYS, + audit_route_study_review_bundle, + build_route_study_review_bundle, + ) + + +SHADOW_REGISTRY_SCHEMA_VERSION = 1 +SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION = "route-study-shadow-public-package-v1" +SHADOW_DESIGN_BINDING_SCHEMA_VERSION = "route-study-shadow-design-binding-v1" +SHADOW_ASSIGNMENT_MANIFEST_SCHEMA_VERSION = "route-study-shadow-assignment-manifest-v1" +SHADOW_CAMPAIGN_SEAL_SCHEMA_VERSION = "route-study-shadow-campaign-seal-v1" +SHADOW_SEED_CAPSULE_SCHEMA_VERSION = "route-study-shadow-seed-capsule-v1" +SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION = "route-study-shadow-assignment-commitment-v1" +SHADOW_CAMPAIGN_CLOSURE_SCHEMA_VERSION = "route-study-shadow-campaign-closure-v1" +SHADOW_SEED_REVEAL_SCHEMA_VERSION = "route-study-shadow-seed-reveal-v1" +SHADOW_ASSIGNMENT_REVEAL_SCHEMA_VERSION = "route-study-shadow-assignment-reveal-v1" +SHADOW_REGISTRY_EVENT_SCHEMA_VERSION = "route-study-shadow-registry-event-v1" +SHADOW_REGISTRY_SNAPSHOT_SCHEMA_VERSION = "route-study-shadow-registry-snapshot-v1" + +SHADOW_ASSIGNMENT_ALGORITHM = "hkdf-sha256-hmac-sha256-whole-policy-bps-v1" +SHADOW_CANONICALIZATION = "rfc8785-jcs-restricted-ijson-integer-v1" +SHADOW_LEGACY_BUNDLE_ENCODING = "route-review-bundle-v1-python-sorted-compact-utf8" +SHADOW_SCHEMA_OBJECTS_SHA256 = "fe12a50841e5f983adce9ef1eb4ebf6635ab079f8c0043f4ca68500b38f9a04c" +SHADOW_TOTAL_ALLOCATION_BPS = 10_000 +SHADOW_CANDIDATE_ALLOCATION_BPS = 5_000 +SHADOW_BLOCK_ID = 0 +MAX_SHADOW_CAMPAIGNS = 1_000 +MAX_SHADOW_COMMITMENTS_PER_CAMPAIGN = 1_000_000 +MAX_SHADOW_VERIFY_BATCH = 1_000 + +_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,159}$") +_SHA256_RE = re.compile(r"^[0-9a-f]{64}$") +_BASE64URL_RE = re.compile(r"^[A-Za-z0-9_-]{43}$") + +_DESIGN_HASH_DOMAIN = b"supermix.route-shadow.design-binding.v1\x00" +_ARM_HASH_DOMAIN = b"supermix.route-shadow.whole-policy-arm.v1\x00" +_SOURCE_POLICY_CLASS_HASH_DOMAIN = b"supermix.route-shadow.source-policy-class.v1\x00" +_MANIFEST_HASH_DOMAIN = b"supermix.route-shadow.assignment-manifest.v1\x00" +_SEED_COMMITMENT_DOMAIN = b"supermix.route-shadow.seed-commitment.v1\x00" +_SEAL_HASH_DOMAIN = b"supermix.route-shadow.campaign-seal.v1\x00" +_PACKAGE_HASH_DOMAIN = b"supermix.route-shadow.public-package.v1\x00" +_IDENTITY_INFO = b"supermix.route-shadow.identity-key.v1\x00" +_ASSIGNMENT_INFO = b"supermix.route-shadow.assignment-key.v1\x00" +_PSEUDONYM_DOMAIN = b"supermix.route-shadow.cluster-pseudonym.v1\x00" +_DRAW_DOMAIN = b"supermix.route-shadow.whole-policy-draw.v1\x00" +_DRAW_HASH_DOMAIN = b"supermix.route-shadow.draw-proof.v1\x00" +_ASSIGNMENT_REVEAL_HASH_DOMAIN = b"supermix.route-shadow.assignment-reveal.v1\x00" +_ASSIGNMENT_COMMITMENT_HASH_DOMAIN = b"supermix.route-shadow.assignment-commitment.v1\x00" +_CLOSURE_HASH_DOMAIN = b"supermix.route-shadow.campaign-closure.v1\x00" +_SEED_REVEAL_HASH_DOMAIN = b"supermix.route-shadow.seed-reveal.v1\x00" +_REGISTRY_EVENT_HASH_DOMAIN = b"supermix.route-shadow.registry-event.v1\x00" +_MAX_SAFE_JSON_INTEGER = (1 << 53) - 1 + +_BOUNDARIES = { + "shadow_only": True, + "ledger_eligible": False, + "executed_logging_propensity": False, + "route_execution_enabled": False, + "model_inference_enabled": False, + "causal_estimate_available": False, + "activation_available": False, + "automatic_promotion_allowed": False, + "causal_design_certified": False, + "pseudonymity_is_anonymity": False, + "post_reveal_identity_unlinkability_guaranteed": False, +} + +_REQUIRED_SHADOW_TABLES = frozenset( + { + "shadow_registry_metadata", + "shadow_campaign_seals", + "shadow_assignment_commitments", + "shadow_campaign_closures", + "shadow_seed_reveals", + "shadow_assignment_reveals", + "shadow_registry_events", + } +) +_REQUIRED_SHADOW_TRIGGERS = frozenset( + { + "shadow_commitments_require_open_campaign", + "shadow_seed_reveal_requires_closure", + "shadow_assignment_reveal_requires_seed", + "shadow_assignment_reveal_campaign_matches_commitment", + "shadow_metadata_no_update", + "shadow_metadata_no_delete", + "shadow_seals_no_update", + "shadow_seals_no_delete", + "shadow_commitments_no_update", + "shadow_commitments_no_delete", + "shadow_closures_no_update", + "shadow_closures_no_delete", + "shadow_seed_reveals_no_update", + "shadow_seed_reveals_no_delete", + "shadow_assignment_reveals_no_update", + "shadow_assignment_reveals_no_delete", + "shadow_events_no_update", + "shadow_events_no_delete", + } +) +_REQUIRED_SHADOW_INDEXES = frozenset( + { + "shadow_commitments_campaign_order_idx", + "shadow_reveals_campaign_order_idx", + } +) + + +class RouteShadowRegistryError(RuntimeError): + """Base class for stable registry failures.""" + + +class ShadowRegistryConflictError(RouteShadowRegistryError): + """Raised when an immutable campaign artifact conflicts with existing state.""" + + +def _require_exact_keys(value: Any, expected: set[str], name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or set(value) != expected: + raise ValueError(f"{name} does not match the v1 schema") + return value + + +def _identifier(value: Any, name: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{name} must be a closed ASCII identifier") + cooked = value.strip() + if not _IDENTIFIER_RE.fullmatch(cooked): + raise ValueError(f"{name} must be a closed ASCII identifier") + return cooked + + +def _sha256(value: Any, name: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + cooked = value.strip() + if not _SHA256_RE.fullmatch(cooked): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + return cooked + + +def _bounded_int(value: Any, name: str, *, minimum: int, maximum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be an integer") + cooked = value + if cooked < minimum or cooked > maximum: + raise ValueError(f"{name} must be between {minimum} and {maximum}") + return cooked + + +def _assert_shadow_json(value: Any, name: str = "shadow artifact") -> None: + if value is None or isinstance(value, bool): + return + if isinstance(value, int): + if abs(value) > _MAX_SAFE_JSON_INTEGER: + raise ValueError(f"{name} must use I-JSON safe integers") + return + if isinstance(value, str): + try: + value.encode("utf-8") + except UnicodeEncodeError as exc: + raise ValueError(f"{name} must not contain lone Unicode surrogates") from exc + return + if isinstance(value, float): + raise ValueError(f"{name} must not contain floating-point values") + if isinstance(value, Mapping): + for key, item in value.items(): + if not isinstance(key, str): + raise ValueError(f"{name} object keys must be strings") + if not key.isascii(): + raise ValueError(f"{name} object keys must be ASCII for portable sorting") + _assert_shadow_json(item, name) + return + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for item in value: + _assert_shadow_json(item, name) + return + raise ValueError(f"{name} contains a non-JSON value") + + +def _canonical_json(value: Any, name: str = "shadow artifact") -> str: + _assert_shadow_json(value, name) + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be canonical finite JSON") from exc + + +def _domain_hash(domain: bytes, value: Any, name: str = "shadow artifact") -> str: + return hashlib.sha256(domain + _canonical_json(value, name).encode("utf-8")).hexdigest() + + +def _schema_integrity(connection: sqlite3.Connection) -> Dict[str, Any]: + rows = connection.execute( + """ + SELECT type, name, sql + FROM sqlite_master + WHERE type IN ('table', 'trigger', 'index') + """ + ).fetchall() + found = { + kind: {str(row["name"]) for row in rows if str(row["type"]) == kind} + for kind in ("table", "trigger", "index") + } + missing_tables = sorted(_REQUIRED_SHADOW_TABLES - found["table"]) + missing_triggers = sorted(_REQUIRED_SHADOW_TRIGGERS - found["trigger"]) + missing_indexes = sorted(_REQUIRED_SHADOW_INDEXES - found["index"]) + required_names = ( + _REQUIRED_SHADOW_TABLES + | _REQUIRED_SHADOW_TRIGGERS + | _REQUIRED_SHADOW_INDEXES + ) + definition_rows = sorted( + ( + { + "type": str(row["type"]), + "name": str(row["name"]), + "sql": " ".join(str(row["sql"] or "").split()), + } + for row in rows + if str(row["name"]) in required_names + ), + key=lambda row: (row["type"], row["name"]), + ) + definition_fingerprint = hashlib.sha256( + _canonical_json(definition_rows, "shadow schema definitions").encode("utf-8") + ).hexdigest() + definitions_ok = hmac.compare_digest( + definition_fingerprint, SHADOW_SCHEMA_OBJECTS_SHA256 + ) + return { + "ok": not (missing_tables or missing_triggers or missing_indexes) + and definitions_ok, + "required_tables": len(_REQUIRED_SHADOW_TABLES), + "required_triggers": len(_REQUIRED_SHADOW_TRIGGERS), + "required_indexes": len(_REQUIRED_SHADOW_INDEXES), + "missing_tables": missing_tables, + "missing_triggers": missing_triggers, + "missing_indexes": missing_indexes, + "definitions_ok": definitions_ok, + "definition_fingerprint": definition_fingerprint, + "expected_definition_fingerprint": SHADOW_SCHEMA_OBJECTS_SHA256, + } + + +def _seed_bytes(value: Any) -> bytes: + if not isinstance(value, (bytes, bytearray)): + raise ValueError("shadow seed material must be bytes") + cooked = bytes(value) + if len(cooked) != 32: + raise ValueError("shadow seed material must contain exactly 32 bytes") + return cooked + + +def generate_shadow_seed() -> bytes: + """Return 256 bits from the operating-system CSPRNG.""" + + return secrets.token_bytes(32) + + +def _encode_seed(seed: bytes) -> str: + return base64.urlsafe_b64encode(_seed_bytes(seed)).decode("ascii").rstrip("=") + + +def _decode_seed(value: Any) -> bytes: + if not isinstance(value, str): + raise ValueError("seed_material_base64url must encode exactly 32 bytes") + cooked = value.strip() + if not _BASE64URL_RE.fullmatch(cooked): + raise ValueError("seed_material_base64url must encode exactly 32 bytes") + try: + seed = base64.urlsafe_b64decode(cooked + "=") + except (ValueError, TypeError) as exc: + raise ValueError("seed_material_base64url is invalid") from exc + seed = _seed_bytes(seed) + if not hmac.compare_digest(_encode_seed(seed), cooked): + raise ValueError("seed_material_base64url is not canonical") + return seed + + +def _reject_duplicate_object_keys(pairs: Sequence[Tuple[str, Any]]) -> Dict[str, Any]: + value: Dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError(f"review bundle JSON contains duplicate object key: {key}") + value[key] = item + return value + + +def _canonical_review_bundle_json(bundle: Any, name: str) -> str: + """Freeze a legacy reviewed bundle without importing its floats into v1 artifacts.""" + + audit_route_study_review_bundle(bundle) + try: + return json.dumps( + bundle, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be canonical finite JSON") from exc + + +def _decode_review_bundle_json(value: Any, name: str) -> Dict[str, Any]: + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must be a canonical JSON string") + try: + decoded = json.loads( + value, + object_pairs_hook=_reject_duplicate_object_keys, + parse_constant=lambda token: (_ for _ in ()).throw( + ValueError(f"{name} contains non-finite number {token}") + ), + ) + except (TypeError, ValueError, json.JSONDecodeError) as exc: + raise ValueError(f"{name} must be a strict canonical JSON string") from exc + if not isinstance(decoded, dict): + raise ValueError(f"{name} must decode to a JSON object") + canonical = _canonical_review_bundle_json(decoded, name) + if canonical != value: + raise ValueError(f"{name} is not in canonical JSON form") + return decoded + + +def _hkdf_expand(seed: bytes, *, salt: bytes, info: bytes) -> bytes: + """Derive one SHA-256-sized subkey using RFC 5869 extract-and-expand.""" + + prk = hmac.new(salt, _seed_bytes(seed), hashlib.sha256).digest() + return hmac.new(prk, info + b"\x01", hashlib.sha256).digest() + + +def _review_bundle_parts(bundle: Any) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]: + verification = audit_route_study_review_bundle(bundle) + if verification.get("verification_level") != "full_source_bound_reconstruction": + raise ValueError("shadow campaigns require full source-bound reconstruction") + if not isinstance(bundle, Mapping): + raise ValueError("review bundle must be a JSON object") + protocol = bundle.get("protocol") + builder = bundle.get("protocol_builder") + if not isinstance(protocol, Mapping) or not isinstance(builder, Mapping): + raise ValueError("review bundle protocol and builder are required") + options = builder.get("options") + if not isinstance(options, Mapping) or set(options) != set(PROTOCOL_OPTION_KEYS): + raise ValueError("review bundle protocol options are incomplete") + return dict(protocol), dict(options), dict(verification) + + +def _whole_policy_arms(protocol: Mapping[str, Any]) -> List[Dict[str, Any]]: + charter = protocol.get("charter") + if not isinstance(charter, Mapping): + raise ValueError("review bundle protocol charter is missing") + source = charter.get("source_studies") + target = charter.get("target_policy_class") + if not isinstance(source, Mapping) or not isinstance(target, Mapping): + raise ValueError("review bundle source and target policy bindings are missing") + common = source.get("common_source_contract") + if not isinstance(common, Mapping): + raise ValueError("review bundle source policy cohort is missing") + support_strata = source.get("support_strata") + if not isinstance(support_strata, list) or not support_strata: + raise ValueError("review bundle source support strata are missing") + frozen_source_strata: List[Dict[str, Any]] = [] + for index, stratum in enumerate(support_strata): + if not isinstance(stratum, Mapping): + raise ValueError(f"source support stratum {index} must be an object") + frozen_source_strata.append( + { + "study_design_hash": _sha256( + stratum.get("study_design_hash"), "source study_design_hash" + ), + "candidate_set_hash": _sha256( + stratum.get("candidate_set_hash"), "source candidate_set_hash" + ), + "distribution_hash": _sha256( + stratum.get("distribution_hash"), "source distribution_hash" + ), + "baseline_action": _identifier( + stratum.get("baseline_action"), "source baseline_action" + ), + } + ) + frozen_source_strata.sort(key=lambda row: row["study_design_hash"]) + source_class_manifest = { + "schema_version": "route-study-shadow-source-policy-class-v1", + "common_source_contract": { + "policy_id": _identifier(common.get("policy_id"), "source policy_id"), + "policy_version": _identifier( + common.get("policy_version"), "source policy_version" + ), + "feature_schema_version": _identifier( + common.get("feature_schema_version"), "source feature_schema_version" + ), + "support_schema_version": _identifier( + common.get("support_schema_version"), "source support_schema_version" + ), + "outcome_contract_schema_version": _identifier( + common.get("outcome_contract_schema_version"), + "source outcome_contract_schema_version", + ), + }, + "support_strata": frozen_source_strata, + } + source_class_hash = _domain_hash( + _SOURCE_POLICY_CLASS_HASH_DOMAIN, + source_class_manifest, + "source whole-policy class", + ) + source_binding = { + "binding_type": "frozen_source_policy_class", + "policy_class_hash": source_class_hash, + "policy_class_manifest": source_class_manifest, + } + target_class_hash = _sha256(target.get("class_hash"), "target class_hash") + target_class_manifest = json.loads( + _canonical_json( + {key: value for key, value in target.items() if key != "class_hash"}, + "target whole-policy class", + ) + ) + if _domain_hash( + b"supermix.route-study.target-class.v1\x00", + target_class_manifest, + "target whole-policy class", + ) != target_class_hash: + # The review-bundle auditor already validates this using the protocol's + # source-of-truth profile catalog. This guard documents the invariant + # without allowing a partially described treatment arm. + raise ValueError("target policy class hash does not bind its full manifest") + target_binding = { + "binding_type": "frozen_target_policy_class", + "policy_class_hash": target_class_hash, + "policy_class_manifest": target_class_manifest, + } + rows = [ + { + "arm_id": "incumbent_source_policy", + "allocation_bps": SHADOW_TOTAL_ALLOCATION_BPS - SHADOW_CANDIDATE_ALLOCATION_BPS, + "policy_binding": source_binding, + }, + { + "arm_id": "candidate_target_policy", + "allocation_bps": SHADOW_CANDIDATE_ALLOCATION_BPS, + "policy_binding": target_binding, + }, + ] + return [ + {**row, "arm_hash": _domain_hash(_ARM_HASH_DOMAIN, row, "whole-policy arm")} + for row in rows + ] + + +def build_shadow_design_binding(review_bundle: Any) -> Dict[str, Any]: + """Bind explicit whole-policy arms without treating route actions as arms.""" + + protocol, options, verification = _review_bundle_parts(review_bundle) + if options.get("seed_commitment") is not None: + raise ValueError("origin review bundle must not already contain seed material") + charter = protocol["charter"] + design = charter["stateful_design"] + population = charter["population"] + if design.get("selected_design_mode") != "sticky_session_cluster": + raise ValueError("shadow registry v1 supports sticky_session_cluster only") + if design.get("assignment_unit") != "session_hash": + raise ValueError("shadow registry v1 requires the session_hash assignment unit") + binding = { + "schema_version": SHADOW_DESIGN_BINDING_SCHEMA_VERSION, + "origin_review_bundle_hash": _sha256( + verification.get("bundle_hash"), "origin review bundle_hash" + ), + "origin_review_bundle_artifact_sha256": hashlib.sha256( + _canonical_review_bundle_json(review_bundle, "origin review bundle").encode("utf-8") + ).hexdigest(), + "origin_protocol_hash": _sha256( + verification.get("protocol_hash"), "origin protocol_hash" + ), + "canonicalization": SHADOW_CANONICALIZATION, + "legacy_review_bundle_encoding": SHADOW_LEGACY_BUNDLE_ENCODING, + "design_mode": "sticky_session_cluster", + "assignment_unit": "study_scoped_session_cluster", + "cluster_key_schema_version": _identifier( + population.get("cluster_key_schema_version"), "cluster_key_schema_version" + ), + "population_rule_id": _identifier( + population.get("population_rule_id"), "population_rule_id" + ), + "population_rule_version": _identifier( + population.get("population_rule_version"), "population_rule_version" + ), + "planned_cluster_ceiling": _bounded_int( + population.get("planned_clusters"), + "planned_cluster_ceiling", + minimum=2, + maximum=MAX_SHADOW_COMMITMENTS_PER_CAMPAIGN, + ), + "admitted_support_strata": [ + _sha256(value, "admitted support stratum") + for value in population.get("admitted_support_strata", []) + ], + "whole_policy_arms": _whole_policy_arms(protocol), + "assignment_algorithm": SHADOW_ASSIGNMENT_ALGORITHM, + "allocation_total_bps": SHADOW_TOTAL_ALLOCATION_BPS, + "boundaries": dict(_BOUNDARIES), + } + if not binding["admitted_support_strata"]: + raise ValueError("shadow design must bind at least one support stratum") + if sum(row["allocation_bps"] for row in binding["whole_policy_arms"]) != SHADOW_TOTAL_ALLOCATION_BPS: + raise ValueError("whole-policy arm allocation must sum to 10000 basis points") + return { + **binding, + "design_binding_hash": _domain_hash( + _DESIGN_HASH_DOMAIN, binding, "shadow design binding" + ), + } + + +def compute_shadow_seed_commitment(design_binding_hash: Any, seed_material: Any) -> str: + design_hash = _sha256(design_binding_hash, "design_binding_hash") + seed = _seed_bytes(seed_material) + return hashlib.sha256( + _SEED_COMMITMENT_DOMAIN + bytes.fromhex(design_hash) + seed + ).hexdigest() + + +def _rebuild_review_bundle_with_seed(review_bundle: Mapping[str, Any], seed_commitment: str) -> Dict[str, Any]: + builder = review_bundle["protocol_builder"] + options = dict(builder["options"]) + options["seed_commitment"] = _sha256(seed_commitment, "seed_commitment") + return build_route_study_review_bundle(review_bundle["source_study_plans"], **options) + + +def _assignment_manifest( + design_binding: Mapping[str, Any], + committed_bundle: Mapping[str, Any], + seed_commitment: str, +) -> Dict[str, Any]: + verification = audit_route_study_review_bundle(committed_bundle) + payload = { + "schema_version": SHADOW_ASSIGNMENT_MANIFEST_SCHEMA_VERSION, + "manifest": { + "version": 1, + "label": "Sealed Shadow Whole-Policy Assignment Manifest", + "state": "shadow_assignment_design", + "design_binding_hash": _sha256( + design_binding.get("design_binding_hash"), "design_binding_hash" + ), + "committed_review_bundle_hash": _sha256( + verification.get("bundle_hash"), "committed review bundle_hash" + ), + "committed_review_bundle_artifact_sha256": hashlib.sha256( + _canonical_review_bundle_json( + committed_bundle, "committed review bundle" + ).encode("utf-8") + ).hexdigest(), + "committed_protocol_hash": _sha256( + verification.get("protocol_hash"), "committed protocol_hash" + ), + "seed_commitment": _sha256(seed_commitment, "seed_commitment"), + "assignment_algorithm": SHADOW_ASSIGNMENT_ALGORITHM, + "canonicalization": SHADOW_CANONICALIZATION, + "legacy_review_bundle_encoding": SHADOW_LEGACY_BUNDLE_ENCODING, + "whole_policy_arms": json.loads( + _canonical_json( + design_binding["whole_policy_arms"], + "whole-policy arms", + ) + ), + "allocation_total_bps": SHADOW_TOTAL_ALLOCATION_BPS, + "block_id": SHADOW_BLOCK_ID, + "boundaries": dict(_BOUNDARIES), + }, + } + return { + **payload, + "manifest_hash": _domain_hash( + _MANIFEST_HASH_DOMAIN, payload, "shadow assignment manifest" + ), + } + + +def _campaign_seal( + design_binding: Mapping[str, Any], + manifest: Mapping[str, Any], + seed_commitment: str, +) -> Dict[str, Any]: + design_hash = _sha256(design_binding.get("design_binding_hash"), "design_binding_hash") + campaign_id = f"shadow:{design_hash}" + payload = { + "schema_version": SHADOW_CAMPAIGN_SEAL_SCHEMA_VERSION, + "seal": { + "version": 1, + "label": "Local Shadow Assignment Seal", + "state": "sealed_shadow_only", + "campaign_id": campaign_id, + "design_binding_hash": design_hash, + "origin_review_bundle_hash": _sha256( + design_binding.get("origin_review_bundle_hash"), "origin review bundle_hash" + ), + "committed_review_bundle_hash": _sha256( + manifest["manifest"]["committed_review_bundle_hash"], + "committed review bundle_hash", + ), + "committed_protocol_hash": _sha256( + manifest["manifest"]["committed_protocol_hash"], "committed protocol_hash" + ), + "manifest_hash": _sha256(manifest.get("manifest_hash"), "manifest_hash"), + "seed_commitment": _sha256(seed_commitment, "seed_commitment"), + "assignment_algorithm": SHADOW_ASSIGNMENT_ALGORITHM, + "canonicalization": SHADOW_CANONICALIZATION, + "private_seed_material_included": False, + "local_append_only_registry_required": True, + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "independent_seed_custody_verified": False, + "nonce_grinding_resistance_proven": False, + "boundaries": dict(_BOUNDARIES), + }, + } + return {**payload, "seal_hash": _domain_hash(_SEAL_HASH_DOMAIN, payload, "campaign seal")} + + +def _seed_capsule(seal: Mapping[str, Any], seed_material: bytes) -> Dict[str, Any]: + return { + "schema_version": SHADOW_SEED_CAPSULE_SCHEMA_VERSION, + "campaign_id": _identifier(seal["seal"]["campaign_id"], "campaign_id"), + "design_binding_hash": _sha256( + seal["seal"]["design_binding_hash"], "design_binding_hash" + ), + "seed_commitment": _sha256(seal["seal"]["seed_commitment"], "seed_commitment"), + "seed_material_base64url": _encode_seed(seed_material), + } + + +def create_shadow_campaign_artifacts(review_bundle: Any, seed_material: Any) -> Dict[str, Any]: + """Create a public campaign package and a separately handled seed capsule.""" + + if not isinstance(review_bundle, Mapping): + raise ValueError("review bundle must be a JSON object") + seed = _seed_bytes(seed_material) + design_binding = build_shadow_design_binding(review_bundle) + seed_commitment = compute_shadow_seed_commitment( + design_binding["design_binding_hash"], seed + ) + committed_bundle = _rebuild_review_bundle_with_seed(review_bundle, seed_commitment) + manifest = _assignment_manifest(design_binding, committed_bundle, seed_commitment) + seal = _campaign_seal(design_binding, manifest, seed_commitment) + payload = { + "schema_version": SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION, + "origin_review_bundle": _canonical_review_bundle_json( + review_bundle, "origin review bundle" + ), + "committed_review_bundle": _canonical_review_bundle_json( + committed_bundle, "committed review bundle" + ), + "design_binding": design_binding, + "assignment_manifest": manifest, + "campaign_seal": seal, + } + public_package = { + **payload, + "public_package_hash": _domain_hash( + _PACKAGE_HASH_DOMAIN, payload, "shadow public package" + ), + } + return { + "public_package": public_package, + "private_seed_capsule": _seed_capsule(seal, seed), + } + + +def audit_shadow_campaign_artifacts(public_package: Any) -> Dict[str, Any]: + package = _require_exact_keys( + public_package, + { + "schema_version", + "origin_review_bundle", + "committed_review_bundle", + "design_binding", + "assignment_manifest", + "campaign_seal", + "public_package_hash", + }, + "shadow public package", + ) + if package.get("schema_version") != SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION: + raise ValueError("unsupported shadow public package schema_version") + origin_encoded = package["origin_review_bundle"] + committed_encoded = package["committed_review_bundle"] + origin = _decode_review_bundle_json(origin_encoded, "origin review bundle") + committed = _decode_review_bundle_json(committed_encoded, "committed review bundle") + origin_protocol, origin_options, origin_verification = _review_bundle_parts(origin) + committed_protocol, committed_options, committed_verification = _review_bundle_parts(committed) + if origin_options.get("seed_commitment") is not None: + raise ValueError("origin review bundle must remain seedless") + seed_commitment = _sha256( + committed_options.get("seed_commitment"), "committed seed_commitment" + ) + expected_options = dict(origin_options) + expected_options["seed_commitment"] = seed_commitment + if committed_options != expected_options: + raise ValueError("committed review bundle changed fields other than seed_commitment") + if committed["source_study_plans"] != origin["source_study_plans"]: + raise ValueError("committed review bundle changed source study plans") + + expected_design = build_shadow_design_binding(origin) + if _canonical_json(package["design_binding"]) != _canonical_json(expected_design): + raise ValueError("shadow design binding does not reconstruct from the origin bundle") + expected_committed = _rebuild_review_bundle_with_seed(origin, seed_commitment) + if _canonical_review_bundle_json( + committed, "committed review bundle" + ) != _canonical_review_bundle_json(expected_committed, "expected committed review bundle"): + raise ValueError("committed review bundle does not reconstruct") + expected_manifest = _assignment_manifest(expected_design, expected_committed, seed_commitment) + if _canonical_json(package["assignment_manifest"]) != _canonical_json(expected_manifest): + raise ValueError("shadow assignment manifest does not reconstruct") + expected_seal = _campaign_seal(expected_design, expected_manifest, seed_commitment) + if _canonical_json(package["campaign_seal"]) != _canonical_json(expected_seal): + raise ValueError("shadow campaign seal does not reconstruct") + + payload = { + "schema_version": package["schema_version"], + "origin_review_bundle": origin_encoded, + "committed_review_bundle": committed_encoded, + "design_binding": package["design_binding"], + "assignment_manifest": package["assignment_manifest"], + "campaign_seal": package["campaign_seal"], + } + supplied_hash = _sha256(package.get("public_package_hash"), "public_package_hash") + expected_hash = _domain_hash(_PACKAGE_HASH_DOMAIN, payload, "shadow public package") + if not hmac.compare_digest(supplied_hash, expected_hash): + raise ValueError("public_package_hash does not match the canonical package") + return { + "ok": True, + "schema_version": SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION, + "campaign_id": expected_seal["seal"]["campaign_id"], + "public_package_hash": supplied_hash, + "origin_review_bundle_hash": origin_verification["bundle_hash"], + "origin_review_bundle_artifact_sha256": hashlib.sha256( + origin_encoded.encode("utf-8") + ).hexdigest(), + "committed_review_bundle_hash": committed_verification["bundle_hash"], + "committed_review_bundle_artifact_sha256": hashlib.sha256( + committed_encoded.encode("utf-8") + ).hexdigest(), + "committed_protocol_hash": committed_verification["protocol_hash"], + "design_binding_hash": expected_design["design_binding_hash"], + "manifest_hash": expected_manifest["manifest_hash"], + "seal_hash": expected_seal["seal_hash"], + "seed_commitment": seed_commitment, + "whole_policy_arm_count": len(expected_manifest["manifest"]["whole_policy_arms"]), + "support_stratum_count": committed_verification["support_stratum_count"], + "verification_level": "full_source_bound_shadow_seal_reconstruction", + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "execution_enabled": False, + "activation_available": False, + } + + +def _seed_from_capsule(public_package: Mapping[str, Any], capsule: Any) -> bytes: + verification = audit_shadow_campaign_artifacts(public_package) + value = _require_exact_keys( + capsule, + { + "schema_version", + "campaign_id", + "design_binding_hash", + "seed_commitment", + "seed_material_base64url", + }, + "shadow seed capsule", + ) + if value.get("schema_version") != SHADOW_SEED_CAPSULE_SCHEMA_VERSION: + raise ValueError("unsupported shadow seed capsule schema_version") + if value.get("campaign_id") != verification["campaign_id"]: + raise ValueError("shadow seed capsule campaign_id does not match") + if value.get("design_binding_hash") != verification["design_binding_hash"]: + raise ValueError("shadow seed capsule design binding does not match") + if value.get("seed_commitment") != verification["seed_commitment"]: + raise ValueError("shadow seed capsule commitment does not match") + seed = _decode_seed(value.get("seed_material_base64url")) + expected = compute_shadow_seed_commitment(verification["design_binding_hash"], seed) + if not hmac.compare_digest(expected, verification["seed_commitment"]): + raise ValueError("shadow seed material does not open the sealed commitment") + return seed + + +def audit_shadow_seed_capsule(public_package: Any, capsule: Any) -> Dict[str, Any]: + seed = _seed_from_capsule(public_package, capsule) + verification = audit_shadow_campaign_artifacts(public_package) + return { + "ok": True, + "schema_version": SHADOW_SEED_CAPSULE_SCHEMA_VERSION, + "campaign_id": verification["campaign_id"], + "seed_commitment": verification["seed_commitment"], + "seed_length_bytes": len(seed), + "seed_commitment_verified": True, + "seed_material_returned": False, + "execution_enabled": False, + "activation_available": False, + } + + +def _cluster_identifier(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("cluster identifier must be a string") + cooked = value + try: + encoded = cooked.encode("utf-8") + except UnicodeEncodeError as exc: + raise ValueError("cluster identifier must be valid UTF-8 text") from exc + if not cooked or len(encoded) > 1024: + raise ValueError("cluster identifier must contain between 1 and 1024 UTF-8 bytes") + return cooked + + +def _assignment_from_pseudonym( + public_package: Mapping[str, Any], + seed: bytes, + cluster_pseudonym: str, + *, + verification: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + verified = ( + dict(verification) + if verification is not None + else audit_shadow_campaign_artifacts(public_package) + ) + pseudonym = _sha256(cluster_pseudonym, "cluster_pseudonym") + design_hash = verified["design_binding_hash"] + seal_hash = verified["seal_hash"] + manifest = public_package["assignment_manifest"]["manifest"] + assignment_key = _hkdf_expand( + seed, + salt=bytes.fromhex(design_hash), + info=_ASSIGNMENT_INFO + bytes.fromhex(seal_hash), + ) + draw_digest = hmac.new( + assignment_key, + _DRAW_DOMAIN + + bytes.fromhex(seal_hash) + + bytes.fromhex(pseudonym) + + SHADOW_BLOCK_ID.to_bytes(8, "big"), + hashlib.sha256, + ).digest() + draw_bucket = int.from_bytes(draw_digest, "big") * SHADOW_TOTAL_ALLOCATION_BPS // (1 << 256) + chosen: Optional[Mapping[str, Any]] = None + cumulative = 0 + for arm in manifest["whole_policy_arms"]: + cumulative += int(arm["allocation_bps"]) + if draw_bucket < cumulative: + chosen = arm + break + if chosen is None: + raise ValueError("whole-policy allocation is incomplete") + reveal_payload = { + "schema_version": SHADOW_ASSIGNMENT_REVEAL_SCHEMA_VERSION, + "assignment": { + "campaign_id": verified["campaign_id"], + "seal_hash": seal_hash, + "manifest_hash": verified["manifest_hash"], + "cluster_pseudonym": pseudonym, + "block_id": SHADOW_BLOCK_ID, + "assignment_algorithm": SHADOW_ASSIGNMENT_ALGORITHM, + "arm_id": chosen["arm_id"], + "arm_hash": chosen["arm_hash"], + "arm_probability_bps": chosen["allocation_bps"], + "draw_bucket_bps": draw_bucket, + "draw_hash": hashlib.sha256(_DRAW_HASH_DOMAIN + draw_digest).hexdigest(), + "boundaries": dict(_BOUNDARIES), + }, + } + reveal = { + **reveal_payload, + "assignment_reveal_hash": _domain_hash( + _ASSIGNMENT_REVEAL_HASH_DOMAIN, + reveal_payload, + "shadow assignment reveal", + ), + } + commitment_payload = { + "schema_version": SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION, + "commitment": { + "campaign_id": verified["campaign_id"], + "seal_hash": seal_hash, + "manifest_hash": verified["manifest_hash"], + "cluster_pseudonym": pseudonym, + "block_id": SHADOW_BLOCK_ID, + "assignment_algorithm": SHADOW_ASSIGNMENT_ALGORITHM, + "seed_commitment": verified["seed_commitment"], + "assignment_reveal_commitment": reveal["assignment_reveal_hash"], + "chosen_arm_withheld_until_reveal": True, + "boundaries": dict(_BOUNDARIES), + }, + } + commitment = { + **commitment_payload, + "commitment_hash": _domain_hash( + _ASSIGNMENT_COMMITMENT_HASH_DOMAIN, + commitment_payload, + "shadow assignment commitment", + ), + } + return {"commitment": commitment, "reveal": reveal} + + +def prepare_shadow_assignment_commitment( + public_package: Any, + seed_capsule: Any, + cluster_identifier: Any, +) -> Dict[str, Any]: + if not isinstance(public_package, Mapping): + raise ValueError("shadow public package must be a JSON object") + seed = _seed_from_capsule(public_package, seed_capsule) + verification = audit_shadow_campaign_artifacts(public_package) + identity_key = _hkdf_expand( + seed, + salt=bytes.fromhex(verification["design_binding_hash"]), + info=_IDENTITY_INFO + bytes.fromhex(verification["seal_hash"]), + ) + pseudonym = hmac.new( + identity_key, + _PSEUDONYM_DOMAIN + _cluster_identifier(cluster_identifier).encode("utf-8"), + hashlib.sha256, + ).hexdigest() + prepared = _assignment_from_pseudonym( + public_package, + seed, + pseudonym, + verification=verification, + ) + return { + "commitment": prepared["commitment"], + "private_reveal": prepared["reveal"], + } + + +def _closure_artifact(public_package: Mapping[str, Any], frozen_count: int) -> Dict[str, Any]: + verification = audit_shadow_campaign_artifacts(public_package) + payload = { + "schema_version": SHADOW_CAMPAIGN_CLOSURE_SCHEMA_VERSION, + "closure": { + "campaign_id": verification["campaign_id"], + "seal_hash": verification["seal_hash"], + "frozen_commitment_count": _bounded_int( + frozen_count, + "frozen_commitment_count", + minimum=0, + maximum=MAX_SHADOW_COMMITMENTS_PER_CAMPAIGN, + ), + "state": "commitments_closed", + "boundaries": dict(_BOUNDARIES), + }, + } + return {**payload, "closure_hash": _domain_hash(_CLOSURE_HASH_DOMAIN, payload, "campaign closure")} + + +def _seed_reveal_artifact(public_package: Mapping[str, Any], seed: bytes) -> Dict[str, Any]: + verification = audit_shadow_campaign_artifacts(public_package) + payload = { + "schema_version": SHADOW_SEED_REVEAL_SCHEMA_VERSION, + "reveal": { + "campaign_id": verification["campaign_id"], + "seal_hash": verification["seal_hash"], + "seed_commitment": verification["seed_commitment"], + "seed_material_base64url": _encode_seed(seed), + "seed_material_fingerprint": hashlib.sha256( + _SEED_REVEAL_HASH_DOMAIN + _seed_bytes(seed) + ).hexdigest(), + "seed_material_revealed": True, + "state": "seed_revealed", + "boundaries": dict(_BOUNDARIES), + }, + } + return { + **payload, + "seed_reveal_hash": _domain_hash(_SEED_REVEAL_HASH_DOMAIN, payload, "seed reveal"), + } + + +def _audit_assignment_commitment_artifact( + value: Any, + *, + verification: Mapping[str, Any], +) -> Dict[str, Any]: + artifact = _require_exact_keys( + value, + {"schema_version", "commitment", "commitment_hash"}, + "shadow assignment commitment", + ) + if artifact.get("schema_version") != SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION: + raise ValueError("unsupported shadow assignment commitment schema_version") + commitment = _require_exact_keys( + artifact.get("commitment"), + { + "campaign_id", + "seal_hash", + "manifest_hash", + "cluster_pseudonym", + "block_id", + "assignment_algorithm", + "seed_commitment", + "assignment_reveal_commitment", + "chosen_arm_withheld_until_reveal", + "boundaries", + }, + "shadow assignment commitment payload", + ) + if commitment.get("campaign_id") != verification.get("campaign_id"): + raise ValueError("assignment commitment campaign_id does not match its seal") + if commitment.get("seal_hash") != verification.get("seal_hash"): + raise ValueError("assignment commitment seal_hash does not match") + if commitment.get("manifest_hash") != verification.get("manifest_hash"): + raise ValueError("assignment commitment manifest_hash does not match") + if commitment.get("seed_commitment") != verification.get("seed_commitment"): + raise ValueError("assignment commitment seed commitment does not match") + _sha256(commitment.get("cluster_pseudonym"), "cluster_pseudonym") + _sha256( + commitment.get("assignment_reveal_commitment"), + "assignment_reveal_commitment", + ) + if _bounded_int( + commitment.get("block_id"), "block_id", minimum=0, maximum=0 + ) != SHADOW_BLOCK_ID: + raise ValueError("assignment commitment block_id is unsupported") + if commitment.get("assignment_algorithm") != SHADOW_ASSIGNMENT_ALGORITHM: + raise ValueError("assignment commitment algorithm is unsupported") + if commitment.get("chosen_arm_withheld_until_reveal") is not True: + raise ValueError("assignment commitment must withhold its arm") + if _canonical_json(commitment.get("boundaries")) != _canonical_json(_BOUNDARIES): + raise ValueError("assignment commitment boundaries do not match") + payload = { + "schema_version": artifact["schema_version"], + "commitment": dict(commitment), + } + expected_hash = _domain_hash( + _ASSIGNMENT_COMMITMENT_HASH_DOMAIN, + payload, + "shadow assignment commitment", + ) + supplied_hash = _sha256(artifact.get("commitment_hash"), "commitment_hash") + if not hmac.compare_digest(expected_hash, supplied_hash): + raise ValueError("assignment commitment hash does not match") + return dict(artifact) + + +def _now_us() -> int: + return time.time_ns() // 1_000 + + +class RouteShadowAssignmentRegistry: + """SQLite registry isolated from executed route decisions and outcomes.""" + + def __init__( + self, + db_path: Any, + *, + timeout_seconds: float = 30.0, + read_only: bool = False, + ) -> None: + self.db_path = Path(db_path) + if str(self.db_path) == ":memory:": + raise ValueError("RouteShadowAssignmentRegistry requires a durable filesystem path") + if not isinstance(read_only, bool): + raise ValueError("read_only must be a boolean") + self.read_only = read_only + if self.read_only: + if not self.db_path.is_file(): + raise FileNotFoundError(f"shadow registry does not exist: {self.db_path}") + else: + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self.timeout_seconds = max(1.0, float(timeout_seconds)) + if self.read_only: + self._validate_existing_schema() + else: + self._initialize() + + def _connect(self) -> sqlite3.Connection: + database = ( + f"{self.db_path.expanduser().resolve().as_uri()}?mode=ro" + if self.read_only + else str(self.db_path) + ) + connection = sqlite3.connect( + database, + timeout=self.timeout_seconds, + isolation_level=None, + uri=self.read_only, + ) + connection.row_factory = sqlite3.Row + connection.execute(f"PRAGMA busy_timeout = {int(self.timeout_seconds * 1000)}") + connection.execute("PRAGMA foreign_keys = ON") + if not self.read_only: + connection.execute("PRAGMA synchronous = FULL") + return connection + + def _validate_existing_schema(self) -> None: + connection = self._connect() + try: + version = int(connection.execute("PRAGMA user_version").fetchone()[0]) + if version != SHADOW_REGISTRY_SCHEMA_VERSION: + raise RouteShadowRegistryError( + f"unsupported shadow registry schema {version}; expected {SHADOW_REGISTRY_SCHEMA_VERSION}" + ) + metadata = connection.execute( + "SELECT value FROM shadow_registry_metadata WHERE key = 'schema_version'" + ).fetchone() + if metadata is None or str(metadata["value"]) != str( + SHADOW_REGISTRY_SCHEMA_VERSION + ): + raise RouteShadowRegistryError("shadow registry schema metadata is missing or invalid") + schema = _schema_integrity(connection) + if schema["missing_tables"]: + raise RouteShadowRegistryError("shadow registry schema tables are incomplete") + except sqlite3.Error as exc: + raise RouteShadowRegistryError("shadow registry schema is incomplete") from exc + finally: + connection.close() + + def _initialize(self) -> None: + connection = self._connect() + try: + version = int(connection.execute("PRAGMA user_version").fetchone()[0]) + if version not in (0, SHADOW_REGISTRY_SCHEMA_VERSION): + raise RouteShadowRegistryError( + f"unsupported shadow registry schema {version}; expected {SHADOW_REGISTRY_SCHEMA_VERSION}" + ) + if version == 0: + existing_objects = connection.execute( + """ + SELECT type, name + FROM sqlite_master + WHERE name NOT LIKE 'sqlite_%' + ORDER BY type, name + """ + ).fetchall() + if existing_objects: + raise RouteShadowRegistryError( + "refusing to initialize a non-empty unversioned SQLite database" + ) + # Reject future schemas before any persistent journal-mode or DDL + # mutation so an older binary cannot rewrite a newer registry. + connection.execute("PRAGMA journal_mode = WAL") + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS shadow_registry_metadata ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS shadow_campaign_seals ( + campaign_id TEXT PRIMARY KEY, + origin_bundle_hash TEXT NOT NULL UNIQUE, + committed_bundle_hash TEXT NOT NULL UNIQUE, + protocol_hash TEXT NOT NULL, + design_binding_hash TEXT NOT NULL UNIQUE, + manifest_hash TEXT NOT NULL UNIQUE, + seal_hash TEXT NOT NULL UNIQUE, + seed_commitment TEXT NOT NULL UNIQUE, + planned_cluster_ceiling INTEGER NOT NULL CHECK (planned_cluster_ceiling >= 2), + public_package_json TEXT NOT NULL, + registered_at_us INTEGER NOT NULL CHECK (registered_at_us > 0) + ); + + CREATE TABLE IF NOT EXISTS shadow_assignment_commitments ( + commitment_hash TEXT PRIMARY KEY, + campaign_id TEXT NOT NULL REFERENCES shadow_campaign_seals(campaign_id), + cluster_pseudonym TEXT NOT NULL, + block_id INTEGER NOT NULL CHECK (block_id = 0), + assignment_reveal_commitment TEXT NOT NULL, + commitment_json TEXT NOT NULL, + committed_at_us INTEGER NOT NULL CHECK (committed_at_us > 0), + UNIQUE (campaign_id, cluster_pseudonym, block_id) + ); + + CREATE TABLE IF NOT EXISTS shadow_campaign_closures ( + campaign_id TEXT PRIMARY KEY REFERENCES shadow_campaign_seals(campaign_id), + frozen_commitment_count INTEGER NOT NULL CHECK (frozen_commitment_count >= 0), + closure_hash TEXT NOT NULL UNIQUE, + closure_json TEXT NOT NULL, + closed_at_us INTEGER NOT NULL CHECK (closed_at_us > 0) + ); + + CREATE TABLE IF NOT EXISTS shadow_seed_reveals ( + campaign_id TEXT PRIMARY KEY REFERENCES shadow_campaign_seals(campaign_id), + seed_material_base64url TEXT NOT NULL, + seed_reveal_hash TEXT NOT NULL UNIQUE, + reveal_json TEXT NOT NULL, + revealed_at_us INTEGER NOT NULL CHECK (revealed_at_us > 0) + ); + + CREATE TABLE IF NOT EXISTS shadow_assignment_reveals ( + commitment_hash TEXT PRIMARY KEY + REFERENCES shadow_assignment_commitments(commitment_hash), + campaign_id TEXT NOT NULL REFERENCES shadow_campaign_seals(campaign_id), + verification_status TEXT NOT NULL CHECK (verification_status IN ('matched', 'mismatch')), + assignment_reveal_hash TEXT NOT NULL, + reveal_json TEXT NOT NULL, + verified_at_us INTEGER NOT NULL CHECK (verified_at_us > 0) + ); + + CREATE TABLE IF NOT EXISTS shadow_registry_events ( + event_sequence INTEGER PRIMARY KEY CHECK (event_sequence >= 1), + event_type TEXT NOT NULL + CHECK (event_type IN ('campaign_sealed', 'assignment_committed', 'commitments_closed', 'seed_revealed', 'assignment_verified')), + artifact_hash TEXT NOT NULL UNIQUE, + previous_event_hash TEXT, + event_hash TEXT NOT NULL UNIQUE, + event_json TEXT NOT NULL, + recorded_at_us INTEGER NOT NULL CHECK (recorded_at_us > 0) + ); + + CREATE INDEX IF NOT EXISTS shadow_commitments_campaign_order_idx + ON shadow_assignment_commitments( + campaign_id, committed_at_us, commitment_hash + ); + CREATE INDEX IF NOT EXISTS shadow_reveals_campaign_order_idx + ON shadow_assignment_reveals( + campaign_id, verified_at_us, commitment_hash + ); + + CREATE TRIGGER IF NOT EXISTS shadow_commitments_require_open_campaign + BEFORE INSERT ON shadow_assignment_commitments + WHEN EXISTS ( + SELECT 1 FROM shadow_campaign_closures c + WHERE c.campaign_id = NEW.campaign_id + ) + BEGIN + SELECT RAISE(ABORT, 'shadow campaign commitments are closed'); + END; + + CREATE TRIGGER IF NOT EXISTS shadow_seed_reveal_requires_closure + BEFORE INSERT ON shadow_seed_reveals + WHEN NOT EXISTS ( + SELECT 1 FROM shadow_campaign_closures c + WHERE c.campaign_id = NEW.campaign_id + ) + BEGIN + SELECT RAISE(ABORT, 'shadow seed reveal requires campaign closure'); + END; + + CREATE TRIGGER IF NOT EXISTS shadow_assignment_reveal_requires_seed + BEFORE INSERT ON shadow_assignment_reveals + WHEN NOT EXISTS ( + SELECT 1 FROM shadow_seed_reveals r + WHERE r.campaign_id = NEW.campaign_id + ) + BEGIN + SELECT RAISE(ABORT, 'shadow assignment reveal requires seed reveal'); + END; + + CREATE TRIGGER IF NOT EXISTS shadow_assignment_reveal_campaign_matches_commitment + BEFORE INSERT ON shadow_assignment_reveals + WHEN NOT EXISTS ( + SELECT 1 FROM shadow_assignment_commitments c + WHERE c.commitment_hash = NEW.commitment_hash + AND c.campaign_id = NEW.campaign_id + ) + BEGIN + SELECT RAISE(ABORT, 'shadow assignment reveal campaign does not match commitment'); + END; + + CREATE TRIGGER IF NOT EXISTS shadow_metadata_no_update + BEFORE UPDATE ON shadow_registry_metadata BEGIN + SELECT RAISE(ABORT, 'shadow registry metadata is append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_metadata_no_delete + BEFORE DELETE ON shadow_registry_metadata BEGIN + SELECT RAISE(ABORT, 'shadow registry metadata is append-only'); + END; + + CREATE TRIGGER IF NOT EXISTS shadow_seals_no_update + BEFORE UPDATE ON shadow_campaign_seals BEGIN + SELECT RAISE(ABORT, 'shadow campaign seals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_seals_no_delete + BEFORE DELETE ON shadow_campaign_seals BEGIN + SELECT RAISE(ABORT, 'shadow campaign seals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_commitments_no_update + BEFORE UPDATE ON shadow_assignment_commitments BEGIN + SELECT RAISE(ABORT, 'shadow assignment commitments are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_commitments_no_delete + BEFORE DELETE ON shadow_assignment_commitments BEGIN + SELECT RAISE(ABORT, 'shadow assignment commitments are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_closures_no_update + BEFORE UPDATE ON shadow_campaign_closures BEGIN + SELECT RAISE(ABORT, 'shadow campaign closures are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_closures_no_delete + BEFORE DELETE ON shadow_campaign_closures BEGIN + SELECT RAISE(ABORT, 'shadow campaign closures are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_seed_reveals_no_update + BEFORE UPDATE ON shadow_seed_reveals BEGIN + SELECT RAISE(ABORT, 'shadow seed reveals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_seed_reveals_no_delete + BEFORE DELETE ON shadow_seed_reveals BEGIN + SELECT RAISE(ABORT, 'shadow seed reveals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_assignment_reveals_no_update + BEFORE UPDATE ON shadow_assignment_reveals BEGIN + SELECT RAISE(ABORT, 'shadow assignment reveals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_assignment_reveals_no_delete + BEFORE DELETE ON shadow_assignment_reveals BEGIN + SELECT RAISE(ABORT, 'shadow assignment reveals are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_events_no_update + BEFORE UPDATE ON shadow_registry_events BEGIN + SELECT RAISE(ABORT, 'shadow registry events are append-only'); + END; + CREATE TRIGGER IF NOT EXISTS shadow_events_no_delete + BEFORE DELETE ON shadow_registry_events BEGIN + SELECT RAISE(ABORT, 'shadow registry events are append-only'); + END; + """ + ) + connection.execute( + "INSERT OR IGNORE INTO shadow_registry_metadata(key, value) VALUES ('schema_version', ?)", + (str(SHADOW_REGISTRY_SCHEMA_VERSION),), + ) + metadata = connection.execute( + "SELECT value FROM shadow_registry_metadata WHERE key = 'schema_version'" + ).fetchone() + if metadata is None or str(metadata["value"]) != str( + SHADOW_REGISTRY_SCHEMA_VERSION + ): + raise RouteShadowRegistryError( + "shadow registry schema metadata is missing or invalid" + ) + if not _schema_integrity(connection)["ok"]: + raise RouteShadowRegistryError( + "shadow registry schema objects are incomplete" + ) + connection.execute(f"PRAGMA user_version = {SHADOW_REGISTRY_SCHEMA_VERSION}") + finally: + connection.close() + + @contextmanager + def _write_transaction(self) -> Iterator[sqlite3.Connection]: + if self.read_only: + raise RouteShadowRegistryError("shadow registry was opened read-only") + connection = self._connect() + try: + connection.execute("BEGIN IMMEDIATE") + yield connection + connection.commit() + except BaseException: + connection.rollback() + raise + finally: + connection.close() + + @contextmanager + def _read_transaction(self) -> Iterator[sqlite3.Connection]: + connection = self._connect() + try: + connection.execute("BEGIN") + yield connection + connection.commit() + except BaseException: + connection.rollback() + raise + finally: + connection.close() + + @staticmethod + def _append_event( + connection: sqlite3.Connection, + event_type: str, + artifact_hash: str, + *, + recorded_at_us: int, + ) -> Dict[str, Any]: + previous = connection.execute( + "SELECT event_sequence, event_hash FROM shadow_registry_events ORDER BY event_sequence DESC LIMIT 1" + ).fetchone() + sequence = int(previous["event_sequence"]) + 1 if previous else 1 + previous_hash = str(previous["event_hash"]) if previous else None + payload = { + "schema_version": SHADOW_REGISTRY_EVENT_SCHEMA_VERSION, + "event_sequence": sequence, + "event_type": _identifier(event_type, "event_type"), + "artifact_hash": _sha256(artifact_hash, "artifact_hash"), + "previous_event_hash": previous_hash, + "boundaries": dict(_BOUNDARIES), + } + event_hash = _domain_hash(_REGISTRY_EVENT_HASH_DOMAIN, payload, "registry event") + connection.execute( + """ + INSERT INTO shadow_registry_events( + event_sequence, event_type, artifact_hash, previous_event_hash, + event_hash, event_json, recorded_at_us + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + sequence, + event_type, + artifact_hash, + previous_hash, + event_hash, + _canonical_json(payload, "registry event"), + recorded_at_us, + ), + ) + return {**payload, "event_hash": event_hash, "recorded_at_us": recorded_at_us} + + def seal_campaign(self, review_bundle: Any, seed_material: Any) -> Dict[str, Any]: + artifacts = create_shadow_campaign_artifacts(review_bundle, seed_material) + public_package = artifacts["public_package"] + capsule = artifacts["private_seed_capsule"] + verification = audit_shadow_campaign_artifacts(public_package) + design = public_package["design_binding"] + registered_at_us = _now_us() + package_json = _canonical_json(public_package, "shadow public package") + created = False + with self._write_transaction() as connection: + existing = connection.execute( + "SELECT * FROM shadow_campaign_seals WHERE origin_bundle_hash = ?", + (verification["origin_review_bundle_hash"],), + ).fetchone() + if existing: + if str(existing["public_package_json"]) != package_json: + raise ShadowRegistryConflictError( + "origin review bundle is already sealed with a different seed commitment" + ) + else: + count = int(connection.execute("SELECT COUNT(*) FROM shadow_campaign_seals").fetchone()[0]) + if count >= MAX_SHADOW_CAMPAIGNS: + raise RouteShadowRegistryError("shadow registry campaign ceiling reached") + connection.execute( + """ + INSERT INTO shadow_campaign_seals( + campaign_id, origin_bundle_hash, committed_bundle_hash, protocol_hash, + design_binding_hash, manifest_hash, seal_hash, seed_commitment, + planned_cluster_ceiling, public_package_json, registered_at_us + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + verification["campaign_id"], + verification["origin_review_bundle_hash"], + verification["committed_review_bundle_hash"], + verification["committed_protocol_hash"], + verification["design_binding_hash"], + verification["manifest_hash"], + verification["seal_hash"], + verification["seed_commitment"], + design["planned_cluster_ceiling"], + package_json, + registered_at_us, + ), + ) + self._append_event( + connection, + "campaign_sealed", + verification["seal_hash"], + recorded_at_us=registered_at_us, + ) + created = True + return { + "ok": True, + "created": created, + "public_package": public_package, + "private_seed_capsule": capsule if created else None, + "private_seed_material_persisted": False, + "seed_capsule_returned_once": created, + "execution_enabled": False, + "activation_available": False, + } + + def _campaign_row(self, connection: sqlite3.Connection, campaign_id: Any) -> sqlite3.Row: + cooked = _identifier(campaign_id, "campaign_id") + row = connection.execute( + "SELECT * FROM shadow_campaign_seals WHERE campaign_id = ?", (cooked,) + ).fetchone() + if row is None: + raise KeyError(f"shadow campaign not found: {cooked}") + return row + + def append_assignment_commitment( + self, + *, + campaign_id: Any, + seed_capsule: Any, + cluster_identifier: Any, + ) -> Dict[str, Any]: + cooked_campaign = _identifier(campaign_id, "campaign_id") + with self._read_transaction() as connection: + row = self._campaign_row(connection, cooked_campaign) + public_package = json.loads(str(row["public_package_json"])) + seed = _seed_from_capsule(public_package, seed_capsule) + identity_key = _hkdf_expand( + seed, + salt=bytes.fromhex(row["design_binding_hash"]), + info=_IDENTITY_INFO + bytes.fromhex(row["seal_hash"]), + ) + pseudonym = hmac.new( + identity_key, + _PSEUDONYM_DOMAIN + _cluster_identifier(cluster_identifier).encode("utf-8"), + hashlib.sha256, + ).hexdigest() + prepared = _assignment_from_pseudonym(public_package, seed, pseudonym) + commitment = prepared["commitment"] + commitment_json = _canonical_json(commitment, "assignment commitment") + committed_at_us = _now_us() + created = False + with self._write_transaction() as connection: + campaign = self._campaign_row(connection, cooked_campaign) + if connection.execute( + "SELECT 1 FROM shadow_campaign_closures WHERE campaign_id = ?", (cooked_campaign,) + ).fetchone(): + raise RouteShadowRegistryError("shadow campaign commitments are closed") + existing = connection.execute( + """ + SELECT * FROM shadow_assignment_commitments + WHERE campaign_id = ? AND cluster_pseudonym = ? AND block_id = ? + """, + (cooked_campaign, pseudonym, SHADOW_BLOCK_ID), + ).fetchone() + if existing: + if str(existing["commitment_json"]) != commitment_json: + raise ShadowRegistryConflictError("cluster already has a conflicting commitment") + else: + count = int( + connection.execute( + "SELECT COUNT(*) FROM shadow_assignment_commitments WHERE campaign_id = ?", + (cooked_campaign,), + ).fetchone()[0] + ) + ceiling = min(int(campaign["planned_cluster_ceiling"]), MAX_SHADOW_COMMITMENTS_PER_CAMPAIGN) + if count >= ceiling: + raise RouteShadowRegistryError("shadow campaign planned cluster ceiling reached") + connection.execute( + """ + INSERT INTO shadow_assignment_commitments( + commitment_hash, campaign_id, cluster_pseudonym, block_id, + assignment_reveal_commitment, commitment_json, committed_at_us + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + commitment["commitment_hash"], + cooked_campaign, + pseudonym, + SHADOW_BLOCK_ID, + commitment["commitment"]["assignment_reveal_commitment"], + commitment_json, + committed_at_us, + ), + ) + self._append_event( + connection, + "assignment_committed", + commitment["commitment_hash"], + recorded_at_us=committed_at_us, + ) + created = True + return { + "ok": True, + "created": created, + "commitment": commitment, + "chosen_arm_revealed": False, + "private_reveal_persisted": False, + "execution_enabled": False, + "activation_available": False, + } + + def close_campaign(self, campaign_id: Any) -> Dict[str, Any]: + cooked = _identifier(campaign_id, "campaign_id") + closed_at_us = _now_us() + with self._write_transaction() as connection: + row = self._campaign_row(connection, cooked) + public_package = json.loads(str(row["public_package_json"])) + count = int( + connection.execute( + "SELECT COUNT(*) FROM shadow_assignment_commitments WHERE campaign_id = ?", + (cooked,), + ).fetchone()[0] + ) + existing = connection.execute( + "SELECT * FROM shadow_campaign_closures WHERE campaign_id = ?", (cooked,) + ).fetchone() + if existing: + return json.loads(str(existing["closure_json"])) + closure = _closure_artifact(public_package, count) + connection.execute( + """ + INSERT INTO shadow_campaign_closures( + campaign_id, frozen_commitment_count, closure_hash, closure_json, closed_at_us + ) VALUES (?, ?, ?, ?, ?) + """, + (cooked, count, closure["closure_hash"], _canonical_json(closure), closed_at_us), + ) + self._append_event( + connection, + "commitments_closed", + closure["closure_hash"], + recorded_at_us=closed_at_us, + ) + return closure + + def reveal_seed(self, *, campaign_id: Any, seed_capsule: Any) -> Dict[str, Any]: + cooked = _identifier(campaign_id, "campaign_id") + revealed_at_us = _now_us() + with self._write_transaction() as connection: + row = self._campaign_row(connection, cooked) + public_package = json.loads(str(row["public_package_json"])) + if not connection.execute( + "SELECT 1 FROM shadow_campaign_closures WHERE campaign_id = ?", (cooked,) + ).fetchone(): + raise RouteShadowRegistryError("shadow seed cannot be revealed before closure") + seed = _seed_from_capsule(public_package, seed_capsule) + seed_encoded = _encode_seed(seed) + reveal = _seed_reveal_artifact(public_package, seed) + existing = connection.execute( + "SELECT * FROM shadow_seed_reveals WHERE campaign_id = ?", (cooked,) + ).fetchone() + if existing: + if str(existing["seed_material_base64url"]) != seed_encoded: + raise ShadowRegistryConflictError("shadow campaign seed was already revealed differently") + return json.loads(str(existing["reveal_json"])) + connection.execute( + """ + INSERT INTO shadow_seed_reveals( + campaign_id, seed_material_base64url, seed_reveal_hash, + reveal_json, revealed_at_us + ) VALUES (?, ?, ?, ?, ?) + """, + ( + cooked, + seed_encoded, + reveal["seed_reveal_hash"], + _canonical_json(reveal), + revealed_at_us, + ), + ) + self._append_event( + connection, + "seed_revealed", + reveal["seed_reveal_hash"], + recorded_at_us=revealed_at_us, + ) + return reveal + + def verify_assignment_reveals( + self, + campaign_id: Any, + *, + batch_size: int = MAX_SHADOW_VERIFY_BATCH, + ) -> Dict[str, Any]: + cooked = _identifier(campaign_id, "campaign_id") + limit = _bounded_int( + batch_size, "batch_size", minimum=1, maximum=MAX_SHADOW_VERIFY_BATCH + ) + verified_at_us = _now_us() + invalid_commitment_artifacts = 0 + with self._write_transaction() as connection: + row = self._campaign_row(connection, cooked) + public_package = json.loads(str(row["public_package_json"])) + package_verification = audit_shadow_campaign_artifacts(public_package) + commitment_count = int( + connection.execute( + "SELECT COUNT(*) FROM shadow_assignment_commitments WHERE campaign_id = ?", + (cooked,), + ).fetchone()[0] + ) + closure_row = connection.execute( + "SELECT * FROM shadow_campaign_closures WHERE campaign_id = ?", (cooked,) + ).fetchone() + if closure_row is None: + raise RouteShadowRegistryError( + "shadow commitments must be closed before assignment verification" + ) + try: + expected_closure = _closure_artifact(public_package, commitment_count) + if ( + str(closure_row["closure_json"]) + != _canonical_json(expected_closure, "campaign closure") + or str(closure_row["closure_hash"]) + != expected_closure["closure_hash"] + or int(closure_row["frozen_commitment_count"]) + != commitment_count + ): + raise ValueError("closure artifact does not match frozen commitments") + except (KeyError, TypeError, ValueError) as exc: + raise RouteShadowRegistryError( + "shadow verification preflight found an invalid campaign closure" + ) from exc + seed_row = connection.execute( + "SELECT * FROM shadow_seed_reveals WHERE campaign_id = ?", (cooked,) + ).fetchone() + if seed_row is None: + raise RouteShadowRegistryError("shadow seed must be revealed before assignment verification") + try: + seed = _decode_seed(seed_row["seed_material_base64url"]) + expected_seed_reveal = _seed_reveal_artifact(public_package, seed) + if ( + str(seed_row["reveal_json"]) + != _canonical_json(expected_seed_reveal, "seed reveal") + or str(seed_row["seed_reveal_hash"]) + != expected_seed_reveal["seed_reveal_hash"] + ): + raise ValueError("seed reveal artifact does not reconstruct") + except (KeyError, TypeError, ValueError) as exc: + raise RouteShadowRegistryError( + "shadow verification preflight found an invalid seed reveal" + ) from exc + pending = connection.execute( + """ + SELECT c.* + FROM shadow_assignment_commitments c + LEFT JOIN shadow_assignment_reveals r ON r.commitment_hash = c.commitment_hash + WHERE c.campaign_id = ? AND r.commitment_hash IS NULL + ORDER BY c.committed_at_us, c.commitment_hash + LIMIT ? + """, + (cooked, limit), + ).fetchall() + matched = 0 + mismatched = 0 + for commitment_row in pending: + commitment_artifact_ok = False + try: + commitment_text = str(commitment_row["commitment_json"]) + commitment_value = json.loads(commitment_text) + if ( + _canonical_json( + commitment_value, "stored assignment commitment" + ) + != commitment_text + ): + raise ValueError("stored commitment is not canonical") + audited = _audit_assignment_commitment_artifact( + commitment_value, + verification=package_verification, + ) + commitment_payload = audited["commitment"] + row_projection = { + "commitment_hash": audited["commitment_hash"], + "campaign_id": commitment_payload["campaign_id"], + "cluster_pseudonym": commitment_payload["cluster_pseudonym"], + "block_id": commitment_payload["block_id"], + "assignment_reveal_commitment": commitment_payload[ + "assignment_reveal_commitment" + ], + } + for field, expected_value in row_projection.items(): + actual_value: Any = commitment_row[field] + if field == "block_id": + actual_value = int(actual_value) + else: + actual_value = str(actual_value) + if actual_value != expected_value: + raise ValueError( + f"stored commitment column {field} does not match" + ) + commitment_artifact_ok = True + except (KeyError, TypeError, ValueError, json.JSONDecodeError): + invalid_commitment_artifacts += 1 + prepared = _assignment_from_pseudonym( + public_package, + seed, + str(commitment_row["cluster_pseudonym"]), + verification=package_verification, + ) + expected_commitment = prepared["commitment"] + reveal = prepared["reveal"] + status = ( + "matched" + if commitment_artifact_ok + and hmac.compare_digest( + expected_commitment["commitment_hash"], + str(commitment_row["commitment_hash"]), + ) + and hmac.compare_digest( + reveal["assignment_reveal_hash"], + str(commitment_row["assignment_reveal_commitment"]), + ) + else "mismatch" + ) + if status == "matched": + matched += 1 + else: + mismatched += 1 + connection.execute( + """ + INSERT INTO shadow_assignment_reveals( + commitment_hash, campaign_id, verification_status, + assignment_reveal_hash, reveal_json, verified_at_us + ) VALUES (?, ?, ?, ?, ?, ?) + """, + ( + commitment_row["commitment_hash"], + cooked, + status, + reveal["assignment_reveal_hash"], + _canonical_json(reveal), + verified_at_us, + ), + ) + self._append_event( + connection, + "assignment_verified", + reveal["assignment_reveal_hash"], + recorded_at_us=verified_at_us, + ) + remaining = int( + connection.execute( + """ + SELECT COUNT(*) + FROM shadow_assignment_commitments c + LEFT JOIN shadow_assignment_reveals r ON r.commitment_hash = c.commitment_hash + WHERE c.campaign_id = ? AND r.commitment_hash IS NULL + """, + (cooked,), + ).fetchone()[0] + ) + total_mismatched = int( + connection.execute( + """ + SELECT COUNT(*) + FROM shadow_assignment_reveals + WHERE campaign_id = ? AND verification_status = 'mismatch' + """, + (cooked,), + ).fetchone()[0] + ) + processing_complete = remaining == 0 + campaign_audit_performed = processing_complete + campaign_audit_ok: Optional[bool] = None + if campaign_audit_performed: + campaign_audit_ok = bool(self.snapshot(cooked)["ok"]) + batch_integrity_ok = invalid_commitment_artifacts == 0 + verification_complete = ( + processing_complete + and total_mismatched == 0 + and campaign_audit_ok is True + ) + return { + "ok": total_mismatched == 0 + and batch_integrity_ok + and (campaign_audit_ok is not False), + "campaign_id": cooked, + "processed": len(pending), + "matched": matched, + "mismatched": mismatched, + "total_mismatched": total_mismatched, + "invalid_commitment_artifacts": invalid_commitment_artifacts, + "remaining": remaining, + "processing_complete": processing_complete, + "verification_complete": verification_complete, + "complete": verification_complete, + "campaign_audit_performed": campaign_audit_performed, + "campaign_audit_ok": campaign_audit_ok, + "execution_enabled": False, + "activation_available": False, + } + + @staticmethod + def _verify_event_chain(rows: Sequence[sqlite3.Row]) -> Dict[str, Any]: + previous_hash: Optional[str] = None + expected_sequence = 1 + for row in rows: + try: + payload = json.loads(str(row["event_json"])) + if not isinstance(payload, Mapping): + raise ValueError("event payload is not an object") + if int(row["event_sequence"]) != expected_sequence: + return {"ok": False, "reason": "event_sequence_gap", "verified_events": expected_sequence - 1} + if payload.get("event_sequence") != expected_sequence: + return {"ok": False, "reason": "event_payload_sequence_mismatch", "verified_events": expected_sequence - 1} + if payload.get("event_type") != str(row["event_type"]): + return {"ok": False, "reason": "event_type_mismatch", "verified_events": expected_sequence - 1} + if payload.get("artifact_hash") != str(row["artifact_hash"]): + return {"ok": False, "reason": "event_artifact_hash_mismatch", "verified_events": expected_sequence - 1} + if payload.get("previous_event_hash") != previous_hash: + return {"ok": False, "reason": "previous_event_hash_mismatch", "verified_events": expected_sequence - 1} + expected_hash = _domain_hash(_REGISTRY_EVENT_HASH_DOMAIN, payload, "registry event") + if not hmac.compare_digest(expected_hash, str(row["event_hash"])): + return {"ok": False, "reason": "event_hash_mismatch", "verified_events": expected_sequence - 1} + except (KeyError, TypeError, ValueError, json.JSONDecodeError): + return { + "ok": False, + "reason": "event_payload_invalid", + "verified_events": expected_sequence - 1, + } + previous_hash = expected_hash + expected_sequence += 1 + return { + "ok": True, + "reason": "verified", + "verified_events": len(rows), + "head_event_hash": previous_hash, + } + + def snapshot(self, campaign_id: Optional[Any] = None) -> Dict[str, Any]: + cooked = _identifier(campaign_id, "campaign_id") if campaign_id is not None else None + with self._read_transaction() as connection: + schema_integrity = _schema_integrity(connection) + if cooked is None: + seal_rows = connection.execute( + "SELECT * FROM shadow_campaign_seals ORDER BY registered_at_us, campaign_id" + ).fetchall() + else: + seal_rows = [self._campaign_row(connection, cooked)] + campaigns: List[Dict[str, Any]] = [] + for seal in seal_rows: + cid = str(seal["campaign_id"]) + commitment_rows = connection.execute( + """ + SELECT * FROM shadow_assignment_commitments + WHERE campaign_id = ? ORDER BY committed_at_us, commitment_hash + """, + (cid,), + ).fetchall() + commitment_count = len(commitment_rows) + closure = connection.execute( + "SELECT * FROM shadow_campaign_closures WHERE campaign_id = ?", (cid,) + ).fetchone() + seed_reveal = connection.execute( + "SELECT * FROM shadow_seed_reveals WHERE campaign_id = ?", (cid,) + ).fetchone() + reveal_rows = connection.execute( + """ + SELECT r.*, c.cluster_pseudonym, c.assignment_reveal_commitment, + c.commitment_json + FROM shadow_assignment_reveals r + JOIN shadow_assignment_commitments c + ON c.commitment_hash = r.commitment_hash + WHERE r.campaign_id = ? + ORDER BY r.verified_at_us, r.commitment_hash + """, + (cid,), + ).fetchall() + revealed = len(reveal_rows) + matched = sum( + 1 for row in reveal_rows if str(row["verification_status"]) == "matched" + ) + mismatched = sum( + 1 for row in reveal_rows if str(row["verification_status"]) == "mismatch" + ) + if closure is None: + lifecycle_state = "accepting_commitments" + elif seed_reveal is None: + lifecycle_state = "commitments_closed" + elif revealed < commitment_count: + lifecycle_state = "seed_revealed" + else: + lifecycle_state = "reveal_verification_complete" + + artifact_audit_errors: List[str] = [] + public_package: Optional[Dict[str, Any]] = None + package_verification: Optional[Dict[str, Any]] = None + try: + package_text = str(seal["public_package_json"]) + decoded_package = json.loads(package_text) + if not isinstance(decoded_package, dict): + raise ValueError("stored public package is not an object") + if _canonical_json(decoded_package, "stored public package") != package_text: + raise ValueError("stored public package is not canonical") + package_verification = audit_shadow_campaign_artifacts(decoded_package) + public_package = decoded_package + expected_seal_fields = { + "campaign_id": package_verification["campaign_id"], + "origin_bundle_hash": package_verification["origin_review_bundle_hash"], + "committed_bundle_hash": package_verification["committed_review_bundle_hash"], + "protocol_hash": package_verification["committed_protocol_hash"], + "design_binding_hash": package_verification["design_binding_hash"], + "manifest_hash": package_verification["manifest_hash"], + "seal_hash": package_verification["seal_hash"], + "seed_commitment": package_verification["seed_commitment"], + "planned_cluster_ceiling": decoded_package["design_binding"][ + "planned_cluster_ceiling" + ], + } + for field, expected_value in expected_seal_fields.items(): + actual_value: Any = seal[field] + if field == "planned_cluster_ceiling": + actual_value = int(actual_value) + else: + actual_value = str(actual_value) + if actual_value != expected_value: + raise ValueError(f"stored seal column {field} does not match its package") + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + artifact_audit_errors.append(f"campaign_seal:{exc}") + + if package_verification is not None: + for commitment_row in commitment_rows: + try: + commitment_text = str(commitment_row["commitment_json"]) + commitment_value = json.loads(commitment_text) + if _canonical_json( + commitment_value, "stored assignment commitment" + ) != commitment_text: + raise ValueError("stored commitment is not canonical") + audited_commitment = _audit_assignment_commitment_artifact( + commitment_value, + verification=package_verification, + ) + payload = audited_commitment["commitment"] + row_checks = { + "commitment_hash": audited_commitment["commitment_hash"], + "campaign_id": payload["campaign_id"], + "cluster_pseudonym": payload["cluster_pseudonym"], + "block_id": payload["block_id"], + "assignment_reveal_commitment": payload[ + "assignment_reveal_commitment" + ], + } + for field, expected_value in row_checks.items(): + actual_value = commitment_row[field] + if field == "block_id": + actual_value = int(actual_value) + else: + actual_value = str(actual_value) + if actual_value != expected_value: + raise ValueError( + f"stored commitment column {field} does not match" + ) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + artifact_audit_errors.append(f"assignment_commitment:{exc}") + + if closure is not None and public_package is not None: + try: + expected_closure = _closure_artifact(public_package, commitment_count) + closure_text = str(closure["closure_json"]) + if _canonical_json(expected_closure) != closure_text: + raise ValueError("closure artifact does not match frozen population") + if str(closure["closure_hash"]) != expected_closure["closure_hash"]: + raise ValueError("closure hash column does not match") + if int(closure["frozen_commitment_count"]) != commitment_count: + raise ValueError("closure count column does not match") + except (KeyError, TypeError, ValueError) as exc: + artifact_audit_errors.append(f"campaign_closure:{exc}") + + revealed_seed: Optional[bytes] = None + if seed_reveal is not None and public_package is not None: + try: + revealed_seed = _decode_seed(seed_reveal["seed_material_base64url"]) + expected_seed_reveal = _seed_reveal_artifact( + public_package, revealed_seed + ) + reveal_text = str(seed_reveal["reveal_json"]) + if _canonical_json(expected_seed_reveal) != reveal_text: + raise ValueError("seed reveal artifact does not reconstruct") + if str(seed_reveal["seed_reveal_hash"]) != expected_seed_reveal[ + "seed_reveal_hash" + ]: + raise ValueError("seed reveal hash column does not match") + except (KeyError, TypeError, ValueError) as exc: + revealed_seed = None + artifact_audit_errors.append(f"seed_reveal:{exc}") + + arm_counts: Dict[str, int] = {} + for reveal_row in reveal_rows: + try: + reveal_text = str(reveal_row["reveal_json"]) + reveal = json.loads(reveal_text) + if revealed_seed is None or public_package is None: + raise ValueError("assignment reveal has no verified seed opening") + expected = _assignment_from_pseudonym( + public_package, + revealed_seed, + str(reveal_row["cluster_pseudonym"]), + verification=package_verification, + ) + expected_reveal = expected["reveal"] + expected_commitment = expected["commitment"] + if _canonical_json(expected_reveal) != reveal_text: + raise ValueError("assignment reveal does not reconstruct") + if str(reveal_row["assignment_reveal_hash"]) != expected_reveal[ + "assignment_reveal_hash" + ]: + raise ValueError("assignment reveal hash column does not match") + expected_status = ( + "matched" + if hmac.compare_digest( + expected_commitment["commitment_hash"], + str(reveal_row["commitment_hash"]), + ) + and hmac.compare_digest( + expected_reveal["assignment_reveal_hash"], + str(reveal_row["assignment_reveal_commitment"]), + ) + else "mismatch" + ) + if str(reveal_row["verification_status"]) != expected_status: + raise ValueError("assignment verification status does not match") + if expected_status == "matched": + arm_id = str(expected_reveal["assignment"]["arm_id"]) + arm_counts[arm_id] = arm_counts.get(arm_id, 0) + 1 + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + artifact_audit_errors.append(f"assignment_reveal:{exc}") + if mismatched: + state = "reveal_verification_failed" + elif artifact_audit_errors: + state = "artifact_verification_failed" + else: + state = lifecycle_state + campaigns.append( + { + "campaign_id": cid, + "state": state, + "lifecycle_state": lifecycle_state, + "origin_review_bundle_hash": str(seal["origin_bundle_hash"]), + "committed_review_bundle_hash": str(seal["committed_bundle_hash"]), + "protocol_hash": str(seal["protocol_hash"]), + "manifest_hash": str(seal["manifest_hash"]), + "seal_hash": str(seal["seal_hash"]), + "seed_commitment": str(seal["seed_commitment"]), + "planned_cluster_ceiling": int(seal["planned_cluster_ceiling"]), + "commitment_count": commitment_count, + "frozen_commitment_count": int(closure["frozen_commitment_count"]) if closure else None, + "seed_revealed": seed_reveal is not None, + "processed_reveal_count": revealed, + "verified_assignment_count": matched, + "matched_assignment_count": matched, + "mismatched_assignment_count": mismatched, + "whole_policy_arm_counts": arm_counts, + "artifact_audit_ok": not artifact_audit_errors, + "artifact_audit_errors": artifact_audit_errors[:8], + "boundaries": dict(_BOUNDARIES), + } + ) + events = connection.execute( + "SELECT * FROM shadow_registry_events ORDER BY event_sequence" + ).fetchall() + chain = self._verify_event_chain(events) + registered_artifacts = sorted( + (str(row["event_type"]), str(row["artifact_hash"])) for row in events + ) + evidence_artifacts: List[Tuple[str, str]] = [] + evidence_artifacts.extend( + ("campaign_sealed", str(row[0])) + for row in connection.execute( + "SELECT seal_hash FROM shadow_campaign_seals" + ).fetchall() + ) + evidence_artifacts.extend( + ("assignment_committed", str(row[0])) + for row in connection.execute( + "SELECT commitment_hash FROM shadow_assignment_commitments" + ).fetchall() + ) + evidence_artifacts.extend( + ("commitments_closed", str(row[0])) + for row in connection.execute( + "SELECT closure_hash FROM shadow_campaign_closures" + ).fetchall() + ) + evidence_artifacts.extend( + ("seed_revealed", str(row[0])) + for row in connection.execute( + "SELECT seed_reveal_hash FROM shadow_seed_reveals" + ).fetchall() + ) + evidence_artifacts.extend( + ("assignment_verified", str(row[0])) + for row in connection.execute( + "SELECT assignment_reveal_hash FROM shadow_assignment_reveals" + ).fetchall() + ) + evidence_artifacts.sort() + artifact_events = { + "ok": registered_artifacts == evidence_artifacts, + "registered_event_artifacts": len(registered_artifacts), + "evidence_artifacts": len(evidence_artifacts), + "reason": ( + "verified" + if registered_artifacts == evidence_artifacts + else "event_evidence_artifact_mismatch" + ), + } + return { + "ok": ( + schema_integrity["ok"] + and chain["ok"] + and artifact_events["ok"] + and all( + row["mismatched_assignment_count"] == 0 + and row["artifact_audit_ok"] + for row in campaigns + ) + ), + "schema_version": SHADOW_REGISTRY_SNAPSHOT_SCHEMA_VERSION, + "registry_schema_version": SHADOW_REGISTRY_SCHEMA_VERSION, + "schema_integrity": schema_integrity, + "campaign_count": len(campaigns), + "campaigns": campaigns, + "event_chain": chain, + "event_artifact_consistency": artifact_events, + "verification_level": "local_append_only_chain_without_external_anchor", + "authenticity_proof_available": False, + "trusted_timestamp_available": False, + "external_transparency_anchor_available": False, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + + +__all__ = [ + "MAX_SHADOW_VERIFY_BATCH", + "SHADOW_ASSIGNMENT_ALGORITHM", + "SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_MANIFEST_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_REVEAL_SCHEMA_VERSION", + "SHADOW_CAMPAIGN_CLOSURE_SCHEMA_VERSION", + "SHADOW_CAMPAIGN_SEAL_SCHEMA_VERSION", + "SHADOW_CANONICALIZATION", + "SHADOW_DESIGN_BINDING_SCHEMA_VERSION", + "SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION", + "SHADOW_LEGACY_BUNDLE_ENCODING", + "SHADOW_REGISTRY_SCHEMA_VERSION", + "SHADOW_REGISTRY_SNAPSHOT_SCHEMA_VERSION", + "SHADOW_SEED_CAPSULE_SCHEMA_VERSION", + "SHADOW_SEED_REVEAL_SCHEMA_VERSION", + "RouteShadowAssignmentRegistry", + "RouteShadowRegistryError", + "ShadowRegistryConflictError", + "audit_shadow_campaign_artifacts", + "audit_shadow_seed_capsule", + "build_shadow_design_binding", + "compute_shadow_seed_commitment", + "create_shadow_campaign_artifacts", + "generate_shadow_seed", + "prepare_shadow_assignment_commitment", +] diff --git a/source/route_policy_study_cli.py b/source/route_policy_study_cli.py new file mode 100644 index 000000000..3f74a40c3 --- /dev/null +++ b/source/route_policy_study_cli.py @@ -0,0 +1,174 @@ +"""Export a read-only adjacent-route rehearsal charter as JSON. + +The CLI consumes an explicit prompt-free post-filter support snapshot. It does +not discover models, sample an assignment, touch the route ledger, or execute a +route. Redirect stdout or use ``--output`` to retain the rehearsal artifact. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Dict, Optional, Sequence + +try: + from .route_policy_explorer import plan_adjacent_route_study +except ImportError: # Direct ``python source/route_policy_study_cli.py`` use. + from route_policy_explorer import plan_adjacent_route_study + + +_INPUT_KEYS = { + "baseline_mode", + "source_contract", + "post_filter_candidates", + "post_filter_exclusions", + "exploration_rate", + "planned_routes", + "scenario_confidence", + "assumed_feedback_rate", + "target_observed_labels", +} + + +def _example_payload() -> Dict[str, Any]: + return { + "baseline_mode": "collective", + "source_contract": { + "policy_id": "auto-route-v2", + "policy_version": "2.0.0", + "feature_schema_version": "route-context-v1", + "support_schema_version": "route-support-v1", + "candidate_set_hash": "a" * 64, + "distribution_hash": "b" * 64, + "outcome_contract_schema_version": "route-outcome-contract-v1", + }, + "post_filter_candidates": [ + { + "action": "off", + "estimated_cost_units": 1.0, + "estimated_model_calls": 1, + "planned_loop_steps": 0, + "latency_tier": "low", + "selected": False, + }, + { + "action": "collective", + "estimated_cost_units": 3.0, + "estimated_model_calls": 3, + "planned_loop_steps": 0, + "latency_tier": "moderate", + "selected": True, + }, + { + "action": "loop", + "estimated_cost_units": 12.0, + "estimated_model_calls": 12, + "planned_loop_steps": 4, + "latency_tier": "frontier", + "selected": False, + }, + { + "action": "collective_loop", + "estimated_cost_units": 24.0, + "estimated_model_calls": 24, + "planned_loop_steps": 4, + "latency_tier": "frontier", + "selected": False, + }, + ], + "post_filter_exclusions": [], + "exploration_rate": 0.10, + "planned_routes": 2_000, + "scenario_confidence": 0.95, + "assumed_feedback_rate": 0.30, + "target_observed_labels": 20, + } + + +def _load_payload(path: str) -> Dict[str, Any]: + if path == "-": + raw = sys.stdin.read() + else: + raw = Path(path).read_text(encoding="utf-8-sig") + try: + payload = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"input is not valid JSON: {exc.msg}") from exc + if not isinstance(payload, dict): + raise ValueError("input JSON must be an object") + unknown = set(payload) - _INPUT_KEYS + if unknown: + raise ValueError( + "input contains unsupported or non-prompt-free fields: " + + ", ".join(sorted(map(str, unknown))) + ) + missing = { + "baseline_mode", + "source_contract", + "post_filter_candidates", + "post_filter_exclusions", + } - set(payload) + if missing: + raise ValueError("input is missing required fields: " + ", ".join(sorted(missing))) + return payload + + +def build_plan(payload: Dict[str, Any]) -> Dict[str, Any]: + """Validate the CLI payload and return the shared pure planner result.""" + + return plan_adjacent_route_study( + payload["baseline_mode"], + payload["post_filter_candidates"], + payload["post_filter_exclusions"], + source_contract=payload["source_contract"], + exploration_rate=payload.get("exploration_rate", 0.10), + planned_routes=payload.get("planned_routes", 2_000), + scenario_confidence=payload.get("scenario_confidence", 0.95), + assumed_feedback_rate=payload.get("assumed_feedback_rate", 0.30), + target_observed_labels=payload.get("target_observed_labels", 20), + ) + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export a prompt-free, non-executing adjacent-route study rehearsal." + ) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--input", help="JSON support snapshot path, or - for stdin") + source.add_argument( + "--example", + action="store_true", + help="render a documented synthetic example; it is not live route evidence", + ) + parser.add_argument("--output", help="optional JSON output path; stdout is always supported") + parser.add_argument("--compact", action="store_true", help="emit compact canonical JSON") + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + try: + payload = _example_payload() if args.example else _load_payload(str(args.input)) + plan = build_plan(payload) + rendered = json.dumps( + plan, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":") if args.compact else None, + indent=None if args.compact else 2, + ) + if args.output: + Path(args.output).write_text(rendered + "\n", encoding="utf-8") + else: + print(rendered) + except (OSError, ValueError) as exc: + print(f"route study error: {exc}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test_route_policy_explorer.py b/test_route_policy_explorer.py new file mode 100644 index 000000000..a95c1c3f8 --- /dev/null +++ b/test_route_policy_explorer.py @@ -0,0 +1,646 @@ +import builtins +import copy +import json +import math + +import pytest + +from source.route_policy_explorer import ( + ACTIVATION_BLOCKERS, + ASSIGNMENT_ALGORITHM, + AUTO_AGENT_MODE_ORDER, + DEFAULT_ASSUMED_FEEDBACK_RATE, + DEFAULT_EXPLORATION_RATE, + DEFAULT_TARGET_OBSERVED_LABELS, + MIN_POSITIVE_EXPLORATION_PROBABILITY, + REHEARSAL_RECEIPT_SCHEMA_VERSION, + REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION, + STUDY_ASSIGNMENT_SCHEMA_VERSION, + STUDY_ID, + STUDY_PLAN_SCHEMA_VERSION, + STUDY_VERSION, + _minimum_simultaneous_label_routes, + _multinomial_at_least_each_probability, + assign_adjacent_route, + plan_adjacent_route_study, + validate_adjacent_route_study, +) + + +_SOURCE_CONTRACT = { + "policy_id": "auto-route-v2", + "policy_version": "2.0.0", + "feature_schema_version": "route-context-v1", + "support_schema_version": "route-support-v1", + "candidate_set_hash": "a" * 64, + "distribution_hash": "b" * 64, + "outcome_contract_schema_version": "route-outcome-contract-v1", +} + + +_ECONOMICS = { + "off": (1.0, 1, 0, "low"), + "collective": (3.0, 3, 0, "moderate"), + "loop": (12.0, 12, 4, "frontier"), + "collective_loop": (24.0, 24, 4, "frontier"), +} + + +def _candidate(action, baseline="collective"): + cost, calls, steps, tier = _ECONOMICS[action] + return { + "action": action, + "estimated_cost_units": cost, + "estimated_model_calls": calls, + "planned_loop_steps": steps, + "latency_tier": tier, + "selected": action == baseline, + } + + +def _partition(*, baseline="collective", feasible=AUTO_AGENT_MODE_ORDER): + feasible_set = set(feasible) + candidates = [_candidate(action, baseline) for action in AUTO_AGENT_MODE_ORDER if action in feasible_set] + exclusions = [ + {"action": action, "reasons": ["capability_or_policy_filter"]} + for action in AUTO_AGENT_MODE_ORDER + if action not in feasible_set + ] + return candidates, exclusions + + +def _plan( + *, + baseline="collective", + feasible=AUTO_AGENT_MODE_ORDER, + source_contract=None, + **kwargs, +): + candidates, exclusions = _partition(baseline=baseline, feasible=feasible) + return plan_adjacent_route_study( + baseline, + candidates, + exclusions, + source_contract=source_contract or _SOURCE_CONTRACT, + **kwargs, + ) + + +def test_default_plan_is_adjacent_canonical_prompt_free_and_shadow_only(): + plan = _plan() + + assert ACTIVATION_BLOCKERS == ( + "target_policy_class_not_precommitted", + "outcome_observation_process_not_validated", + "population_definition_not_precommitted", + "session_carryover_not_addressed", + "interference_not_addressed", + "stopping_rules_not_precommitted", + "preassignment_seed_commitment_not_sealed", + "external_ope_not_validated", + ) + + assert plan["schema_version"] == STUDY_PLAN_SCHEMA_VERSION + assert plan["study"] == { + "study_id": STUDY_ID, + "study_version": STUDY_VERSION, + "label": "Bounded-Exposure Adjacent-Route Rehearsal v1", + "design_kind": "bounded_exposure_adjacent_rehearsal", + } + assert len(plan["design_hash"]) == 64 + int(plan["design_hash"], 16) + + charter = plan["charter"] + assert charter["source_contract"] == _SOURCE_CONTRACT + assert charter["prompt_free_contract"] == { + "prompt_free": True, + "raw_prompt_included": False, + "raw_session_id_included": False, + "free_form_text_fields_allowed": False, + "canonical_json": "sorted_keys_compact_utf8_no_nan", + } + assert charter["enrollment"]["eligible"] is True + assert charter["enrollment"]["baseline_action"] == "collective" + assert charter["enrollment"]["adjacent_feasible_actions"] == ["off", "loop"] + design = charter["probability_design"] + assert design["decision_type"] == "randomized" + assert design["probability_stage"] == "post_filter" + assert design["requested_exploration_rate"] == DEFAULT_EXPLORATION_RATE + assert design["applied_exploration_rate"] == DEFAULT_EXPLORATION_RATE + assert design["eligible_actions"] == ["off", "collective", "loop"] + assert design["action_probabilities"] == { + "off": 0.05, + "collective": 0.9, + "loop": 0.05, + } + assert design["minimum_positive_exploration_probability"] == MIN_POSITIVE_EXPLORATION_PROBABILITY + assert sum(design["action_probabilities"].values()) == pytest.approx(1.0) + assert [row["action"] for row in charter["candidates"]] == ["off", "collective", "loop"] + assert charter["exclusions"] == [ + { + "action": "collective_loop", + "reasons": ["nonadjacent_exploration_guard"], + "post_filter_feasible": True, + } + ] + boundaries = charter["causal_boundaries"] + assert boundaries["deployment"] == "shadow_only" + assert boundaries["execution_enabled"] is False + assert boundaries["execution_performed"] is False + assert boundaries["io_performed"] is False + assert boundaries["off_policy_estimate_computed"] is False + assert boundaries["automatic_promotion_allowed"] is False + assert boundaries["baseline_performance_guarantee"] is False + assert boundaries["logging_policy_optimality_claim"] is False + assert boundaries["ledger_eligible"] is False + assert boundaries["preassignment_commitment_sealed"] is False + assert boundaries["nonce_grinding_resistant"] is False + assert boundaries["activation_blockers"] == list(ACTIVATION_BLOCKERS) + json.dumps(plan, sort_keys=True, allow_nan=False) + + +def test_plan_is_order_invariant_and_does_not_mutate_inputs(): + candidates, exclusions = _partition() + candidates.reverse() + before_candidates = copy.deepcopy(candidates) + before_exclusions = copy.deepcopy(exclusions) + + shuffled = plan_adjacent_route_study( + "collective", candidates, exclusions, source_contract=_SOURCE_CONTRACT + ) + canonical = _plan() + + assert shuffled == canonical + assert candidates == before_candidates + assert exclusions == before_exclusions + + +def test_plan_validation_survives_json_javascript_integral_float_transport(): + plan = _plan() + + def javascript_numbers(value): + if isinstance(value, dict): + return {key: javascript_numbers(item) for key, item in value.items()} + if isinstance(value, list): + return [javascript_numbers(item) for item in value] + if isinstance(value, float) and value.is_integer(): + return int(value) + return value + + transported = javascript_numbers(json.loads(json.dumps(plan))) + assert isinstance(plan["charter"]["candidates"][0]["estimated_cost_units"], float) + assert isinstance( + transported["charter"]["candidates"][0]["estimated_cost_units"], int + ) + + validated = validate_adjacent_route_study(transported) + assert validated["design_hash"] == plan["design_hash"] + assert validated["plan"] == plan + + tampered = copy.deepcopy(transported) + tampered["charter"]["candidates"][0]["estimated_model_calls"] += 1 + with pytest.raises(ValueError, match="design_hash does not match"): + validate_adjacent_route_study(tampered) + + +def test_source_contract_is_strict_and_bound_into_the_design_hash(): + baseline = _plan() + changed_contract = {**_SOURCE_CONTRACT, "policy_version": "2.0.1"} + changed = _plan(source_contract=changed_contract) + + assert baseline["charter"]["source_contract"] == _SOURCE_CONTRACT + assert changed["charter"]["source_contract"] == changed_contract + assert baseline["design_hash"] != changed["design_hash"] + + missing = dict(_SOURCE_CONTRACT) + missing.pop("candidate_set_hash") + with pytest.raises(ValueError, match="every required provenance field"): + _plan(source_contract=missing) + malformed = {**_SOURCE_CONTRACT, "distribution_hash": "ABC"} + with pytest.raises(ValueError, match="lowercase SHA-256"): + _plan(source_contract=malformed) + stale_schema = {**_SOURCE_CONTRACT, "support_schema_version": "route-support-v0"} + with pytest.raises(ValueError, match="must be route-support-v1"): + _plan(source_contract=stale_schema) + + tampered = copy.deepcopy(baseline) + tampered["charter"]["source_contract"]["policy_version"] = "9.9.9" + with pytest.raises(ValueError, match="design_hash"): + assign_adjacent_route(tampered, "source-contract-tamper-nonce") + + +def test_resource_and_traffic_forecasts_are_expected_and_not_power_claims(): + plan = _plan(planned_routes=100, scenario_confidence=0.95) + charter = plan["charter"] + traffic = charter["traffic_scenario"] + forecast = charter["resource_forecast"] + + assert traffic["analysis_type"] == "traffic_capacity_scenario_not_statistical_power" + assert traffic["expected"] == { + "routes_by_action": {"off": 5.0, "collective": 90.0, "loop": 5.0}, + "exploration_routes": 10.0, + } + assert traffic["high_probability"]["method"] == ( + "exact_binomial_marginal_upper_quantiles_with_bonferroni_union_bound" + ) + assert traffic["high_probability"]["simultaneous_confidence_at_least"] == 0.95 + assert traffic["high_probability"]["upper_exploration_routes"] >= 10 + assert all( + traffic["high_probability"]["upper_routes_by_action"][action] + >= math.floor(expected) + for action, expected in traffic["expected"]["routes_by_action"].items() + ) + assert traffic["not_power_analysis"] is True + assert traffic["power_analysis_performed"] is False + assert traffic["minimum_detectable_effect"] is None + assert traffic["sample_size_recommendation"] is None + assert traffic["policy_value_estimate"] is None + label_scenario = traffic["observed_label_scenario"] + assert label_scenario["analysis_type"] == ( + "simultaneous_alternate_label_traffic_not_power_or_mnar_correction" + ) + assert label_scenario["target_scope"] == ( + "at_least_target_observed_labels_on_every_alternate_action" + ) + assert label_scenario["alternate_actions"] == ["off", "loop"] + assert label_scenario["alternate_action_propensities"] == { + "off": 0.05, + "loop": 0.05, + } + assert label_scenario["weakest_alternate_propensity"] == 0.05 + assert label_scenario["assumed_feedback_rate"] == DEFAULT_ASSUMED_FEEDBACK_RATE + assert label_scenario["per_route_observed_label_probability_by_alternate_action"] == { + "off": 0.015, + "loop": 0.015, + } + assert ( + label_scenario["target_observed_labels_per_alternate_action"] + == DEFAULT_TARGET_OBSERVED_LABELS + ) + assert label_scenario[ + "expected_observed_labels_at_planned_routes_by_alternate_action" + ] == {"off": 1.5, "loop": 1.5} + assert label_scenario["expected_routes_for_target_by_alternate_action"] == { + "off": 1333.333333, + "loop": 1333.333333, + } + exact = label_scenario["exact_simultaneous_target"] + assert exact["method"] == "exact_joint_multinomial_tail_inversion_two_alternates" + minimum_routes = exact[ + "minimum_routes_for_target_on_every_alternate_action" + ] + assert minimum_routes > 1333 + assert exact["probability_at_minimum_routes"] >= 0.95 + assert exact["probability_at_previous_route_count"] < 0.95 + assert label_scenario["missingness_identification_performed"] is False + assert label_scenario["mnar_correction_performed"] is False + assert label_scenario["not_power_analysis"] is True + + assert forecast["baseline_per_route"] == { + "cost_units": 3.0, + "model_calls": 3, + "loop_steps": 0, + } + assert forecast["expected_per_route"] == { + "cost_units": 3.35, + "model_calls": 3.35, + "loop_steps": 0.2, + } + assert forecast["expected_for_planned_routes"] == { + "cost_units": 335.0, + "model_calls": 335.0, + "loop_steps": 20.0, + } + assert forecast["expected_increment_vs_baseline"] == { + "cost_units": 35.0, + "model_calls": 35.0, + "loop_steps": 20.0, + } + assert forecast["forecast_only"] is True + assert forecast["budget_guarantee"] is False + assert forecast["high_probability_conservative_capacity_upper"]["cost_units"] >= 335.0 + + +def test_linear_joint_recurrence_matches_exhaustive_small_multinomials(): + def exhaustive(trials, first, second, target): + remainder = 1.0 - first - second + total = 0.0 + for first_count in range(target, trials + 1): + for second_count in range(target, trials - first_count + 1): + remainder_count = trials - first_count - second_count + coefficient = math.factorial(trials) / ( + math.factorial(first_count) + * math.factorial(second_count) + * math.factorial(remainder_count) + ) + total += coefficient * (first**first_count) * (second**second_count) * ( + remainder**remainder_count + ) + return total + + for trials in range(2, 11): + for target in range(1, min(3, trials // 2) + 1): + for probabilities in ((0.05, 0.05), (0.1, 0.2), (0.25, 0.15)): + actual = _multinomial_at_least_each_probability( + trials, probabilities, target + ) + assert actual == pytest.approx( + exhaustive(trials, *probabilities, target), abs=1e-12 + ) + + +def test_maximum_label_target_keeps_exact_joint_inversion_practical(): + plan = _plan( + planned_routes=100_000, + assumed_feedback_rate=0.30, + target_observed_labels=1_000, + ) + exact = plan["charter"]["traffic_scenario"]["observed_label_scenario"][ + "exact_simultaneous_target" + ] + assert exact["method"] == "exact_joint_multinomial_tail_inversion_two_alternates" + assert exact["probability_at_minimum_routes"] >= 0.95 + assert exact["probability_at_previous_route_count"] < 0.95 + + +def test_tiny_label_probability_large_n_inversion_avoids_lgamma_cancellation(): + minimum_routes = _minimum_simultaneous_label_routes( + label_probabilities=(1e-8, 1e-8), + target_labels=10, + confidence=0.95, + ) + assert minimum_routes == 1_706_028_428 + assert _multinomial_at_least_each_probability( + minimum_routes, (1e-8, 1e-8), 10 + ) >= 0.95 + assert _multinomial_at_least_each_probability( + minimum_routes - 1, (1e-8, 1e-8), 10 + ) < 0.95 + + +def test_one_neighbor_receives_full_alternate_mass_after_post_filtering(): + plan = _plan( + baseline="off", + feasible=("off", "collective"), + ) + + assert plan["charter"]["enrollment"]["adjacent_feasible_actions"] == ["collective"] + assert plan["charter"]["probability_design"]["action_probabilities"] == { + "off": 0.9, + "collective": 0.1, + } + assert plan["charter"]["probability_design"]["minimum_positive_exploration_probability"] == 0.1 + exclusions = {row["action"]: row for row in plan["charter"]["exclusions"]} + assert exclusions["loop"]["post_filter_feasible"] is False + assert exclusions["collective_loop"]["post_filter_feasible"] is False + exact = plan["charter"]["traffic_scenario"]["observed_label_scenario"][ + "exact_simultaneous_target" + ] + assert exact["method"] == "exact_binomial_tail_inversion_single_alternate" + assert exact["probability_at_minimum_routes"] >= 0.95 + assert exact["probability_at_previous_route_count"] < 0.95 + + +def test_nonadjacent_only_support_is_not_enrolled_and_cannot_be_assigned(): + plan = _plan( + baseline="off", + feasible=("off", "loop"), + ) + + enrollment = plan["charter"]["enrollment"] + design = plan["charter"]["probability_design"] + assert enrollment == { + "eligible": False, + "reason": "no_feasible_adjacent_action", + "baseline_action": "off", + "adjacent_feasible_actions": [], + "maximum_adjacent_neighbors": 2, + } + assert design["decision_type"] == "deterministic_not_enrolled" + assert design["applied_exploration_rate"] == 0.0 + assert design["action_probabilities"] == {"off": 1.0} + assert plan["charter"]["traffic_scenario"]["expected"]["exploration_routes"] == 0.0 + label_scenario = plan["charter"]["traffic_scenario"]["observed_label_scenario"] + assert label_scenario["weakest_alternate_propensity"] is None + assert label_scenario["alternate_actions"] == [] + assert label_scenario["per_route_observed_label_probability_by_alternate_action"] == {} + assert label_scenario["expected_routes_for_target_by_alternate_action"] == {} + assert label_scenario["exact_simultaneous_target"]["method"] == ( + "not_applicable_no_alternate_actions" + ) + assert label_scenario["exact_simultaneous_target"][ + "minimum_routes_for_target_on_every_alternate_action" + ] is None + with pytest.raises(ValueError, match="not eligible"): + assign_adjacent_route(plan, "non-enrolled-assignment-nonce") + + +@pytest.mark.parametrize( + ("mutation", "match"), + [ + (lambda candidates, exclusions: candidates.append(copy.deepcopy(candidates[0])), "duplicate"), + (lambda candidates, exclusions: candidates[0].update({"prompt": "private"}), "unsupported fields"), + (lambda candidates, exclusions: candidates[0].update({"estimated_cost_units": float("nan")}), "finite"), + (lambda candidates, exclusions: candidates[0].update({"estimated_model_calls": True}), "integer"), + (lambda candidates, exclusions: candidates[0].update({"selected": True}), "selected flags"), + (lambda candidates, exclusions: exclusions.append({"action": "off", "reasons": ["capability_or_policy_filter"]}), "disjoint"), + (lambda candidates, exclusions: candidates.pop(), "partition all route actions"), + ], +) +def test_strict_candidate_partition_validation(mutation, match): + candidates, exclusions = _partition() + mutation(candidates, exclusions) + with pytest.raises(ValueError, match=match): + plan_adjacent_route_study( + "collective", candidates, exclusions, source_contract=_SOURCE_CONTRACT + ) + + +def test_strict_exclusion_reason_and_planning_parameter_validation(): + candidates, exclusions = _partition(baseline="off", feasible=("off", "collective")) + bad_reasons = copy.deepcopy(exclusions) + bad_reasons[0]["reasons"] = ["prompt said this was risky"] + with pytest.raises(ValueError, match="versioned reason codes"): + plan_adjacent_route_study( + "off", candidates, bad_reasons, source_contract=_SOURCE_CONTRACT + ) + with pytest.raises(ValueError, match="at or above 0.05"): + plan_adjacent_route_study( + "collective", + *_partition(), + source_contract=_SOURCE_CONTRACT, + exploration_rate=0.09, + ) + with pytest.raises(ValueError, match=r"\(0, 0.20\]"): + plan_adjacent_route_study( + "collective", + *_partition(), + source_contract=_SOURCE_CONTRACT, + exploration_rate=0.21, + ) + with pytest.raises(ValueError, match="planned_routes"): + plan_adjacent_route_study( + "collective", *_partition(), source_contract=_SOURCE_CONTRACT, planned_routes=0 + ) + with pytest.raises(ValueError, match="scenario_confidence"): + plan_adjacent_route_study( + "collective", + *_partition(), + source_contract=_SOURCE_CONTRACT, + scenario_confidence=1.0, + ) + with pytest.raises(ValueError, match="assumed_feedback_rate"): + plan_adjacent_route_study( + "collective", + *_partition(), + source_contract=_SOURCE_CONTRACT, + assumed_feedback_rate=0.0, + ) + with pytest.raises(ValueError, match="target_observed_labels"): + plan_adjacent_route_study( + "collective", + *_partition(), + source_contract=_SOURCE_CONTRACT, + target_observed_labels=0, + ) + + +def test_feedback_assumption_and_target_change_label_traffic_and_design_hash(): + baseline = _plan( + planned_routes=1000, + assumed_feedback_rate=0.5, + target_observed_labels=10, + ) + alternative = _plan( + planned_routes=1000, + assumed_feedback_rate=0.25, + target_observed_labels=10, + ) + + scenario = baseline["charter"]["traffic_scenario"]["observed_label_scenario"] + assert scenario["per_route_observed_label_probability_by_alternate_action"] == { + "off": 0.025, + "loop": 0.025, + } + assert scenario["expected_observed_labels_at_planned_routes_by_alternate_action"] == { + "off": 25.0, + "loop": 25.0, + } + assert scenario["expected_routes_for_target_by_alternate_action"] == { + "off": 400.0, + "loop": 400.0, + } + assert baseline["design_hash"] != alternative["design_hash"] + assert ( + scenario["exact_simultaneous_target"][ + "minimum_routes_for_target_on_every_alternate_action" + ] + < alternative["charter"]["traffic_scenario"]["observed_label_scenario"][ + "exact_simultaneous_target" + ]["minimum_routes_for_target_on_every_alternate_action"] + ) + + +def test_planned_routes_accepts_browser_maximum_of_one_hundred_thousand(): + plan = _plan( + planned_routes=100_000, + assumed_feedback_rate=1.0, + target_observed_labels=1, + ) + assert plan["charter"]["traffic_scenario"]["planned_routes"] == 100_000 + + +def test_assign_is_nonce_deterministic_and_returns_non_ledger_rehearsal_artifacts(): + plan = _plan() + nonce = "route-assignment-nonce-0001" + + first = assign_adjacent_route(plan, nonce) + second = assign_adjacent_route(copy.deepcopy(plan), nonce) + + assert first == second + assert first["schema_version"] == STUDY_ASSIGNMENT_SCHEMA_VERSION + assert first["study_id"] == STUDY_ID + assert first["study_version"] == STUDY_VERSION + assert first["design_hash"] == plan["design_hash"] + receipt = first["rehearsal_assignment_receipt"] + proposal = first["rehearsal_support_proposal"] + assert receipt["schema_version"] == REHEARSAL_RECEIPT_SCHEMA_VERSION + assert receipt["rehearsal_only"] is True + assert receipt["assignment_algorithm"] == ASSIGNMENT_ALGORITHM + assert len(receipt["assignment_nonce_hash"]) == 64 + assert len(receipt["assignment_draw_hash"]) == 64 + assert 0.0 <= float(receipt["assignment_draw_unit_interval"]) < 1.0 + assert receipt["proposed_action"] in {"off", "collective", "loop"} + assert receipt["proposed_action_probability"] == proposal["action_probabilities"][ + receipt["proposed_action"] + ] + assert nonce not in json.dumps(first, sort_keys=True) + serialized = json.dumps(first, sort_keys=True) + assert "logging_support" not in serialized + assert '"assignment_commitment":' not in serialized + + assert proposal["schema_version"] == REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION + assert proposal["decision_type"] == "rehearsal_randomized_non_ledger" + assert proposal["probability_stage"] == "post_filter_rehearsal_only" + assert proposal["source_contract"] == _SOURCE_CONTRACT + assert proposal["ledger_eligible"] is False + assert proposal["preassignment_commitment_sealed"] is False + assert proposal["nonce_grinding_resistant"] is False + assert proposal["activation_blockers"] == list(ACTIVATION_BLOCKERS) + assert first["ledger_eligible"] is False + assert first["preassignment_commitment_sealed"] is False + assert first["nonce_grinding_resistant"] is False + assert first["activation_blockers"] == list(ACTIVATION_BLOCKERS) + assert first["side_effects"] == { + "io_performed": False, + "ledger_write_performed": False, + "execution_started": False, + "model_inference_started": False, + } + assert first["causal_boundaries"]["execution_enabled"] is False + assert first["causal_boundaries"]["off_policy_estimate_computed"] is False + assert first["causal_boundaries"]["automatic_promotion_allowed"] is False + + +def test_fixed_nonce_space_reaches_baseline_and_each_adjacent_action(): + plan = _plan() + observed = {} + for index in range(10_000): + nonce = f"fixed-study-nonce-{index:05d}" + assigned = assign_adjacent_route(plan, nonce) + receipt = assigned["rehearsal_assignment_receipt"] + observed.setdefault(receipt["proposed_action"], receipt) + if set(observed) == {"off", "collective", "loop"}: + break + + assert set(observed) == {"off", "collective", "loop"} + assert observed["off"]["is_exploration_proposal"] is True + assert observed["collective"]["is_exploration_proposal"] is False + assert observed["loop"]["is_exploration_proposal"] is True + + +def test_nonce_changes_rehearsal_receipt_and_plan_tampering_fails_closed(): + plan = _plan() + first = assign_adjacent_route(plan, "route-assignment-nonce-alpha") + second = assign_adjacent_route(plan, "route-assignment-nonce-bravo") + first_receipt = first["rehearsal_assignment_receipt"] + second_receipt = second["rehearsal_assignment_receipt"] + assert first_receipt["assignment_draw_hash"] != second_receipt["assignment_draw_hash"] + assert first_receipt["assignment_nonce_hash"] != second_receipt["assignment_nonce_hash"] + + tampered = copy.deepcopy(plan) + tampered["charter"]["probability_design"]["action_probabilities"]["off"] = 0.1 + with pytest.raises(ValueError, match="design_hash"): + assign_adjacent_route(tampered, "route-assignment-nonce-alpha") + with pytest.raises(ValueError, match="between 16 and 240"): + assign_adjacent_route(plan, "short") + + +def test_planning_and_assignment_perform_no_file_io(monkeypatch): + def forbidden_open(*args, **kwargs): + raise AssertionError("planner attempted file I/O") + + monkeypatch.setattr(builtins, "open", forbidden_open) + plan = _plan(planned_routes=25) + assignment = assign_adjacent_route(plan, "file-io-free-assignment-nonce") + assert assignment["side_effects"]["io_performed"] is False diff --git a/test_route_policy_lab.py b/test_route_policy_lab.py new file mode 100644 index 000000000..33ee53b1f --- /dev/null +++ b/test_route_policy_lab.py @@ -0,0 +1,1008 @@ +from copy import deepcopy +from dataclasses import FrozenInstanceError + +import pytest + +from source.route_policy_ledger import ( + EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION, + build_logging_support_envelope, +) +from source.route_policy_lab import ( + AUTO_AGENT_MODE_ORDER, + POLICY_PROFILES, + RoutePolicyProfile, + analyze_route_policy, + ordered_allowed_modes, + propensity_readiness, + select_profile_action, +) + + +def _usage( + route_id, + *, + score, + selected, + allowed=None, + cost=1.0, + latency=100.0, + policy_extra=None, + status="completed", + success=True, + executed=None, + fingerprint_valid=True, + fingerprint_reason=None, + outcome_contract_precommitted=True, +): + policy = { + "score": score, + "selected_agent_mode": selected, + } + if allowed is not None: + policy["allowed_agent_modes"] = allowed + policy.update(policy_extra or {}) + row = { + "route_id": route_id, + "selected_agent_mode": selected, + "executed_agent_mode": selected if executed is None else executed, + "status": status, + "success": success, + "outcome_contracts_precommitted_at_begin": outcome_contract_precommitted, + "auto_agent_policy": policy, + "route_economics": { + "actual": { + "cost_units": cost, + "elapsed_ms": latency, + } + }, + } + if fingerprint_valid is not None: + row["decision_record_fingerprint_valid"] = bool(fingerprint_valid) + row["decision_record_fingerprint_reason"] = ( + fingerprint_reason + or ("verified" if fingerprint_valid else "fingerprint_mismatch") + ) + return row + + +def _feedback(route_id, quality): + return { + "route_id": route_id, + "feedback_axes": {"quality": quality}, + "rating": "up" if quality == 1 else "down", + } + + +def _stochastic_policy(*, selected, context, version="route-v2", probabilities=None): + probabilities = probabilities or {"off": 0.5, "collective": 0.5} + eligible_actions = list(probabilities) + positive_actions = [action for action, probability in probabilities.items() if probability > 0.0] + randomized = len(positive_actions) >= 2 + logging_support = build_logging_support_envelope( + { + "schema_version": "route-support-v1", + "decision_type": "randomized" if randomized else "deterministic", + "probability_stage": "post_filter", + "sampler": { + "name": "test-randomized-assignment" if randomized else "test-argmax", + "version": "1", + "exploration_rate": 0.5 if randomized else 0.0, + "assignment_unit": "route", + "assignment_commitment": ( + f"{EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION}:" + + "ab" * 32 + if randomized + else None + ), + }, + "candidates": [{"action": action} for action in eligible_actions], + "exclusions": [], + }, + eligible_modes=eligible_actions, + action_probabilities=probabilities, + chosen_mode=selected, + ) + return { + "score": 2, + "selected_agent_mode": selected, + "allowed_agent_modes": ["off", "collective"], + "policy_id": "auto-route-v2", + "policy_version": version, + "feature_schema_version": "route-context-v1", + "decision_context": context, + "decision_type": "stochastic", + "probability_stage": "post_filter", + "eligible_actions": eligible_actions, + "post_filter_action_probabilities": probabilities, + "logging_propensity": probabilities[selected], + "logging_support": logging_support, + "candidate_set_hash": logging_support["candidate_set_hash"], + "distribution_hash": logging_support["distribution_hash"], + } + + +def _durable_lifecycle(route_count: int): + return { + "lifecycle": { + "counts": { + "started": route_count, + "completed": route_count, + "failed": 0, + "inflight": 0, + } + }, + "analysis_window": {"truncated": False}, + } + + +def _balanced_randomized_evidence(route_count: int = 40): + assert route_count >= 4 and route_count % 4 == 0 + context = {"action_mode": "chat", "budget": "balanced"} + usage_rows = [] + feedback_rows = [] + midpoint = route_count // 2 + for index in range(route_count): + target_off = index < midpoint + score = 1 if target_off else 2 + within_group = index if target_off else index - midpoint + target = "off" if target_off else "collective" + selected = target if within_group % 2 == 0 else ("collective" if target == "off" else "off") + policy = _stochastic_policy(selected=selected, context=context) + policy["score"] = score + route_id = f"route-{index}" + usage_rows.append( + _usage( + route_id, + score=score, + selected=selected, + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + feedback_rows.append(_feedback(route_id, 1 if index % 3 else -1)) + return context, usage_rows, feedback_rows + + +def test_builtin_profiles_are_immutable_and_monotonic() -> None: + assert tuple(POLICY_PROFILES) == ("efficiency", "balanced", "quality_first") + assert POLICY_PROFILES["balanced"].thresholds == { + "off": 0, + "collective": 2, + "loop": 4, + "collective_loop": 5, + } + with pytest.raises(TypeError): + POLICY_PROFILES["new"] = RoutePolicyProfile("new", 1, 2, 3) + with pytest.raises(TypeError): + POLICY_PROFILES["balanced"].thresholds["loop"] = 99 + with pytest.raises(FrozenInstanceError): + POLICY_PROFILES["balanced"].loop_min_score = 99 + with pytest.raises(ValueError): + RoutePolicyProfile("bad", 4, 3, 5) + + +def test_profile_selection_preserves_neighbor_order_and_allowed_modes() -> None: + assert AUTO_AGENT_MODE_ORDER == ("off", "collective", "loop", "collective_loop") + assert ordered_allowed_modes(["loop", "off", "collective", "loop", "invalid"]) == ( + "off", + "collective", + "loop", + ) + assert select_profile_action(1, AUTO_AGENT_MODE_ORDER, "balanced") == "off" + assert select_profile_action(2, AUTO_AGENT_MODE_ORDER, "balanced") == "collective" + assert select_profile_action(4, AUTO_AGENT_MODE_ORDER, "balanced") == "loop" + assert select_profile_action(5, AUTO_AGENT_MODE_ORDER, "balanced") == "collective_loop" + assert select_profile_action(4, ["off", "collective", "collective_loop"], "balanced") == "collective" + assert select_profile_action(6, ["off", "collective", "collective_loop"], "balanced") == "collective_loop" + assert select_profile_action(3, AUTO_AGENT_MODE_ORDER, "efficiency") == "collective" + assert select_profile_action(3, AUTO_AGENT_MODE_ORDER, "quality_first") == "loop" + assert select_profile_action("not-a-score", [], "balanced") == "off" + + +def test_exact_join_support_candidate_agreement_and_matched_observed_metrics() -> None: + usage_rows = [ + _usage( + "route-1", + score=2, + selected="collective", + allowed=list(AUTO_AGENT_MODE_ORDER), + cost=2.0, + latency=120.0, + ), + _usage( + "route-2", + score=4, + selected="off", + allowed=list(AUTO_AGENT_MODE_ORDER), + cost=1.0, + latency=60.0, + ), + _usage("route-3", score=5, selected="collective_loop", cost=8.0, latency=900.0), + {"route_id": "", "selected_agent_mode": "off"}, + ] + feedback_rows = [ + _feedback("route-1", 1), + { + "route_id": "route-2", + "feedback_axes": {"quality": None, "cost_pressure": 1}, + "rating": "down", + }, + _feedback("feedback-only", -1), + ] + + report = analyze_route_policy(usage_rows, feedback_rows, profile="balanced") + + assert report["analysis_kind"] == "associational_matched_route_replay" + assert report["causal_interpretation"]["causal"] is False + assert report["support"]["usage"] == { + "rows": 4, + "rows_with_route_id": 3, + "rows_without_route_id": 1, + "unique_route_ids": 3, + "duplicate_route_ids": 0, + } + assert report["support"]["exact_joined_route_ids"] == 2 + assert report["support"]["exact_usage_join_coverage"] == pytest.approx(2 / 3, abs=1e-6) + assert report["support"]["exact_feedback_join_coverage"] == pytest.approx(2 / 3, abs=1e-6) + assert report["candidate_action_agreement"] == { + "evaluable_routes": 2, + "matched_routes": 1, + "changed_routes": 1, + "agreement_rate": 0.5, + } + matched = report["matched_observed"] + assert matched["quality_sample_count"] == 1 + assert matched["approval_rate"] == 1.0 + assert matched["avg_cost_units"] == 2.0 + assert matched["avg_elapsed_ms"] == 120.0 + assert matched["causal"] is False + + +def test_historical_deterministic_logs_are_noncausal_and_block_promotion() -> None: + report = analyze_route_policy( + [_usage("route-1", score=4, selected="off")], + [_feedback("route-1", 1)], + profile="balanced", + ) + + readiness = report["propensity_readiness"] + assert readiness["valid_routes"] == 0 + assert readiness["off_policy_estimate_computed"] is False + assert readiness["invalid_reason_counts"]["expected_policy_version_required"] == 1 + assert readiness["invalid_reason_counts"]["expected_context_required"] == 1 + assert readiness["invalid_reason_counts"]["post_filter_probability_vector_missing"] == 1 + gate = report["promotion_gate"] + assert gate["status"] == "blocked" + assert gate["deployment"] == "shadow_only" + assert gate["reason_code"] == "no_valid_randomized_overlap" + assert gate["automatic_promotion_allowed"] is False + assert gate["requires_validated_external_ope"] is True + assert gate["causal_claim"] is False + assert report["evaluation_readiness"]["target_overlap"]["effective_sample_size"] == 0.0 + assert all("off-policy estimate" not in warning.lower() or "not" in warning.lower() for warning in report["warnings"]) + + +def test_propensity_readiness_requires_exact_version_context_and_positive_overlap() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + policy = _stochastic_policy(selected="collective", context=context) + ready = propensity_readiness( + policy, + chosen_action="collective", + target_action="off", + expected_policy_version="route-v2", + expected_context=context, + ) + assert ready["ready"] is True + assert ready["chosen_probability"] == 0.5 + assert ready["target_probability"] == 0.5 + assert ready["off_policy_estimate"] is False + + wrong_version = propensity_readiness( + policy, + chosen_action="collective", + expected_policy_version="route-v3", + expected_context=context, + ) + assert "policy_version_mismatch" in wrong_version["reasons"] + + wrong_context = propensity_readiness( + policy, + chosen_action="collective", + expected_policy_version="route-v2", + expected_context={"action_mode": "image", "budget": "balanced"}, + ) + assert "decision_context_mismatch" in wrong_context["reasons"] + + deterministic = _stochastic_policy( + selected="off", + context=context, + probabilities={"off": 1.0, "collective": 0.0}, + ) + no_overlap = propensity_readiness( + deterministic, + chosen_action="off", + target_action="collective", + expected_policy_version="route-v2", + expected_context=context, + ) + assert no_overlap["ready"] is False + assert "insufficient_randomized_support" in no_overlap["reasons"] + assert "target_action_has_no_positive_probability" in no_overlap["reasons"] + + +@pytest.mark.parametrize( + ("mutation", "reason"), + [ + ({"decision_type": "deterministic"}, "decision_not_explicitly_stochastic"), + ({"probability_stage": "pre_filter"}, "probabilities_not_post_filter"), + ({"eligible_actions": ["off"]}, "eligible_actions_invalid"), + ( + {"post_filter_action_probabilities": {"off": 0.7, "collective": 0.7}}, + "probabilities_do_not_sum_to_one", + ), + ( + {"post_filter_action_probabilities": {"off": 1.0}}, + "probability_vector_not_exactly_eligible_actions", + ), + ], +) +def test_propensity_readiness_rejects_malformed_logging(mutation, reason) -> None: + context = "chat:balanced" + policy = _stochastic_policy(selected="off", context=context) + policy.update(mutation) + result = propensity_readiness( + policy, + chosen_action="off", + expected_policy_version="route-v2", + expected_context=context, + ) + assert result["ready"] is False + assert reason in result["reasons"] + + +def test_valid_propensity_rows_only_make_external_evaluation_input_ready() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + usage_rows = [] + feedback_rows = [] + for index, selected in enumerate(("off", "collective"), start=1): + policy = _stochastic_policy(selected=selected, context=context) + usage_rows.append( + _usage( + f"route-{index}", + score=2, + selected=selected, + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + feedback_rows.append(_feedback(f"route-{index}", 1)) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + profile="balanced", + expected_policy_version="route-v2", + expected_context=context, + min_overlap_routes=2, + min_global_effective_sample_size=1, + min_per_action_effective_sample_size=1, + durable_lifecycle=_durable_lifecycle(2), + ) + + assert report["propensity_readiness"]["valid_routes"] == 2 + assert report["propensity_readiness"]["ready_as_external_ope_input"] is True + assert report["propensity_readiness"]["off_policy_estimate_computed"] is False + assert report["candidate_action_agreement"]["changed_routes"] == 1 + assert report["promotion_gate"]["status"] == "external_ope_required" + assert report["promotion_gate"]["deployment"] == "shadow_only" + assert report["promotion_gate"]["automatic_promotion_allowed"] is False + + +def test_matched_approval_interval_is_reproducible_and_descriptive_only() -> None: + usage_rows = [] + feedback_rows = [] + for index, quality in enumerate((1, 1, 1, -1), start=1): + usage_rows.append(_usage(f"route-{index}", score=2, selected="collective")) + feedback_rows.append(_feedback(f"route-{index}", quality)) + + first = analyze_route_policy(usage_rows, feedback_rows) + second = analyze_route_policy(usage_rows, feedback_rows) + interval = first["matched_observed"]["approval_interval"] + + assert first == second + assert first["matched_observed"]["approval_rate"] == 0.75 + assert interval["method"] == "descriptive_unweighted_wilson" + assert interval["causal"] is False + assert 0.0 < interval["lower_bound"] < 0.75 < interval["upper_bound"] < 1.0 + + +def test_duplicate_route_ids_use_last_row_and_are_reported() -> None: + usage_rows = [ + _usage("route-1", score=1, selected="off", cost=99.0), + _usage("route-1", score=2, selected="collective", cost=3.0), + ] + report = analyze_route_policy(usage_rows, [_feedback("route-1", 1)]) + + assert report["support"]["usage"]["duplicate_route_ids"] == 1 + assert report["candidate_action_agreement"]["matched_routes"] == 1 + assert report["matched_observed"]["avg_cost_units"] == 3.0 + + +def test_invalid_profile_and_overlap_floor_fail_closed() -> None: + with pytest.raises(ValueError, match="unknown route policy profile"): + analyze_route_policy([], [], profile="turbo") + with pytest.raises(ValueError, match="at least 1"): + analyze_route_policy([], [], min_overlap_routes=0) + + +def test_probability_stage_is_required() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + policy = _stochastic_policy(selected="off", context=context) + policy.pop("probability_stage") + + result = propensity_readiness( + policy, + chosen_action="off", + target_action="collective", + expected_policy_version="route-v2", + expected_context=context, + ) + + assert result["ready"] is False + assert "probabilities_not_post_filter" in result["reasons"] + + +def test_readiness_checks_all_usage_rows_not_only_feedback_join() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + usage_rows = [] + feedback_rows = [] + for index in range(120): + route_id = f"route-{index}" + if index < 20: + selected = "collective" if index % 2 == 0 else "off" + policy = _stochastic_policy(selected=selected, context=context) + feedback_rows.append(_feedback(route_id, 1)) + else: + selected = "off" + policy = {"score": 2, "selected_agent_mode": selected} + usage_rows.append( + _usage( + route_id, + score=2, + selected=selected, + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_version="route-v2", + expected_context=context, + min_global_effective_sample_size=5, + min_per_action_effective_sample_size=5, + ) + integrity = report["evaluation_readiness"]["logging_integrity"] + + assert integrity["checked_usage_routes"] == 120 + assert integrity["valid_routes"] == 20 + assert integrity["invalid_routes"] == 100 + assert integrity["valid_rate"] == pytest.approx(1 / 6, abs=1e-6) + assert report["evaluation_readiness"]["outcome_observation"]["quality_observed_routes"] == 20 + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + assert "logging_integrity_incomplete" in report["promotion_gate"]["blocking_reason_codes"] + + +def test_zero_candidate_matches_yield_zero_overlap_ess() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + probabilities = {"off": 0.999, "collective": 0.001} + usage_rows = [] + feedback_rows = [] + for index in range(20): + route_id = f"route-{index}" + policy = _stochastic_policy( + selected="off", + context=context, + probabilities=probabilities, + ) + usage_rows.append( + _usage( + route_id, + score=2, + selected="off", + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + feedback_rows.append(_feedback(route_id, 1)) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_version="route-v2", + expected_context=context, + ) + overlap = report["evaluation_readiness"]["target_overlap"] + + assert report["evaluation_readiness"]["logging_integrity"]["valid_routes"] == 20 + assert overlap["effective_sample_size"] == 0.0 + assert overlap["nonzero_weight_routes"] == 0 + assert overlap["minimum_target_probability"] == 0.001 + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + assert "insufficient_global_overlap_ess" in report["promotion_gate"]["blocking_reason_codes"] + + +def test_per_action_overlap_blocks_when_global_ess_hides_unseen_action() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + usage_rows = [] + feedback_rows = [] + for index in range(20): + target_off = index >= 10 + score = 1 if target_off else 2 + selected = "collective" + policy = _stochastic_policy(selected=selected, context=context) + policy["score"] = score + route_id = f"route-{index}" + usage_rows.append( + _usage( + route_id, + score=score, + selected=selected, + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + feedback_rows.append(_feedback(route_id, 1)) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_version="route-v2", + expected_context=context, + min_global_effective_sample_size=5, + min_per_action_effective_sample_size=5, + ) + overlap = report["evaluation_readiness"]["target_overlap"] + + assert overlap["effective_sample_size"] == 10.0 + assert overlap["per_action"]["collective"]["effective_sample_size"] == 10.0 + assert overlap["per_action"]["off"]["effective_sample_size"] == 0.0 + assert overlap["weakest_target_action"] == "off" + assert "insufficient_per_action_overlap" in report["promotion_gate"]["blocking_reason_codes"] + + +def test_unknown_quality_missingness_blocks_external_ope_readiness() -> None: + context = {"action_mode": "chat", "budget": "balanced"} + usage_rows = [] + feedback_rows = [] + for index in range(40): + route_id = f"route-{index}" + selected = "collective" if index % 2 == 0 else "off" + policy = _stochastic_policy(selected=selected, context=context) + usage_rows.append( + _usage( + route_id, + score=2, + selected=selected, + allowed=["off", "collective"], + policy_extra=policy, + ) + ) + if index < 10: + feedback_rows.append(_feedback(route_id, 1)) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_version="route-v2", + expected_context=context, + min_global_effective_sample_size=10, + min_per_action_effective_sample_size=10, + ) + outcome = report["evaluation_readiness"]["outcome_observation"] + + assert outcome["quality_observed_routes"] == 10 + assert outcome["quality_missing_routes"] == 30 + assert outcome["observation_propensity_logged"] is False + assert outcome["ready"] is False + assert "unknown_quality_observation_process" in report["promotion_gate"]["blocking_reason_codes"] + + +def test_balanced_overlap_advances_only_to_external_ope_required() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + readiness = report["evaluation_readiness"] + + assert readiness["target_overlap"]["effective_sample_size"] == 20.0 + assert readiness["target_overlap"]["per_action"]["off"]["effective_sample_size"] == 10.0 + assert readiness["target_overlap"]["per_action"]["collective"]["effective_sample_size"] == 10.0 + assert readiness["ready_for_external_ope"] is True + assert readiness["policy_value_estimated"] is False + assert report["promotion_gate"]["status"] == "external_ope_required" + assert report["promotion_gate"]["automatic_promotion_allowed"] is False + + +def test_route_specific_context_integrity_detects_mismatch() -> None: + first_context = {"action_mode": "chat", "budget": "balanced", "score": 2} + second_context = {"action_mode": "chat", "budget": "balanced", "score": 3} + first = _stochastic_policy(selected="collective", context=first_context) + second = _stochastic_policy(selected="collective", context=second_context) + usage_rows = [ + _usage("route-1", score=2, selected="collective", policy_extra=first), + _usage("route-2", score=2, selected="collective", policy_extra=second), + ] + + report = analyze_route_policy( + usage_rows, + [_feedback("route-1", 1), _feedback("route-2", 1)], + expected_policy_version="route-v2", + expected_context_by_route_id={ + "route-1": first_context, + "route-2": {"action_mode": "chat", "budget": "balanced", "score": 99}, + }, + min_global_effective_sample_size=1, + min_per_action_effective_sample_size=1, + ) + integrity = report["evaluation_readiness"]["logging_integrity"] + + assert integrity["valid_routes"] == 1 + assert integrity["invalid_routes"] == 1 + assert integrity["invalid_reason_counts"]["decision_context_mismatch"] == 1 + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_legacy_compatibility_evidence_never_fabricates_lifecycle_readiness() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=None, + ) + + lifecycle = report["evaluation_readiness"]["lifecycle_integrity"] + assert lifecycle["durable_lifecycle_present"] is False + assert lifecycle["reconciled"] is False + assert report["promotion_gate"]["checks"]["durable_lifecycle_present"] is False + assert "durable_lifecycle_required" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_minimum_overlap_route_floor_is_enforced() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_overlap_routes=1000, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + + assert report["promotion_gate"]["checks"]["minimum_overlap_routes_met"] is False + assert "insufficient_overlap_routes" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +@pytest.mark.parametrize( + ("mutation", "reason"), + [ + ("missing", "logging_support_missing"), + ("commitment", "assignment_commitment_missing"), + ("candidate_hash", "candidate_set_hash_mismatch"), + ("distribution_hash", "distribution_hash_mismatch"), + ], +) +def test_randomized_labels_require_a_verified_support_envelope(mutation, reason) -> None: + context = {"action_mode": "chat", "budget": "balanced"} + policy = _stochastic_policy(selected="collective", context=context) + if mutation == "missing": + policy.pop("logging_support") + elif mutation == "commitment": + policy["logging_support"] = deepcopy(policy["logging_support"]) + policy["logging_support"]["sampler"]["assignment_commitment"] = None + elif mutation == "candidate_hash": + policy["logging_support"] = deepcopy(policy["logging_support"]) + policy["logging_support"]["candidate_set_hash"] = "0" * 64 + else: + policy["logging_support"] = deepcopy(policy["logging_support"]) + policy["logging_support"]["distribution_hash"] = "0" * 64 + + result = propensity_readiness( + policy, + chosen_action="collective", + target_action="off", + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + ) + + assert result["ready"] is False + assert reason in result["reasons"] + + +def test_unversioned_observation_propensity_cannot_clear_missing_quality() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + for row in feedback_rows: + row.pop("feedback_axes", None) + row["rating"] = "" + row["observation_propensity"] = 0.75 + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + observation = report["evaluation_readiness"]["outcome_observation"] + + assert observation["raw_observation_propensity_routes"] == 40 + assert observation["versioned_observation_propensity_routes"] == 0 + assert observation["observation_propensity_logged"] is False + assert observation["ready"] is False + assert "unknown_quality_observation_process" in report["promotion_gate"]["blocking_reason_codes"] + + +@pytest.mark.parametrize("duplicate_population", ["usage", "feedback"]) +def test_duplicate_route_ids_in_durable_population_fail_closed(duplicate_population) -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + if duplicate_population == "usage": + usage_rows.append(deepcopy(usage_rows[0])) + else: + feedback_rows.append(deepcopy(feedback_rows[0])) + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + population = report["evaluation_readiness"]["population_integrity"] + + assert population[f"duplicate_{duplicate_population}_route_ids"] == 1 + assert population["complete"] is False + assert report["promotion_gate"]["checks"]["population_integrity_complete"] is False + assert "duplicate_route_ids_in_durable_population" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +@pytest.mark.parametrize( + ("mutation", "diagnostic"), + [ + ("missing_score", "missing_or_invalid_score_routes"), + ("invalid_score", "missing_or_invalid_score_routes"), + ("missing_mode", "missing_or_invalid_chosen_mode_routes"), + ("invalid_mode", "missing_or_invalid_chosen_mode_routes"), + ], +) +def test_unevaluable_route_in_durable_population_fails_closed(mutation, diagnostic) -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + row = usage_rows[0] + if mutation == "missing_score": + row["auto_agent_policy"].pop("score") + elif mutation == "invalid_score": + row["auto_agent_policy"]["score"] = "not-a-score" + elif mutation == "missing_mode": + row.pop("selected_agent_mode") + row["auto_agent_policy"].pop("selected_agent_mode") + else: + row["selected_agent_mode"] = "not-a-mode" + row["auto_agent_policy"]["selected_agent_mode"] = "not-a-mode" + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + population = report["evaluation_readiness"]["population_integrity"] + + assert population["unevaluable_usage_routes"] == 1 + assert population[diagnostic] == 1 + assert population["complete"] is False + assert "unevaluable_routes_in_durable_population" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +@pytest.mark.parametrize( + ("status", "success", "completed", "failed", "inflight"), + [ + ("failed", False, 39, 1, 0), + ("inflight", None, 39, 0, 1), + ], +) +def test_quality_on_non_successful_durable_route_is_rejected( + status, success, completed, failed, inflight +) -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + usage_rows[0]["status"] = status + usage_rows[0]["success"] = success + feedback_rows[0]["route_status"] = status + lifecycle = { + "lifecycle": { + "counts": { + "started": 40, + "completed": completed, + "failed": failed, + "inflight": inflight, + } + }, + "analysis_window": {"truncated": False}, + } + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=lifecycle, + ) + outcome = report["evaluation_readiness"]["outcome_observation"] + + assert outcome["raw_quality_outcome_routes"] == 40 + assert outcome["quality_observed_routes"] == 39 + assert outcome["invalid_quality_outcome_routes"] == 1 + assert outcome["quality_on_non_successful_routes"] == 1 + assert outcome["evidence_integrity_complete"] is False + assert report["promotion_gate"]["checks"]["outcome_evidence_integrity"] is False + assert "quality_outcome_on_ineligible_route" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_chosen_and_executed_mode_mismatch_fails_execution_and_outcome_integrity() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + chosen = usage_rows[0]["selected_agent_mode"] + usage_rows[0]["executed_agent_mode"] = "collective" if chosen == "off" else "off" + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + execution = report["evaluation_readiness"]["execution_integrity"] + outcome = report["evaluation_readiness"]["outcome_observation"] + + assert execution["chosen_executed_mode_mismatch_routes"] == 1 + assert execution["unverified_successful_execution_routes"] == 1 + assert execution["complete"] is False + assert outcome["quality_on_unverified_execution_routes"] == 1 + assert report["promotion_gate"]["checks"]["execution_integrity_complete"] is False + assert "chosen_executed_mode_mismatch" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_zero_quality_outcomes_never_certify_with_versioned_observation_propensities() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + for row in feedback_rows: + row.pop("feedback_axes", None) + row["rating"] = "" + row["observation_propensity"] = 0.75 + row["observation_policy_id"] = "quality-observation-v1" + row["observation_policy_version"] = "1" + row["outcome_definition_version"] = "quality-v1" + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + min_global_effective_sample_size=20, + min_per_action_effective_sample_size=10, + durable_lifecycle=_durable_lifecycle(40), + ) + observation = report["evaluation_readiness"]["outcome_observation"] + + assert observation["observation_propensity_logged"] is True + assert observation["versioned_observation_propensity_routes"] == 40 + assert observation["quality_observed_routes"] == 0 + assert observation["has_observed_quality_outcomes"] is False + assert observation["evidence_integrity_complete"] is True + assert observation["ready"] is False + assert report["promotion_gate"]["checks"]["quality_observation_ready"] is False + assert "no_observed_quality_outcomes" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + assert report["evaluation_readiness"]["policy_value_estimated"] is False + assert report["promotion_gate"]["automatic_promotion_allowed"] is False + + +def test_invalid_decision_record_fingerprint_blocks_durable_readiness() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + usage_rows[0]["decision_record_fingerprint_valid"] = False + usage_rows[0]["decision_record_fingerprint_reason"] = "fingerprint_mismatch" + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + durable_lifecycle=_durable_lifecycle(len(usage_rows)), + ) + + integrity = report["evaluation_readiness"]["logging_integrity"] + assert integrity["decision_record_fingerprint_complete"] is False + assert integrity["decision_record_fingerprint_valid_routes"] == len(usage_rows) - 1 + assert integrity["decision_record_fingerprint_invalid_reason_counts"] == { + "fingerprint_mismatch": 1, + } + assert report["promotion_gate"]["checks"]["logging_integrity_complete"] is False + assert "decision_record_fingerprint_invalid" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_missing_precommitted_outcome_contract_blocks_durable_readiness() -> None: + context, usage_rows, feedback_rows = _balanced_randomized_evidence() + usage_rows[0]["outcome_contracts_precommitted_at_begin"] = False + + report = analyze_route_policy( + usage_rows, + feedback_rows, + expected_policy_id="auto-route-v2", + expected_policy_version="route-v2", + expected_feature_schema="route-context-v1", + expected_context=context, + durable_lifecycle=_durable_lifecycle(len(usage_rows)), + ) + + observation = report["evaluation_readiness"]["outcome_observation"] + assert observation["outcome_contract_required"] is True + assert observation["precommitted_outcome_contract_routes"] == len(usage_rows) - 1 + assert observation["missing_or_invalid_outcome_contract_routes"] == 1 + assert observation["outcome_contract_integrity_complete"] is False + assert observation["evidence_integrity_complete"] is False + assert report["promotion_gate"]["checks"]["outcome_evidence_integrity"] is False + assert "outcome_contract_not_precommitted" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + assert report["promotion_gate"]["automatic_promotion_allowed"] is False diff --git a/test_route_policy_ledger.py b/test_route_policy_ledger.py new file mode 100644 index 000000000..02c9833d0 --- /dev/null +++ b/test_route_policy_ledger.py @@ -0,0 +1,1201 @@ +import json +import sqlite3 +import uuid +from concurrent.futures import ThreadPoolExecutor + +import pytest + +import source.route_policy_ledger as route_policy_ledger +from source.route_policy_ledger import ( + DECISION_FINGERPRINT_SCHEMA_VERSION, + LEDGER_SCHEMA_VERSION, + OUTCOME_CONTRACT_SCHEMA_VERSION, + OUTCOME_MATURITY_SCHEMA_VERSION, + OUTCOME_NAMES, + DecisionNotFoundError, + EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION, + LedgerConflictError, + RoutePolicyLedger, + build_route_outcome_contracts, + hash_session_identity, +) + + +EXECUTED_ASSIGNMENT_COMMITMENT = ( + f"{EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION}:" + "1" * 64 +) + + +def _begin(ledger, session_id="session-secret", **overrides): + values = { + "session_id": session_id, + "policy_name": "auto-route", + "policy_version": "auto-route-v2", + "policy_schema_version": "2", + "decision_context": {"action_mode": "chat", "budget": "balanced"}, + "eligible_modes": ["off", "collective"], + "chosen_mode": "collective", + "action_probabilities": {"off": 0.2, "collective": 0.8}, + "logging_support": { + "schema_version": "route-support-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "test_rng", + "version": "1", + "exploration_rate": 0.2, + "assignment_unit": "route", + "assignment_commitment": EXECUTED_ASSIGNMENT_COMMITMENT, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + }, + "estimated_economics": {"estimated_model_calls": 3, "estimated_cost_units": 3.0}, + } + values.update(overrides) + return ledger.begin_decision(**values) + + +def _create_v1_database(path, *, probabilities_json='{"off":1.0,"collective":0.0}'): + route_id = str(uuid.uuid4()) + connection = sqlite3.connect(path) + connection.executescript( + """ + CREATE TABLE ledger_metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL); + CREATE TABLE session_counters ( + session_hash TEXT PRIMARY KEY, + next_sequence INTEGER NOT NULL + ); + CREATE TABLE route_decisions ( + route_id TEXT PRIMARY KEY, + session_hash TEXT NOT NULL, + session_sequence INTEGER NOT NULL, + ledger_schema_version INTEGER NOT NULL, + policy_name TEXT NOT NULL, + policy_version TEXT NOT NULL, + policy_schema_version TEXT NOT NULL, + decision_context_json TEXT NOT NULL, + eligible_modes_json TEXT NOT NULL, + action_probabilities_json TEXT NOT NULL, + chosen_mode TEXT NOT NULL, + executed_mode TEXT, + estimated_economics_json TEXT NOT NULL, + actual_economics_json TEXT, + status TEXT NOT NULL, + success INTEGER, + error_category TEXT, + error_message TEXT, + started_at REAL NOT NULL, + completed_at REAL, + UNIQUE (session_hash, session_sequence) + ); + CREATE TABLE route_feedback_revisions ( + route_id TEXT NOT NULL REFERENCES route_decisions(route_id) ON DELETE CASCADE, + revision INTEGER NOT NULL, + idempotency_key TEXT NOT NULL, + feedback_json TEXT NOT NULL, + recorded_at REAL NOT NULL, + PRIMARY KEY (route_id, revision), + UNIQUE (route_id, idempotency_key) + ); + PRAGMA user_version = 1; + """ + ) + connection.execute( + """ + INSERT INTO route_decisions( + route_id, session_hash, session_sequence, ledger_schema_version, + policy_name, policy_version, policy_schema_version, + decision_context_json, eligible_modes_json, action_probabilities_json, + chosen_mode, executed_mode, estimated_economics_json, actual_economics_json, + status, success, started_at, completed_at + ) VALUES (?, ?, 1, 1, 'auto-route', 'v1', 'context-v1', ?, ?, ?, + 'off', 'off', '{}', '{"cost_units":1.0}', 'completed', 1, 1.0, 2.0) + """, + ( + route_id, + hash_session_identity("legacy-session"), + '{"action_mode":"text","score":1,"allowed_agent_modes":["off","collective"]}', + '["off","collective"]', + probabilities_json, + ), + ) + connection.execute( + """ + INSERT INTO route_feedback_revisions( + route_id, revision, idempotency_key, feedback_json, recorded_at + ) VALUES (?, 1, 'legacy-feedback', '{"rating":"up"}', 3.0) + """, + (route_id,), + ) + connection.commit() + connection.close() + return route_id + + +def _create_v2_database(path): + route_id = _create_v1_database(path) + connection = sqlite3.connect(path) + connection.execute( + "UPDATE route_decisions SET ledger_schema_version = 2 WHERE route_id = ?", + (route_id,), + ) + connection.execute("PRAGMA user_version = 2") + connection.commit() + connection.close() + return route_id + + +def test_begin_decision_persists_provenance_json_and_hashed_session(tmp_path) -> None: + database = tmp_path / "route-policy.sqlite3" + ledger = RoutePolicyLedger(database) + + row = _begin(ledger) + + assert uuid.UUID(row["route_id"]) + assert row["session_hash"] == hash_session_identity("session-secret") + assert row["session_hash"] != "session-secret" + assert row["session_sequence"] == 1 + assert row["ledger_schema_version"] == LEDGER_SCHEMA_VERSION + assert row["policy_name"] == "auto-route" + assert row["policy_version"] == "auto-route-v2" + assert row["policy_schema_version"] == "2" + assert row["decision_context"] == {"action_mode": "chat", "budget": "balanced"} + assert row["eligible_modes"] == ["off", "collective"] + assert row["action_probabilities"] == {"collective": 0.8, "off": 0.2} + assert row["decision_type"] == "randomized" + assert row["probability_stage"] == "post_filter" + assert row["chosen_probability"] == 0.8 + assert len(row["candidate_set_hash"]) == 64 + assert len(row["distribution_hash"]) == 64 + assert len(row["decision_record_fingerprint"]) == 64 + assert row["decision_record_fingerprint_valid"] is True + assert row["decision_record_fingerprint_reason"] == "verified" + assert ( + row["logging_support"]["decision_record_fingerprint_schema_version"] + == DECISION_FINGERPRINT_SCHEMA_VERSION + ) + assert ( + row["logging_support"]["decision_record_fingerprint"] + == row["decision_record_fingerprint"] + ) + assert row["chosen_mode"] == "collective" + assert row["estimated_economics"]["estimated_model_calls"] == 3 + assert row["status"] == "inflight" + assert row["success"] is None + assert row["feedback_status"] == "unknown" + assert row["started_at_utc"].endswith("Z") + + connection = sqlite3.connect(database) + stored = connection.execute("SELECT session_hash FROM route_decisions").fetchone()[0] + columns = [column[1] for column in connection.execute("PRAGMA table_info(route_decisions)")] + journal_mode = connection.execute("PRAGMA journal_mode").fetchone()[0] + connection.close() + assert stored == row["session_hash"] + assert "session_id" not in columns + assert journal_mode.lower() == "wal" + + +def test_two_phase_success_failure_and_inflight_survive_restart(tmp_path) -> None: + database = tmp_path / "restart.sqlite3" + ledger = RoutePolicyLedger(database) + successful = _begin(ledger, route_id=str(uuid.uuid4())) + failed = _begin(ledger, route_id=str(uuid.uuid4())) + inflight = _begin(ledger, route_id=str(uuid.uuid4())) + + complete = ledger.complete_decision( + successful["route_id"], + success=True, + executed_mode="collective", + actual_economics={"model_calls": 2, "cost_units": 2.0, "elapsed_ms": 40.0}, + ) + failure = ledger.complete_decision( + failed["route_id"], + success=False, + executed_mode="collective", + actual_economics={"model_calls": 1, "cost_units": 1.0, "elapsed_ms": 20.0}, + error_category="provider_error", + error_message="upstream request failed", + ) + + assert complete["status"] == "completed" + assert complete["success"] is True + assert complete["actual_economics"]["cost_units"] == 2.0 + assert complete["completed_at_utc"].endswith("Z") + assert failure["status"] == "failed" + assert failure["success"] is False + assert failure["error_category"] == "provider_error" + + reopened = RoutePolicyLedger(database) + assert reopened.get_decision(inflight["route_id"])["status"] == "inflight" + assert reopened.get_decision(failed["route_id"])["error_message"] == "upstream request failed" + assert reopened.report()["counts"] == {"started": 3, "completed": 1, "failed": 1, "inflight": 1} + + +def test_complete_decision_is_idempotent_but_conflicts_fail_closed(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "complete.sqlite3") + decision = _begin(ledger) + arguments = { + "success": True, + "executed_mode": "collective", + "actual_economics": {"model_calls": 2}, + } + + first = ledger.complete_decision(decision["route_id"], **arguments) + second = ledger.complete_decision(decision["route_id"], **arguments) + + assert second["completed_at"] == first["completed_at"] + with pytest.raises(LedgerConflictError): + ledger.complete_decision( + decision["route_id"], + success=True, + executed_mode="off", + actual_economics={"model_calls": 1}, + ) + with pytest.raises(ValueError, match="error_category"): + ledger.complete_decision(str(uuid.uuid4()), success=False) + with pytest.raises(DecisionNotFoundError): + ledger.complete_decision(str(uuid.uuid4()), success=True) + + +def test_session_sequences_are_monotonic_and_independent(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "sequences.sqlite3") + + a1 = _begin(ledger, session_id="session-a") + b1 = _begin(ledger, session_id="session-b") + a2 = _begin(ledger, session_id="session-a") + + assert [a1["session_sequence"], a2["session_sequence"]] == [1, 2] + assert b1["session_sequence"] == 1 + assert ledger.report(session_id="session-a")["counts"]["started"] == 2 + assert ledger.report(session_id="session-b")["counts"]["started"] == 1 + + +def test_feedback_revisions_are_idempotent_and_missing_is_unknown(tmp_path) -> None: + database = tmp_path / "feedback.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + assert ledger.get_decision(decision["route_id"])["feedback_status"] == "unknown" + first = ledger.record_feedback( + decision["route_id"], {"rating": "up", "intent": "good"}, idempotency_key="ui-submit-1" + ) + duplicate = ledger.record_feedback( + decision["route_id"], {"intent": "good", "rating": "up"}, idempotency_key="ui-submit-1" + ) + second = ledger.record_feedback( + decision["route_id"], {"rating": "down", "intent": "too_slow"}, idempotency_key="ui-submit-2" + ) + + assert first["revision"] == 1 + assert first["idempotent"] is False + assert duplicate["revision"] == 1 + assert duplicate["recorded_at"] == first["recorded_at"] + assert duplicate["idempotent"] is True + assert second["revision"] == 2 + with pytest.raises(LedgerConflictError): + ledger.record_feedback( + decision["route_id"], {"rating": "down"}, idempotency_key="ui-submit-1" + ) + + reopened = RoutePolicyLedger(database) + row = reopened.get_decision(decision["route_id"]) + assert row["feedback_status"] == "known" + assert row["feedback_revision_count"] == 2 + assert row["latest_feedback"]["feedback"]["intent"] == "too_slow" + assert [item["revision"] for item in reopened.feedback_history(decision["route_id"])] == [1, 2] + + +def test_content_derived_retry_does_not_suppress_a_later_deliberate_revision(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "feedback-content-retry.sqlite3") + decision = _begin(ledger) + + first_up = ledger.record_feedback(decision["route_id"], {"rating": "up"}) + retry_up = ledger.record_feedback(decision["route_id"], {"rating": "up"}) + down = ledger.record_feedback(decision["route_id"], {"rating": "down"}) + second_up = ledger.record_feedback(decision["route_id"], {"rating": "up"}) + + assert first_up["revision"] == retry_up["revision"] == 1 + assert retry_up["idempotent"] is True + assert down["revision"] == 2 + assert second_up["revision"] == 3 + assert second_up["idempotent"] is False + assert [row["feedback"]["rating"] for row in ledger.feedback_history(decision["route_id"])] == [ + "up", + "down", + "up", + ] + + +def test_explicit_feedback_request_id_remains_retry_idempotent_across_revisions(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "feedback-explicit-retry.sqlite3") + decision = _begin(ledger) + + first = ledger.record_feedback( + decision["route_id"], {"rating": "up"}, idempotency_key="request-1" + ) + ledger.record_feedback(decision["route_id"], {"rating": "down"}) + retry = ledger.record_feedback( + decision["route_id"], {"rating": "up"}, idempotency_key="request-1" + ) + + assert retry["revision"] == first["revision"] == 1 + assert retry["recorded_at"] == first["recorded_at"] + assert retry["idempotent"] is True + assert len(ledger.feedback_history(decision["route_id"])) == 2 + + +def test_decision_feedback_projection_uses_one_snapshot_and_one_aggregate_query( + tmp_path, monkeypatch +) -> None: + ledger = RoutePolicyLedger(tmp_path / "feedback-projection.sqlite3") + decision = _begin(ledger) + ledger.record_feedback(decision["route_id"], {"rating": "up"}) + ledger.record_feedback(decision["route_id"], {"rating": "down"}) + + statements = [] + original_connect = ledger._connect + + def traced_connect(): + connection = original_connect() + connection.set_trace_callback(statements.append) + return connection + + monkeypatch.setattr(ledger, "_connect", traced_connect) + projected = ledger.get_decision(decision["route_id"]) + + feedback_queries = [ + statement + for statement in statements + if "FROM route_feedback_revisions" in statement + ] + assert projected["feedback_revision_count"] == 2 + assert projected["latest_feedback"]["revision"] == 2 + assert projected["latest_feedback"]["feedback"]["rating"] == "down" + assert statements[0] == "BEGIN" + assert len(feedback_queries) == 1 + assert "COUNT(*) AS revision_count" in feedback_queries[0] + assert "MAX(revision) AS latest_revision" in feedback_queries[0] + + +def test_report_and_snapshot_feedback_coverage_use_unknown_semantics(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "report.sqlite3") + completed = _begin(ledger) + failed = _begin(ledger) + inflight = _begin(ledger) + ledger.complete_decision(completed["route_id"], success=True) + ledger.complete_decision(failed["route_id"], success=False, error_category="runtime_error") + ledger.record_feedback(completed["route_id"], {"rating": "up"}) + ledger.record_feedback(completed["route_id"], {"rating": "down"}, idempotency_key="revision-2") + ledger.record_feedback(inflight["route_id"], {"rating": "up"}) + + report = ledger.report() + assert report["counts"] == {"started": 3, "completed": 1, "failed": 1, "inflight": 1} + assert report["feedback_coverage"] == { + "known": 2, + "unknown": 1, + "coverage_rate": pytest.approx(2 / 3, abs=1e-6), + "revision_count": 3, + "terminal_known": 1, + "terminal_unknown": 1, + "terminal_coverage_rate": 0.5, + "missing_feedback_semantics": "unknown", + } + + snapshot = ledger.snapshot(limit=2) + assert snapshot["counts"] == report["counts"] + assert len(snapshot["recent_decisions"]) == 2 + assert snapshot["recent_decisions"][0]["session_sequence"] > snapshot["recent_decisions"][1]["session_sequence"] + + +def test_concurrent_begins_allocate_unique_monotonic_session_sequences(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "concurrent.sqlite3", timeout_seconds=10) + + def begin(index): + return _begin( + ledger, + session_id="shared-session", + route_id=str(uuid.uuid4()), + decision_context={"index": index}, + ) + + with ThreadPoolExecutor(max_workers=8) as executor: + decisions = list(executor.map(begin, range(24))) + + assert sorted(row["session_sequence"] for row in decisions) == list(range(1, 25)) + assert len({row["route_id"] for row in decisions}) == 24 + assert ledger.report(session_id="shared-session")["counts"] == { + "started": 24, + "completed": 0, + "failed": 0, + "inflight": 24, + } + + +def test_concurrent_feedback_revisions_are_unique_and_complete(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "concurrent-feedback.sqlite3", timeout_seconds=10) + decision = _begin(ledger) + + def revise(index): + return ledger.record_feedback( + decision["route_id"], + {"rating": "up", "index": index}, + idempotency_key=f"revision-{index}", + ) + + with ThreadPoolExecutor(max_workers=6) as executor: + results = list(executor.map(revise, range(12))) + + assert sorted(row["revision"] for row in results) == list(range(1, 13)) + assert len(ledger.feedback_history(decision["route_id"])) == 12 + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"route_id": "not-a-uuid"}, "valid UUID"), + ({"eligible_modes": ["off", "off"]}, "duplicates"), + ({"chosen_mode": "loop"}, "present in eligible_modes"), + ({"action_probabilities": {"off": 1.0}}, "exactly the eligible"), + ({"action_probabilities": {"off": 0.7, "collective": 0.7}}, "sum to 1"), + ({"action_probabilities": {"off": 1.0, "collective": 0.0}}, "positive logged probability"), + ({"decision_context": {"score": float("nan")}}, "JSON serializable"), + ], +) +def test_begin_validation_fails_before_writing(tmp_path, overrides, message) -> None: + ledger = RoutePolicyLedger(tmp_path / "validation.sqlite3") + with pytest.raises(ValueError, match=message): + _begin(ledger, **overrides) + assert ledger.report()["counts"]["started"] == 0 + + +def test_supplied_uuid_and_duplicate_route_conflict(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "route-id.sqlite3") + supplied = uuid.uuid4().hex + first = _begin(ledger, route_id=supplied) + assert first["route_id"] == supplied + with pytest.raises(LedgerConflictError): + _begin(ledger, route_id=supplied) + assert ledger.report()["counts"]["started"] == 1 + + +def test_unknown_route_feedback_and_durable_path_validation(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "unknown.sqlite3") + with pytest.raises(DecisionNotFoundError): + ledger.record_feedback(str(uuid.uuid4()), {"rating": "up"}) + with pytest.raises(DecisionNotFoundError): + ledger.get_decision(str(uuid.uuid4())) + with pytest.raises(ValueError, match="durable filesystem path"): + RoutePolicyLedger(":memory:") + + +def test_v1_migration_preserves_lifecycle_feedback_and_backfills_support(tmp_path) -> None: + database = tmp_path / "legacy-v1.sqlite3" + route_id = _create_v1_database(database) + + ledger = RoutePolicyLedger(database) + row = ledger.get_decision(route_id) + + assert row["status"] == "completed" + assert row["success"] is True + assert row["feedback_revision_count"] == 1 + assert row["latest_feedback"]["feedback"]["rating"] == "up" + assert row["decision_type"] == "deterministic" + assert row["probability_stage"] == "post_filter" + assert row["chosen_probability"] == 1.0 + assert row["logging_support"]["migration_source"] == "ledger_schema_v1" + assert row["logging_support"]["decision_record_fingerprint"] is None + assert row["decision_record_fingerprint"] is None + assert row["decision_record_fingerprint_valid"] is False + assert row["decision_record_fingerprint_reason"] == "legacy_unverifiable" + + connection = sqlite3.connect(database) + assert connection.execute("PRAGMA user_version").fetchone()[0] == LEDGER_SCHEMA_VERSION + assert connection.execute("SELECT COUNT(*) FROM route_decision_support").fetchone()[0] == 1 + connection.close() + + +@pytest.mark.parametrize( + "probabilities_json", + [ + "{}", + "[]", + "null", + "42", + '"scalar"', + "{malformed", + '{"off":NaN,"collective":0.0}', + '{"off":Infinity,"collective":0.0}', + '{"off":1e999,"collective":0.0}', + ], +) +def test_v1_malformed_probability_vector_is_never_upgraded_to_randomized( + tmp_path, probabilities_json +) -> None: + database = tmp_path / "legacy-malformed.sqlite3" + route_id = _create_v1_database(database, probabilities_json=probabilities_json) + + row = RoutePolicyLedger(database).get_decision(route_id) + + assert row["action_probabilities"] == {} + assert row["decision_type"] == "legacy_unknown" + assert row["chosen_probability"] is None + assert row["logging_support"]["migration_source"] == "ledger_schema_v1" + + +def test_support_hashes_are_canonical_and_mapping_order_independent(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "hashes.sqlite3") + first = _begin(ledger) + second_support = { + "exclusions": [], + "candidates": [{"action": "off"}, {"action": "collective"}], + "sampler": { + "assignment_commitment": EXECUTED_ASSIGNMENT_COMMITMENT, + "assignment_unit": "route", + "exploration_rate": 0.2, + "version": "1", + "name": "test_rng", + }, + "probability_stage": "post_filter", + "decision_type": "randomized", + "schema_version": "route-support-v1", + } + second = _begin( + ledger, + decision_context={"budget": "balanced", "action_mode": "chat"}, + action_probabilities={"collective": 0.8, "off": 0.2}, + logging_support=second_support, + ) + + assert second["candidate_set_hash"] == first["candidate_set_hash"] + assert second["distribution_hash"] == first["distribution_hash"] + assert second["decision_record_fingerprint"] == first["decision_record_fingerprint"] + + +@pytest.mark.parametrize( + ("statement", "replacement"), + [ + ( + "UPDATE route_decisions SET decision_context_json = ? WHERE route_id = ?", + '{"action_mode":"chat","budget":"efficiency"}', + ), + ( + "UPDATE route_decisions SET policy_name = ? WHERE route_id = ?", + "tampered-policy", + ), + ( + "UPDATE route_decisions SET policy_version = ? WHERE route_id = ?", + "tampered-version", + ), + ( + "UPDATE route_decisions SET policy_schema_version = ? WHERE route_id = ?", + "tampered-schema", + ), + ( + "UPDATE route_decisions SET eligible_modes_json = ? WHERE route_id = ?", + '["collective","off"]', + ), + ( + "UPDATE route_decisions SET action_probabilities_json = ? WHERE route_id = ?", + '{"collective":0.7,"off":0.3}', + ), + ( + "UPDATE route_decisions SET chosen_mode = ? WHERE route_id = ?", + "off", + ), + ], +) +def test_decision_record_fingerprint_detects_sqlite_record_tampering( + tmp_path, statement, replacement +) -> None: + database = tmp_path / "fingerprint-tamper.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + connection.execute(statement, (replacement, decision["route_id"])) + connection.commit() + connection.close() + + tampered = ledger.get_decision(decision["route_id"]) + assert tampered["decision_record_fingerprint"] == decision["decision_record_fingerprint"] + assert tampered["decision_record_fingerprint_valid"] is False + assert tampered["decision_record_fingerprint_reason"] == "fingerprint_mismatch" + + +def test_decision_record_fingerprint_detects_support_projection_tampering(tmp_path) -> None: + database = tmp_path / "fingerprint-support-tamper.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + connection.execute( + "UPDATE route_decision_support SET candidate_set_hash = ? WHERE route_id = ?", + ("0" * 64, decision["route_id"]), + ) + connection.commit() + connection.close() + + tampered = ledger.get_decision(decision["route_id"]) + assert tampered["decision_record_fingerprint_valid"] is False + assert tampered["decision_record_fingerprint_reason"] == "support_projection_mismatch" + + +def test_existing_v2_row_without_fingerprint_is_never_retroactively_fabricated(tmp_path) -> None: + database = tmp_path / "pre-fingerprint-v2.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + envelope = json.loads( + connection.execute( + "SELECT logging_envelope_json FROM route_decision_support WHERE route_id = ?", + (decision["route_id"],), + ).fetchone()[0] + ) + envelope.pop("decision_record_fingerprint", None) + envelope.pop("decision_record_fingerprint_schema_version", None) + connection.execute( + "UPDATE route_decision_support SET logging_envelope_json = ? WHERE route_id = ?", + (json.dumps(envelope, sort_keys=True, separators=(",", ":")), decision["route_id"]), + ) + connection.commit() + connection.close() + + reopened = RoutePolicyLedger(database) + unverifiable = reopened.get_decision(decision["route_id"]) + assert unverifiable["decision_record_fingerprint"] is None + assert unverifiable["decision_record_fingerprint_valid"] is False + assert unverifiable["decision_record_fingerprint_reason"] == "missing_unverifiable" + + +def test_randomized_label_requires_assignment_commitment_and_v2_requires_vector(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "support-validation.sqlite3") + support = { + "schema_version": "route-support-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "rng", + "version": "1", + "exploration_rate": 0.1, + "assignment_unit": "route", + "assignment_commitment": None, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + } + with pytest.raises(ValueError, match="assignment commitment"): + _begin(ledger, logging_support=support) + with pytest.raises(ValueError, match="action_probabilities are required"): + _begin(ledger, action_probabilities=None) + assert ledger.report()["counts"]["started"] == 0 + + +@pytest.mark.parametrize( + ("location", "key", "value", "message"), + [ + ("root", "shadow_only", True, "shadow-only"), + ("root", "rehearsal_only", True, "rehearsal-only"), + ("root", "ledger_eligible", False, "ledger_eligible=false"), + ("sampler", "shadow_only", True, "shadow-only"), + ("sampler", "rehearsal_only", True, "rehearsal-only"), + ("sampler", "ledger_eligible", False, "ledger_eligible=false"), + ], +) +def test_executed_ledger_rejects_shadow_or_rehearsal_support_flags( + tmp_path, location, key, value, message +) -> None: + ledger = RoutePolicyLedger(tmp_path / f"shadow-flag-{location}-{key}.sqlite3") + support = { + "schema_version": "route-support-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "rng", + "version": "1", + "exploration_rate": 0.2, + "assignment_unit": "route", + "assignment_commitment": EXECUTED_ASSIGNMENT_COMMITMENT, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + } + target = support if location == "root" else support["sampler"] + target[key] = value + with pytest.raises(ValueError, match=message): + _begin(ledger, logging_support=support) + assert ledger.report()["counts"]["started"] == 0 + + +@pytest.mark.parametrize( + "commitment", + [ + "shadow:campaign:assignment", + "shadow-v1:opaque-receipt", + "route-study-shadow-assignment-commitment-v1:abc", + "route-shadow:abc", + ], +) +def test_executed_ledger_rejects_reserved_shadow_commitment_namespaces( + tmp_path, commitment +) -> None: + ledger = RoutePolicyLedger(tmp_path / "shadow-commitment.sqlite3") + support = { + "schema_version": "route-support-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "rng", + "version": "1", + "exploration_rate": 0.2, + "assignment_unit": "route", + "assignment_commitment": commitment, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + } + with pytest.raises(ValueError, match="shadow assignment commitments"): + _begin(ledger, logging_support=support) + assert ledger.report()["counts"]["started"] == 0 + + +def test_executed_ledger_rejects_real_bare_shadow_commitment_hash(tmp_path) -> None: + from source.route_policy_protocol import build_route_study_review_bundle_from_input + from source.route_policy_protocol_cli import _example_bundle_input + from source.route_policy_shadow_registry import ( + create_shadow_campaign_artifacts, + prepare_shadow_assignment_commitment, + ) + + bundle = build_route_study_review_bundle_from_input(_example_bundle_input()) + artifacts = create_shadow_campaign_artifacts(bundle, bytes(range(32))) + shadow_commitment = prepare_shadow_assignment_commitment( + artifacts["public_package"], + artifacts["private_seed_capsule"], + "ledger-boundary-cluster", + )["commitment"]["commitment_hash"] + support = { + "schema_version": "route-support-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "rng", + "version": "1", + "exploration_rate": 0.2, + "assignment_unit": "route", + "assignment_commitment": shadow_commitment, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + } + with pytest.raises(ValueError, match="route-execution-assignment-v1"): + _begin(RoutePolicyLedger(tmp_path / "bare-shadow-hash.sqlite3"), logging_support=support) + + +def test_executed_ledger_rejects_unknown_support_schema(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "support-schema.sqlite3") + support = { + "schema_version": "route-study-shadow-assignment-commitment-v1", + "decision_type": "randomized", + "probability_stage": "post_filter", + "sampler": { + "name": "rng", + "version": "1", + "exploration_rate": 0.2, + "assignment_unit": "route", + "assignment_commitment": EXECUTED_ASSIGNMENT_COMMITMENT, + }, + "candidates": [{"action": "off"}, {"action": "collective"}], + "exclusions": [], + } + with pytest.raises(ValueError, match="schema_version must be route-support-v1"): + _begin(ledger, logging_support=support) + assert ledger.report()["counts"]["started"] == 0 + + +def test_policy_evidence_snapshot_is_prompt_free_and_separates_missing_feedback(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "evidence.sqlite3") + decision = ledger.begin_decision( + session_id="private-session", + policy_name="auto-route-v2", + policy_version="2.0.0", + policy_schema_version="route-context-v1", + decision_context={ + "action_mode": "text", + "budget_profile": "balanced", + "score": 2, + "allowed_agent_modes": ["off", "collective"], + "prompt": "private raw prompt", + "nested": { + "user_prompt": "private nested prompt", + "prompt_length": 18, + }, + }, + eligible_modes=["off", "collective"], + chosen_mode="collective", + action_probabilities={"off": 0.0, "collective": 1.0}, + ) + ledger.complete_decision(decision["route_id"], success=True, executed_mode="collective") + ledger.record_feedback( + decision["route_id"], + {"rating": "up", "feedback_axes": {"quality": 1}, "observation_status": "observed"}, + ) + + snapshot = ledger.policy_evidence_snapshot( + session_id="private-session", + policy_name="auto-route-v2", + policy_version="2.0.0", + ) + + assert snapshot["analysis_window"]["included_decisions"] == 1 + assert snapshot["analysis_window"]["truncated"] is False + assert snapshot["usage_rows"][0]["auto_agent_policy"]["decision_type"] == "deterministic" + assert snapshot["usage_rows"][0]["decision_record_fingerprint_valid"] is True + assert ( + snapshot["usage_rows"][0]["auto_agent_policy"][ + "decision_record_fingerprint_reason" + ] + == "verified" + ) + assert ( + snapshot["usage_rows"][0]["auto_agent_policy"][ + "decision_record_fingerprint" + ] + == decision["decision_record_fingerprint"] + ) + assert snapshot["feedback_rows"][0]["feedback_axes"]["quality"] == 1 + assert snapshot["expected_context_by_route_id"][decision["route_id"]]["score"] == 2 + assert snapshot["expected_context_by_route_id"][decision["route_id"]]["nested"] == { + "prompt_length": 18 + } + assert "private raw prompt" not in str(snapshot) + assert "private nested prompt" not in str(snapshot) + assert "private-session" not in str(snapshot) + assert "prompt" not in snapshot["usage_rows"][0] + + +def test_route_outcome_contract_builder_is_complete_canonical_and_strict() -> None: + contracts = build_route_outcome_contracts() + + assert tuple(contracts) == OUTCOME_NAMES + assert {contract["schema_version"] for contract in contracts.values()} == { + OUTCOME_CONTRACT_SCHEMA_VERSION + } + assert all(contract["precommitted"] is True for contract in contracts.values()) + assert all(contract["commitment_source"] == "safe_default" for contract in contracts.values()) + assert all(len(contract["contract_hash"]) == 64 for contract in contracts.values()) + assert contracts["user_quality_rating"]["value_type"] == "ordinal" + assert contracts["user_quality_rating"]["unit"] == "signed_unit_interval" + + custom = {name: {} for name in OUTCOME_NAMES} + custom["user_quality_rating"] = { + "outcome_definition_version": "quality-v2", + "observation_policy_id": "quality-dialog", + "observation_policy_version": "2", + "maturity_delay_seconds": 30, + } + built = build_route_outcome_contracts(custom) + assert built["user_quality_rating"]["outcome_definition_version"] == "quality-v2" + assert built["user_quality_rating"]["maturity_delay_seconds"] == 30.0 + assert all(contract["commitment_source"] == "caller" for contract in built.values()) + + with pytest.raises(ValueError, match="exactly the four"): + build_route_outcome_contracts({"route_success": {}}) + custom["cost"] = {"unit": "dollars"} + with pytest.raises(ValueError, match="value_type or unit"): + build_route_outcome_contracts(custom) + + +def test_begin_atomically_precommits_four_append_only_outcome_contracts(tmp_path) -> None: + database = tmp_path / "contracts.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + assert set(decision["outcome_contracts"]) == set(OUTCOME_NAMES) + assert decision["outcome_contracts_precommitted_at_begin"] is True + assert decision["outcome_contracts_defaulted_at_begin"] is True + assert decision["outcome_contract_commitment_source"] == "safe_default" + assert all( + contract["contract_hash_valid"] is True + and contract["contract_hash_reason"] == "verified" + and contract["commitment_timing_valid"] is True + and contract["commitment_timing_reason"] == "verified" + for contract in decision["outcome_contracts"].values() + ) + + connection = sqlite3.connect(database) + assert connection.execute("SELECT COUNT(*) FROM route_outcome_contracts").fetchone()[0] == 4 + with pytest.raises(sqlite3.IntegrityError, match="append-only"): + connection.execute( + "UPDATE route_outcome_contracts SET unit = 'changed' WHERE route_id = ?", + (decision["route_id"],), + ) + connection.close() + + bad = {name: {} for name in OUTCOME_NAMES} + bad["latency"] = {"maturity_delay_seconds": float("inf")} + with pytest.raises(ValueError, match="finite non-negative"): + _begin(ledger, route_id=str(uuid.uuid4()), outcome_contracts=bad) + assert ledger.report()["counts"]["started"] == 1 + + +def test_completion_emits_three_events_with_nested_economics_compatibility(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "completion-events.sqlite3") + decision = _begin(ledger) + completed = ledger.complete_decision( + decision["route_id"], + success=True, + executed_mode="collective", + actual_economics={"actual": {"cost_units": 2.5, "elapsed_ms": 41.0}}, + ) + + events = {event["outcome_name"]: event for event in completed["outcome_events"]} + assert set(events) == {"route_success", "cost", "latency"} + assert events["route_success"]["observation_status"] == "observed" + assert events["route_success"]["value"] is True + assert events["cost"]["value"] == 2.5 + assert events["latency"]["value"] == 41.0 + assert {event["event_source"] for event in events.values()} == {"route_completion"} + assert completed["actual_economics"]["actual"]["cost_units"] == 2.5 + + second = _begin(ledger, route_id=str(uuid.uuid4())) + invalid = ledger.complete_decision( + second["route_id"], + success=False, + error_category="provider_error", + actual_economics={"cost_units": -1, "elapsed_ms": -20}, + ) + invalid_events = {event["outcome_name"]: event for event in invalid["outcome_events"]} + assert invalid_events["route_success"]["value"] is False + assert invalid_events["cost"]["observation_status"] == "not_observed" + assert invalid_events["cost"]["value"] is None + assert invalid_events["latency"]["observation_status"] == "not_observed" + + +def test_feedback_emits_exactly_one_quality_event_per_revision(tmp_path) -> None: + ledger = RoutePolicyLedger(tmp_path / "quality-events.sqlite3") + decision = _begin(ledger) + + first = ledger.record_feedback( + decision["route_id"], + {"rating": "up", "feedback_intent": "good"}, + idempotency_key="quality-1", + ) + retry = ledger.record_feedback( + decision["route_id"], + {"feedback_intent": "good", "rating": "up"}, + idempotency_key="quality-1", + ) + ledger.record_feedback( + decision["route_id"], + {"rating": "down", "feedback_intent": "too_slow"}, + idempotency_key="quality-2", + ) + ledger.record_feedback( + decision["route_id"], + {"feedback_axes": {"quality": -0.25}}, + idempotency_key="quality-3", + ) + ledger.record_feedback( + decision["route_id"], + {"feedback_axes": {"quality": 2.0}}, + idempotency_key="quality-4", + ) + + assert retry["revision"] == first["revision"] == 1 + assert retry["idempotent"] is True + events = ledger.get_decision(decision["route_id"])["outcome_events"] + assert len(events) == 4 + assert [event["event_key"] for event in events] == [ + "feedback_revision:1", + "feedback_revision:2", + "feedback_revision:3", + "feedback_revision:4", + ] + assert events[0]["value"] == 1.0 + assert events[0]["metadata"]["raw_rating"] == "up" + assert events[1]["observation_status"] == "not_observed" + assert events[1]["value"] is None + assert events[2]["value"] == -0.25 + assert events[3]["observation_status"] == "not_observed" + + +@pytest.mark.parametrize("creator", [_create_v1_database, _create_v2_database]) +def test_v1_v2_migration_backfills_posthoc_contracts_and_descriptive_events( + tmp_path, creator +) -> None: + database = tmp_path / f"{creator.__name__}.sqlite3" + route_id = creator(database) + + ledger = RoutePolicyLedger(database) + decision = ledger.get_decision(route_id) + + assert set(decision["outcome_contracts"]) == set(OUTCOME_NAMES) + assert decision["outcome_contracts_precommitted_at_begin"] is False + assert decision["outcome_contract_commitment_source"] == "legacy_posthoc" + assert all( + contract["precommitted"] is False + and contract["commitment_source"] == "legacy_posthoc" + and contract["contract_hash_valid"] is True + for contract in decision["outcome_contracts"].values() + ) + events = decision["outcome_events"] + assert len(events) == 4 + assert {event["event_source"] for event in events} == {"legacy_posthoc"} + assert {event["outcome_name"] for event in events} == set(OUTCOME_NAMES) + assert next(event for event in events if event["outcome_name"] == "cost")["value"] == 1.0 + assert ( + next(event for event in events if event["outcome_name"] == "latency")[ + "observation_status" + ] + == "not_observed" + ) + + # Reopening repeats neither contracts nor descriptive observations. + reopened = RoutePolicyLedger(database) + assert len(reopened.get_decision(route_id)["outcome_events"]) == 4 + connection = sqlite3.connect(database) + assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 + assert connection.execute("SELECT COUNT(*) FROM route_outcome_contracts").fetchone()[0] == 4 + assert connection.execute("SELECT COUNT(*) FROM route_outcome_observation_events").fetchone()[0] == 4 + connection.close() + + +def test_policy_evidence_snapshot_uses_one_finite_as_of_for_maturity(monkeypatch, tmp_path) -> None: + clock = {"now": 100.0} + monkeypatch.setattr(route_policy_ledger.time, "time", lambda: clock["now"]) + ledger = RoutePolicyLedger(tmp_path / "fixed-as-of.sqlite3") + contracts = {name: {} for name in OUTCOME_NAMES} + contracts["user_quality_rating"] = {"maturity_delay_seconds": 30} + decision = _begin(ledger, outcome_contracts=contracts) + clock["now"] = 110.0 + ledger.complete_decision( + decision["route_id"], + success=True, + actual_economics={"cost_units": 1.0, "elapsed_ms": 5.0}, + ) + clock["now"] = 120.0 + ledger.record_feedback(decision["route_id"], {"rating": "up"}) + + early = ledger.policy_evidence_snapshot(as_of=105.0) + assert early["as_of"] == 105.0 + assert early["analysis_window"]["as_of"] == 105.0 + assert early["lifecycle"]["counts"] == { + "started": 1, + "completed": 0, + "failed": 0, + "inflight": 1, + } + assert early["usage_rows"][0]["outcome_events"] == [] + maturity = early["outcome_contract_maturity"] + assert maturity["schema_version"] == OUTCOME_MATURITY_SCHEMA_VERSION + assert maturity["policy_value_estimate"] is None + assert maturity["causal_identification"] == "not_performed" + assert maturity["missingness_identification"] == "not_performed" + assert maturity["by_outcome"]["route_success"]["mature_contract_count"] == 1 + assert maturity["by_outcome"]["user_quality_rating"]["pending_contract_count"] == 1 + + completed = ledger.policy_evidence_snapshot(as_of=115.0) + assert completed["lifecycle"]["counts"]["completed"] == 1 + assert len(completed["usage_rows"][0]["outcome_events"]) == 3 + assert completed["lifecycle"]["feedback_coverage"]["known"] == 0 + + with_feedback = ledger.policy_evidence_snapshot(as_of=125.0) + assert len(with_feedback["usage_rows"][0]["outcome_events"]) == 4 + assert with_feedback["lifecycle"]["feedback_coverage"]["known"] == 1 + with pytest.raises(ValueError, match="finite timestamp"): + ledger.policy_evidence_snapshot(as_of=float("nan")) + + +def test_contract_hash_tampering_fails_closed_in_maturity(tmp_path) -> None: + database = tmp_path / "contract-tamper.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + connection.execute("DROP TRIGGER route_outcome_contracts_no_update") + connection.execute( + "UPDATE route_outcome_contracts SET contract_json = '{}' WHERE route_id = ? AND outcome_name = 'cost'", + (decision["route_id"],), + ) + connection.commit() + connection.close() + + tampered = ledger.get_decision(decision["route_id"]) + assert tampered["outcome_contracts"]["cost"]["contract_hash_valid"] is False + assert ( + tampered["outcome_contracts"]["cost"]["contract_hash_reason"] + == "contract_projection_mismatch" + ) + assert tampered["outcome_contracts_precommitted_at_begin"] is False + maturity = ledger.policy_evidence_snapshot( + as_of=decision["started_at"] + 2.0 + )["outcome_contract_maturity"] + assert maturity["precommitted_routes"] == 0 + assert maturity["complete_contract_sets"] == 0 + assert maturity["by_outcome"]["cost"]["invalid_contract_count"] == 1 + assert maturity["by_outcome"]["cost"]["precommitted_count"] == 0 + assert maturity["by_outcome"]["cost"]["mature_contract_count"] == 0 + + +def test_contract_projection_type_tampering_fails_closed(tmp_path) -> None: + database = tmp_path / "contract-type-tamper.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + connection.execute("DROP TRIGGER route_outcome_contracts_no_update") + raw = json.loads( + connection.execute( + "SELECT contract_json FROM route_outcome_contracts " + "WHERE route_id = ? AND outcome_name = 'cost'", + (decision["route_id"],), + ).fetchone()[0] + ) + raw["precommitted"] = 1 + connection.execute( + "UPDATE route_outcome_contracts SET contract_json = ? " + "WHERE route_id = ? AND outcome_name = 'cost'", + (json.dumps(raw, sort_keys=True, separators=(",", ":")), decision["route_id"]), + ) + connection.commit() + connection.close() + + tampered = ledger.get_decision(decision["route_id"]) + assert tampered["outcome_contracts"]["cost"]["contract_hash_valid"] is False + assert ( + tampered["outcome_contracts"]["cost"]["contract_hash_reason"] + == "contract_projection_mismatch" + ) + assert tampered["outcome_contracts_precommitted_at_begin"] is False + + +def test_late_outcome_contract_commitment_fails_closed(tmp_path) -> None: + database = tmp_path / "contract-timing-tamper.sqlite3" + ledger = RoutePolicyLedger(database) + decision = _begin(ledger) + + connection = sqlite3.connect(database) + connection.execute("DROP TRIGGER route_outcome_contracts_no_update") + connection.execute( + "UPDATE route_outcome_contracts SET committed_at = ? " + "WHERE route_id = ? AND outcome_name = 'cost'", + (decision["started_at"] + 1.0, decision["route_id"]), + ) + connection.commit() + connection.close() + + tampered = ledger.get_decision(decision["route_id"]) + cost_contract = tampered["outcome_contracts"]["cost"] + assert cost_contract["contract_hash_valid"] is True + assert cost_contract["commitment_timing_valid"] is False + assert cost_contract["commitment_timing_reason"] == "committed_after_decision_start" + assert tampered["outcome_contracts_precommitted_at_begin"] is False + maturity = ledger.policy_evidence_snapshot( + as_of=decision["started_at"] + 2.0 + )["outcome_contract_maturity"] + assert maturity["precommitted_routes"] == 0 + assert maturity["by_outcome"]["cost"]["late_commitment_count"] == 1 + assert maturity["by_outcome"]["cost"]["precommitted_count"] == 0 + assert maturity["by_outcome"]["cost"]["mature_contract_count"] == 0 diff --git a/test_route_policy_protocol.py b/test_route_policy_protocol.py new file mode 100644 index 000000000..669a0727d --- /dev/null +++ b/test_route_policy_protocol.py @@ -0,0 +1,435 @@ +import builtins +import copy +import json + +import pytest + +from source.route_policy_explorer import ACTIVATION_BLOCKERS, plan_adjacent_route_study +from source.route_policy_protocol import ( + PROTOCOL_BUILD_INPUT_SCHEMA_VERSION, + PROTOCOL_SCHEMA_VERSION, + REVIEW_BUNDLE_SCHEMA_VERSION, + _PROTOCOL_HASH_DOMAIN, + _REVIEW_BUNDLE_HASH_DOMAIN, + _domain_hash, + audit_route_study_review_bundle, + audit_route_study_protocol, + build_route_study_review_bundle, + build_route_study_protocol, +) +from source.route_policy_study_cli import _example_payload + + +def _study(*, candidate_hash="a", distribution_hash="b"): + payload = _example_payload() + payload["source_contract"] = { + **payload["source_contract"], + "candidate_set_hash": candidate_hash * 64, + "distribution_hash": distribution_hash * 64, + } + return plan_adjacent_route_study( + payload["baseline_mode"], + payload["post_filter_candidates"], + payload["post_filter_exclusions"], + source_contract=payload["source_contract"], + exploration_rate=payload["exploration_rate"], + planned_routes=payload["planned_routes"], + scenario_confidence=payload["scenario_confidence"], + assumed_feedback_rate=payload["assumed_feedback_rate"], + target_observed_labels=payload["target_observed_labels"], + ) + + +def _rehash_protocol(protocol): + protocol["protocol_hash"] = _domain_hash( + _PROTOCOL_HASH_DOMAIN, + { + "schema_version": protocol["schema_version"], + "protocol": protocol["protocol"], + "charter": protocol["charter"], + }, + ) + + +def _rehash_bundle(bundle): + bundle["bundle_hash"] = _domain_hash( + _REVIEW_BUNDLE_HASH_DOMAIN, + { + "schema_version": bundle["schema_version"], + "bundle": bundle["bundle"], + "protocol_builder": bundle["protocol_builder"], + "source_study_plans": bundle["source_study_plans"], + "protocol": bundle["protocol"], + }, + ) + + +def test_default_protocol_is_prompt_free_canonical_and_fail_closed(): + protocol = build_route_study_protocol(_study()) + + assert protocol["schema_version"] == PROTOCOL_SCHEMA_VERSION + assert protocol["protocol"] == { + "version": "1.0.0", + "label": "Stateful Route Experiment Preflight v1", + "state": "draft_for_independent_review", + "activation_available": False, + } + assert len(protocol["protocol_hash"]) == 64 + int(protocol["protocol_hash"], 16) + + charter = protocol["charter"] + assert charter["source_studies"]["support_stratum_count"] == 1 + assert charter["source_studies"]["source_plans_embedded"] is False + stratum = charter["source_studies"]["support_strata"][0] + assert stratum["executed_logging_propensities"] is False + assert stratum["eligible_actions"] == ["off", "collective", "loop"] + assert charter["target_policy_class"]["profile_name"] == "balanced" + assert charter["target_policy_class"]["thresholds"] == { + "off": 0, + "collective": 2, + "loop": 4, + "collective_loop": 5, + } + assert charter["target_policy_class"]["optimality_claim"] is False + assert charter["population"]["one_study_policy_per_cluster"] is True + assert charter["population"]["population_precommitted"] is False + assert charter["stopping_and_resources"]["outcome_dependent_stopping_allowed"] is False + assert charter["stopping_and_resources"]["automatic_promotion_allowed"] is False + assert charter["randomness"]["assignment_implementation_available"] is False + assert charter["randomness"]["assignment_performed"] is False + assert charter["randomness"]["seed_commitment"] is None + assert charter["external_evaluation"]["causal_estimate_available"] is False + assert charter["outcome_observation"]["observation_process_validated"] is False + assert set(charter["outcome_observation"]["contracts"]) == { + "route_success", + "user_quality_rating", + "cost", + "latency", + } + assert all( + contract["precommitted"] is True + and contract["commitment_source"] == "protocol_draft" + for contract in charter["outcome_observation"]["contracts"].values() + ) + assert [row["code"] for row in charter["blocker_register"]] == list( + ACTIVATION_BLOCKERS + ) + assert all(row["activation_blocking"] is True for row in charter["blocker_register"]) + boundaries = charter["causal_boundaries"] + assert boundaries["activation_blockers"] == list(ACTIVATION_BLOCKERS) + assert boundaries["declarations_are_not_validation"] is True + assert boundaries["execution_enabled"] is False + assert boundaries["ledger_write_performed"] is False + assert boundaries["model_inference_performed"] is False + assert boundaries["off_policy_estimate_computed"] is False + assert boundaries["automatic_promotion_allowed"] is False + assert audit_route_study_protocol(protocol)["ok"] is True + json.dumps(protocol, sort_keys=True, allow_nan=False) + + +def test_default_design_surfaces_unknown_state_and_screens_out_route_assignment(): + screen = build_route_study_protocol(_study())["charter"]["stateful_design"] + + assert screen["selected_design_mode"] == "sticky_session_cluster" + assert screen["assignment_unit"] == "session_hash" + assert screen["selected_design_status"] == "declaration_incomplete" + assert screen["carryover_validated"] is False + assert screen["interference_validated"] is False + route = screen["candidate_screen"][0] + assert route == { + "design_mode": "route_randomization", + "status": "screened_out", + "blocking_reasons": [ + "stateful_product_campaigns_cannot_randomize_independent_routes_in_v1" + ], + "assignment_implementation_available": False, + } + assert set(screen["selected_design_blocking_reasons"]) == { + "carryover_scope_unknown", + "interference_scope_unknown", + "temporal_variation_unknown", + } + + +def test_switchback_declarations_remain_unvalidated_and_require_washout(): + no_washout = build_route_study_protocol( + _study(), + design_mode="clustered_switchback", + carryover_scope="within_session", + interference_scope="shared_resource", + temporal_variation="nonstationary", + block_length_routes=20, + washout_routes=0, + ) + screen = no_washout["charter"]["stateful_design"] + assert screen["assignment_unit"] == "session_hash_x_time_block" + assert screen["selected_design_status"] == "incompatible_or_unvalidated_assumptions" + assert set(screen["selected_design_blocking_reasons"]) == { + "positive_washout_required_for_declared_carryover", + "interference_graph_and_exposure_mapping_not_validated", + "block_balance_and_mixing_not_validated", + } + + with_washout = build_route_study_protocol( + _study(), + design_mode="clustered_switchback", + carryover_scope="within_session", + interference_scope="none_declared", + temporal_variation="stable_declared", + block_length_routes=20, + washout_routes=4, + ) + selected = with_washout["charter"]["stateful_design"] + assert selected["selected_design_status"] == "assumptions_declared_unvalidated" + assert selected["selected_design_blocking_reasons"] == [] + assert selected["causal_design_certified"] is False + assert with_washout["protocol"]["activation_available"] is False + + +def test_profile_population_resources_and_seed_commitment_are_bound_but_unsealed(): + base = build_route_study_protocol(_study()) + changed = build_route_study_protocol( + _study(), + target_policy_profile="quality_first", + population_rule_id="local-consenting-session-clusters", + population_rule_version="2", + cluster_key_schema_version="study-scoped-session-hash-v2", + planned_clusters=480, + max_routes_per_cluster=12, + analysis_every_clusters=60, + seed_commitment="c" * 64, + external_estimator_id="independent-cluster-aipw-v1", + external_reviewer_id="review-board-v1", + ) + + assert changed["protocol_hash"] != base["protocol_hash"] + assert changed["charter"]["target_policy_class"]["profile_name"] == "quality_first" + assert changed["charter"]["population"]["planned_clusters"] == 480 + assert changed["charter"]["stopping_and_resources"]["planned_route_ceiling"] == 5760 + randomness = changed["charter"]["randomness"] + assert randomness["seed_commitment"] == "c" * 64 + assert randomness["seed_material_included"] is False + assert randomness["seed_commitment_sealed"] is False + assert randomness["nonce_grinding_resistant"] is False + seed_blocker = next( + row + for row in changed["charter"]["blocker_register"] + if row["code"] == "preassignment_seed_commitment_not_sealed" + ) + assert seed_blocker["status"] == "drafted_unsealed" + assert changed["charter"]["external_evaluation"]["validated"] is False + + +def test_multiple_strata_are_deduplicated_order_invariant_and_share_source_cohort(): + first = _study(candidate_hash="a", distribution_hash="b") + second = _study(candidate_hash="c", distribution_hash="d") + + forward = build_route_study_protocol([first, second]) + reverse = build_route_study_protocol([second, first]) + assert forward == reverse + assert forward["charter"]["source_studies"]["support_stratum_count"] == 2 + hashes = [ + row["study_design_hash"] + for row in forward["charter"]["source_studies"]["support_strata"] + ] + assert hashes == sorted(hashes) + + with pytest.raises(ValueError, match="duplicate design hashes"): + build_route_study_protocol([first, first]) + + incompatible = copy.deepcopy(second) + incompatible["charter"]["source_contract"]["policy_version"] = "3.0.0" + # Rebuilding gives it a valid explorer hash but a different source cohort. + payload = _example_payload() + payload["source_contract"]["policy_version"] = "3.0.0" + incompatible = plan_adjacent_route_study( + payload["baseline_mode"], + payload["post_filter_candidates"], + payload["post_filter_exclusions"], + source_contract=payload["source_contract"], + ) + with pytest.raises(ValueError, match="one source policy/schema cohort"): + build_route_study_protocol([first, incompatible]) + + +def test_audit_rejects_tampering_activation_and_schema_drift(): + protocol = build_route_study_protocol(_study()) + + tampered = copy.deepcopy(protocol) + tampered["charter"]["target_policy_class"]["thresholds"]["loop"] = 999 + with pytest.raises(ValueError, match="frozen profile"): + audit_route_study_protocol(tampered) + + activated = copy.deepcopy(protocol) + activated["charter"]["causal_boundaries"]["execution_enabled"] = True + with pytest.raises(ValueError, match="fail-closed"): + audit_route_study_protocol(activated) + + extra = copy.deepcopy(protocol) + extra["activate"] = True + with pytest.raises(ValueError, match="top-level"): + audit_route_study_protocol(extra) + + rehashed_unsafe = copy.deepcopy(protocol) + rehashed_unsafe["charter"]["randomness"]["assignment_performed"] = True + rehashed_unsafe["protocol_hash"] = _domain_hash( + _PROTOCOL_HASH_DOMAIN, + { + "schema_version": rehashed_unsafe["schema_version"], + "protocol": rehashed_unsafe["protocol"], + "charter": rehashed_unsafe["charter"], + }, + ) + with pytest.raises(ValueError, match="non-assigning and unsealed"): + audit_route_study_protocol(rehashed_unsafe) + + +def test_builder_rejects_unsafe_or_ambiguous_protocol_inputs(): + plan = _study() + with pytest.raises(ValueError, match="design_mode must be one of"): + build_route_study_protocol(plan, design_mode="route_randomization") + with pytest.raises(ValueError, match="carryover_scope must be one of"): + build_route_study_protocol(plan, carryover_scope="probably fine") + with pytest.raises(ValueError, match="prompt-free identifier"): + build_route_study_protocol(plan, population_rule_id="people who ask about health") + with pytest.raises(ValueError, match="lowercase SHA-256"): + build_route_study_protocol(plan, seed_commitment="ABC") + with pytest.raises(ValueError, match="smaller than block_length"): + build_route_study_protocol(plan, block_length_routes=10, washout_routes=10) + + +def test_builder_does_not_use_io_randomness_or_mutate_source_plan(monkeypatch): + plan = _study() + before = copy.deepcopy(plan) + + def fail(*_args, **_kwargs): + raise AssertionError("protocol preflight must not perform file I/O") + + monkeypatch.setattr(builtins, "open", fail) + protocol = build_route_study_protocol(plan) + assert plan == before + assert protocol["charter"]["causal_boundaries"]["io_performed"] is False + assert protocol["charter"]["randomness"]["assignment_performed"] is False + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + ( + lambda draft: draft["charter"]["stopping_and_resources"].update( + outcome_dependent_stopping_allowed=True, + automatic_promotion_allowed=True, + ), + "stopping rules", + ), + ( + lambda draft: draft["charter"]["population"].update( + population_precommitted=True, + cluster_map_validated=True, + ), + "population", + ), + ( + lambda draft: draft["charter"]["outcome_observation"]["contracts"][ + "user_quality_rating" + ].update(unit="invented_quality_proxy"), + "outcome contracts", + ), + ( + lambda draft: draft["charter"]["external_evaluation"].update( + validated=True, + winner_available=True, + ), + "external evaluation", + ), + ], +) +def test_structural_audit_rejects_semantic_mutations_even_after_public_rehash(mutate, message): + draft = build_route_study_protocol(_study()) + mutate(draft) + _rehash_protocol(draft) + + with pytest.raises(ValueError, match=message): + audit_route_study_protocol(draft) + + +def test_review_bundle_is_source_bound_order_invariant_and_explicitly_unsigned(): + first = _study(candidate_hash="a", distribution_hash="b") + second = _study(candidate_hash="c", distribution_hash="d") + + forward = build_route_study_review_bundle( + [first, second], + target_policy_profile="quality_first", + carryover_scope="within_session", + interference_scope="none_declared", + temporal_variation="stable_declared", + ) + reverse = build_route_study_review_bundle( + [second, first], + target_policy_profile="quality_first", + carryover_scope="within_session", + interference_scope="none_declared", + temporal_variation="stable_declared", + ) + + assert forward == reverse + assert forward["schema_version"] == REVIEW_BUNDLE_SCHEMA_VERSION + assert forward["protocol_builder"]["schema_version"] == ( + PROTOCOL_BUILD_INPUT_SCHEMA_VERSION + ) + assert forward["bundle"]["verification_level"] == ( + "full_source_bound_reconstruction" + ) + assert forward["bundle"]["authenticity_proof_available"] is False + assert forward["bundle"]["trusted_timestamp_available"] is False + audit = audit_route_study_review_bundle(forward) + assert audit["ok"] is True + assert audit["support_stratum_count"] == 2 + assert audit["source_plan_reconstruction_performed"] is True + assert audit["activation_available"] is False + + +def test_review_bundle_reconstruction_rejects_rehashed_source_binding_substitution(): + bundle = build_route_study_review_bundle( + [_study(candidate_hash="a", distribution_hash="b"), _study(candidate_hash="c", distribution_hash="d")] + ) + protocol = bundle["protocol"] + old_hash = protocol["charter"]["source_studies"]["support_strata"][0][ + "study_design_hash" + ] + replacement = "f" * 64 if old_hash != "f" * 64 else "e" * 64 + protocol["charter"]["source_studies"]["support_strata"][0][ + "study_design_hash" + ] = replacement + protocol["charter"]["population"]["admitted_support_strata"][0] = replacement + # Keep structural ordering canonical so only full source reconstruction exposes it. + paired = list( + zip( + protocol["charter"]["source_studies"]["support_strata"], + protocol["charter"]["population"]["admitted_support_strata"], + ) + ) + paired.sort(key=lambda pair: pair[0]["study_design_hash"]) + protocol["charter"]["source_studies"]["support_strata"] = [pair[0] for pair in paired] + protocol["charter"]["population"]["admitted_support_strata"] = [ + pair[1] for pair in paired + ] + _rehash_protocol(protocol) + assert audit_route_study_protocol(protocol)["verification_level"] == ( + "structural_without_source_plans" + ) + _rehash_bundle(bundle) + + with pytest.raises(ValueError, match="does not reconstruct"): + audit_route_study_review_bundle(bundle) + + +def test_review_bundle_rejects_reordered_or_noncanonical_source_plans_after_rehash(): + bundle = build_route_study_review_bundle( + [_study(candidate_hash="a", distribution_hash="b"), _study(candidate_hash="c", distribution_hash="d")] + ) + bundle["source_study_plans"].reverse() + _rehash_bundle(bundle) + + with pytest.raises(ValueError, match="not canonical or canonically ordered"): + audit_route_study_review_bundle(bundle) diff --git a/test_route_policy_protocol_cli.py b/test_route_policy_protocol_cli.py new file mode 100644 index 000000000..76a86d40d --- /dev/null +++ b/test_route_policy_protocol_cli.py @@ -0,0 +1,105 @@ +import json + +from source.route_policy_protocol_cli import main + + +def test_cli_example_exports_fail_closed_protocol(capsys): + assert main(["--example", "--compact"]) == 0 + captured = capsys.readouterr() + assert captured.err == "" + protocol = json.loads(captured.out) + assert protocol["protocol"]["label"] == "Stateful Route Experiment Preflight v1" + assert protocol["protocol"]["activation_available"] is False + assert protocol["charter"]["stateful_design"]["selected_design_status"] == ( + "declaration_incomplete" + ) + assert protocol["charter"]["randomness"]["assignment_performed"] is False + exact_contract = protocol["charter"]["source_studies"]["common_source_contract"] + assert exact_contract["outcome_contract_schema_version"] == "route-outcome-contract-v1" + + +def test_cli_build_and_audit_round_trip(tmp_path, capsys): + protocol_path = tmp_path / "protocol.json" + assert main(["--example", "--output", str(protocol_path)]) == 0 + assert capsys.readouterr().out == "" + + assert main(["--audit", str(protocol_path), "--compact"]) == 0 + captured = capsys.readouterr() + audit = json.loads(captured.out) + assert captured.err == "" + assert audit["ok"] is True + assert audit["state"] == "draft_for_independent_review" + assert audit["activation_available"] is False + assert len(audit["activation_blockers"]) == 8 + + +def test_cli_rejects_prompt_bearing_input(tmp_path, capsys): + input_path = tmp_path / "unsafe.json" + input_path.write_text(json.dumps({"study_plans": [], "prompt": "private text"}), encoding="utf-8") + assert main(["--input", str(input_path)]) == 2 + captured = capsys.readouterr() + assert captured.out == "" + assert "unsupported or non-prompt-free fields: prompt" in captured.err + + +def test_cli_rejects_tampered_protocol_audit(tmp_path, capsys): + protocol_path = tmp_path / "protocol.json" + assert main(["--example", "--output", str(protocol_path)]) == 0 + capsys.readouterr() + protocol = json.loads(protocol_path.read_text(encoding="utf-8")) + protocol["charter"]["causal_boundaries"]["automatic_promotion_allowed"] = True + protocol_path.write_text(json.dumps(protocol), encoding="utf-8") + + assert main(["--audit", str(protocol_path)]) == 2 + captured = capsys.readouterr() + assert captured.out == "" + assert "fail-closed" in captured.err + + +def test_cli_accepts_windows_utf8_bom_protocol_files(tmp_path, capsys): + protocol_path = tmp_path / "protocol.json" + assert main(["--example", "--output", str(protocol_path)]) == 0 + capsys.readouterr() + raw = protocol_path.read_text(encoding="utf-8") + protocol_path.write_text(raw, encoding="utf-8-sig") + + assert main(["--audit", str(protocol_path), "--compact"]) == 0 + captured = capsys.readouterr() + assert json.loads(captured.out)["ok"] is True + assert captured.err == "" + + +def test_cli_example_bundle_round_trips_with_full_source_reconstruction(tmp_path, capsys): + bundle_path = tmp_path / "review-bundle.json" + assert main(["--example-bundle", "--output", str(bundle_path)]) == 0 + assert capsys.readouterr().out == "" + bundle = json.loads(bundle_path.read_text(encoding="utf-8")) + assert bundle["schema_version"] == "route-study-review-bundle-v1" + assert len(bundle["source_study_plans"]) == 2 + assert bundle["bundle"]["authenticity_proof_available"] is False + + bundle_path.write_text( + bundle_path.read_text(encoding="utf-8"), encoding="utf-8-sig" + ) + assert main(["--audit-bundle", str(bundle_path), "--compact"]) == 0 + captured = capsys.readouterr() + audit = json.loads(captured.out) + assert captured.err == "" + assert audit["ok"] is True + assert audit["support_stratum_count"] == 2 + assert audit["verification_level"] == "full_source_bound_reconstruction" + assert audit["source_plan_reconstruction_performed"] is True + + +def test_cli_bundle_audit_rejects_source_plan_substitution(tmp_path, capsys): + bundle_path = tmp_path / "review-bundle.json" + assert main(["--example-bundle", "--output", str(bundle_path)]) == 0 + capsys.readouterr() + bundle = json.loads(bundle_path.read_text(encoding="utf-8")) + bundle["source_study_plans"][0]["design_hash"] = "f" * 64 + bundle_path.write_text(json.dumps(bundle), encoding="utf-8") + + assert main(["--audit-bundle", str(bundle_path)]) == 2 + captured = capsys.readouterr() + assert captured.out == "" + assert "design_hash does not match" in captured.err diff --git a/test_route_policy_shadow_cli.py b/test_route_policy_shadow_cli.py new file mode 100644 index 000000000..47917d8b7 --- /dev/null +++ b/test_route_policy_shadow_cli.py @@ -0,0 +1,310 @@ +import json +from pathlib import Path + +from source.route_policy_protocol import build_route_study_review_bundle_from_input +from source.route_policy_protocol_cli import _example_bundle_input +from source.route_policy_shadow_cli import main +from source.route_policy_shadow_registry import ( + create_shadow_campaign_artifacts, + generate_shadow_seed, +) + + +def _write_json(path: Path, value): + path.write_text(json.dumps(value), encoding="utf-8") + + +def _bundle_file(tmp_path: Path) -> Path: + path = tmp_path / "review-bundle.json" + _write_json(path, build_route_study_review_bundle_from_input(_example_bundle_input())) + return path + + +def _run_success(capsys, argv): + assert main(argv) == 0 + captured = capsys.readouterr() + assert captured.err == "" + return json.loads(captured.out), captured.out + + +def test_cli_runs_sealed_commit_close_reveal_verify_flow_without_secret_or_raw_id_leaks( + tmp_path, capsys +): + bundle_path = _bundle_file(tmp_path) + registry_path = tmp_path / "shadow.sqlite3" + capsule_path = tmp_path / "private-seed.json" + + sealed, seal_stdout = _run_success( + capsys, + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(bundle_path), + "--seed-output", + str(capsule_path), + "--compact", + ], + ) + capsule = json.loads(capsule_path.read_text(encoding="utf-8")) + private_seed = capsule["seed_material_base64url"] + campaign_id = sealed["public_package"]["campaign_seal"]["seal"]["campaign_id"] + assert sealed["created"] is True + assert sealed["private_seed_capsule_written"] is True + assert sealed["private_seed_material_returned"] is False + assert '"private_seed_capsule":{' not in seal_stdout + assert private_seed not in seal_stdout + + cluster_identifier = "tenant-raw-id-that-must-not-be-exported" + cluster_path = tmp_path / "cluster.json" + _write_json(cluster_path, {"cluster_identifier": cluster_identifier}) + committed, commit_stdout = _run_success( + capsys, + [ + "commit", + "--registry", + str(registry_path), + "--campaign", + campaign_id, + "--seed-input", + str(capsule_path), + "--cluster-input", + str(cluster_path), + "--compact", + ], + ) + assert committed["created"] is True + assert committed["chosen_arm_revealed"] is False + assert cluster_identifier not in commit_stdout + assert private_seed not in commit_stdout + assert "arm_id" not in commit_stdout + + # The separate registry, WAL, and shared-memory files must not contain the + # private seed before the explicit post-closure reveal. + for path in tmp_path.glob("shadow.sqlite3*"): + assert private_seed.encode("ascii") not in path.read_bytes() + + closed, _ = _run_success( + capsys, + ["close", "--registry", str(registry_path), "--campaign", campaign_id, "--compact"], + ) + assert closed["closure"]["frozen_commitment_count"] == 1 + + revealed, reveal_stdout = _run_success( + capsys, + [ + "reveal", + "--registry", + str(registry_path), + "--campaign", + campaign_id, + "--seed-input", + str(capsule_path), + "--compact", + ], + ) + assert revealed["seed_material_revealed_in_registry"] is True + assert revealed["seed_material_returned"] is False + assert private_seed not in reveal_stdout + + verified, _ = _run_success( + capsys, + ["verify", "--registry", str(registry_path), "--campaign", campaign_id, "--compact"], + ) + assert verified == { + "activation_available": False, + "campaign_id": campaign_id, + "campaign_audit_ok": True, + "campaign_audit_performed": True, + "complete": True, + "execution_enabled": False, + "invalid_commitment_artifacts": 0, + "matched": 1, + "mismatched": 0, + "total_mismatched": 0, + "ok": True, + "processed": 1, + "processing_complete": True, + "remaining": 0, + "verification_complete": True, + } + + status, _ = _run_success( + capsys, + ["status", "--registry", str(registry_path), "--campaign", campaign_id, "--compact"], + ) + assert status["campaign_count"] == 1 + assert status["campaigns"][0]["state"] == "reveal_verification_complete" + assert status["campaigns"][0]["matched_assignment_count"] == 1 + assert status["read_only"] is True + assert status["execution_enabled"] is False + assert status["automatic_promotion_allowed"] is False + + # Status is a verifier, not a mutation command. Durable database bytes and + # logical evidence must therefore remain unchanged across repeated reads. + before_database = registry_path.read_bytes() + before_rows = status + status_again, _ = _run_success( + capsys, + ["status", "--registry", str(registry_path), "--campaign", campaign_id, "--compact"], + ) + assert status_again == before_rows + assert registry_path.read_bytes() == before_database + + +def test_seal_recovers_exclusive_capsule_and_is_idempotent(tmp_path, capsys): + bundle_path = _bundle_file(tmp_path) + bundle = json.loads(bundle_path.read_text(encoding="utf-8")) + capsule_path = tmp_path / "private-seed.json" + registry_path = tmp_path / "shadow.sqlite3" + artifacts = create_shadow_campaign_artifacts(bundle, generate_shadow_seed()) + _write_json(capsule_path, artifacts["private_seed_capsule"]) + + first, first_stdout = _run_success( + capsys, + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(bundle_path), + "--seed-output", + str(capsule_path), + "--compact", + ], + ) + assert first["created"] is True + assert first["private_seed_capsule_written"] is False + assert first["private_seed_capsule_recovered"] is True + assert artifacts["private_seed_capsule"]["seed_material_base64url"] not in first_stdout + + second, _ = _run_success( + capsys, + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(bundle_path), + "--seed-output", + str(capsule_path), + "--compact", + ], + ) + assert second["created"] is False + assert second["private_seed_capsule_recovered"] is True + + capsule_path.unlink() + assert ( + main( + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(bundle_path), + "--seed-output", + str(capsule_path), + ] + ) + == 2 + ) + captured = capsys.readouterr() + assert captured.out == "" + assert "recover its original seed capsule" in captured.err + assert not capsule_path.exists() + + +def test_cli_strict_json_rejects_duplicate_keys_and_non_finite_values(tmp_path, capsys): + registry_path = tmp_path / "shadow.sqlite3" + capsule_path = tmp_path / "private-seed.json" + duplicate_bundle = tmp_path / "duplicate.json" + duplicate_bundle.write_text('{"schema_version":"one","schema_version":"two"}', encoding="utf-8") + + assert ( + main( + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(duplicate_bundle), + "--seed-output", + str(capsule_path), + ] + ) + == 2 + ) + captured = capsys.readouterr() + assert captured.out == "" + assert "duplicate object key: schema_version" in captured.err + assert not registry_path.exists() + assert not capsule_path.exists() + + non_finite = tmp_path / "non-finite.json" + non_finite.write_text('{"value":NaN}', encoding="utf-8") + assert ( + main( + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(non_finite), + "--seed-output", + str(capsule_path), + ] + ) + == 2 + ) + captured = capsys.readouterr() + assert captured.out == "" + assert "non-finite number: NaN" in captured.err + + +def test_commit_rejects_ambiguous_cluster_input_without_echoing_identifier(tmp_path, capsys): + bundle_path = _bundle_file(tmp_path) + registry_path = tmp_path / "shadow.sqlite3" + capsule_path = tmp_path / "private-seed.json" + sealed, _ = _run_success( + capsys, + [ + "seal", + "--registry", + str(registry_path), + "--bundle", + str(bundle_path), + "--seed-output", + str(capsule_path), + "--compact", + ], + ) + campaign_id = sealed["public_package"]["campaign_seal"]["seal"]["campaign_id"] + cluster_path = tmp_path / "cluster.json" + cluster_path.write_text( + '{"cluster_identifier":"secret-a","cluster_identifier":"secret-b"}', + encoding="utf-8", + ) + + assert ( + main( + [ + "commit", + "--registry", + str(registry_path), + "--campaign", + campaign_id, + "--seed-input", + str(capsule_path), + "--cluster-input", + str(cluster_path), + ] + ) + == 2 + ) + captured = capsys.readouterr() + assert captured.out == "" + assert "duplicate object key: cluster_identifier" in captured.err + assert "secret-a" not in captured.err + assert "secret-b" not in captured.err diff --git a/test_route_policy_shadow_registry.py b/test_route_policy_shadow_registry.py new file mode 100644 index 000000000..bd81b1f75 --- /dev/null +++ b/test_route_policy_shadow_registry.py @@ -0,0 +1,853 @@ +import base64 +import copy +import json +import sqlite3 +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier + +import pytest + +from source.route_policy_explorer import plan_adjacent_route_study +from source.route_policy_protocol import build_route_study_review_bundle +from source.route_policy_shadow_registry import ( + SHADOW_ASSIGNMENT_ALGORITHM, + RouteShadowAssignmentRegistry, + RouteShadowRegistryError, + ShadowRegistryConflictError, + _PACKAGE_HASH_DOMAIN, + _canonical_json, + _domain_hash, + audit_shadow_campaign_artifacts, + audit_shadow_seed_capsule, + prepare_shadow_assignment_commitment, +) +from source.route_policy_study_cli import _example_payload + + +SEED = bytes.fromhex( + "04b12fc8e9714e36aa4c0f39bfc81460e959f6f99b7264e614c76d612b835f58" +) +OTHER_SEED = bytes.fromhex( + "e8134bf1a09987438cc8f0d7dfe8c6d507d635f8a10b049e4f33d744a0e4d90a" +) + + +def _review_bundle(*, candidate="a", distribution="b", planned_clusters=8): + payload = _example_payload() + payload["source_contract"] = { + **payload["source_contract"], + "candidate_set_hash": candidate * 64, + "distribution_hash": distribution * 64, + } + study = plan_adjacent_route_study( + payload["baseline_mode"], + payload["post_filter_candidates"], + payload["post_filter_exclusions"], + source_contract=payload["source_contract"], + exploration_rate=payload["exploration_rate"], + planned_routes=payload["planned_routes"], + scenario_confidence=payload["scenario_confidence"], + assumed_feedback_rate=payload["assumed_feedback_rate"], + target_observed_labels=payload["target_observed_labels"], + ) + return build_route_study_review_bundle( + study, + planned_clusters=planned_clusters, + analysis_every_clusters=max(1, planned_clusters // 2), + ) + + +def _sealed(tmp_path: Path, *, seed=SEED, bundle=None): + db_path = tmp_path / "route-policy-shadow-registry.sqlite3" + registry = RouteShadowAssignmentRegistry(db_path) + result = registry.seal_campaign(bundle or _review_bundle(), seed) + package = result["public_package"] + return { + "db_path": db_path, + "registry": registry, + "seal_result": result, + "package": package, + "capsule": result["private_seed_capsule"], + "campaign_id": package["campaign_seal"]["seal"]["campaign_id"], + } + + +def _assert_no_floats(value): + if isinstance(value, dict): + for item in value.values(): + _assert_no_floats(item) + elif isinstance(value, list): + for item in value: + _assert_no_floats(item) + else: + assert not isinstance(value, float) + + +def _rehash_public_package(package): + payload = {key: value for key, value in package.items() if key != "public_package_hash"} + package["public_package_hash"] = _domain_hash( + _PACKAGE_HASH_DOMAIN, + payload, + "shadow public package", + ) + + +def _database_files(db_path: Path): + return [path for path in db_path.parent.glob(f"{db_path.name}*") if path.is_file()] + + +def test_valid_lifecycle_is_batch_reconstructable_and_stays_shadow_only(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + package = sealed["package"] + capsule = sealed["capsule"] + campaign_id = sealed["campaign_id"] + + assert sealed["seal_result"]["created"] is True + assert sealed["seal_result"]["private_seed_material_persisted"] is False + assert audit_shadow_campaign_artifacts(package)["verification_level"] == ( + "full_source_bound_shadow_seal_reconstruction" + ) + assert audit_shadow_seed_capsule(package, capsule)["seed_material_returned"] is False + + first = registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=capsule, + cluster_identifier="raw-session-alpha", + ) + second = registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=capsule, + cluster_identifier="raw-session-beta", + ) + for result in (first, second): + assert result["created"] is True + assert result["chosen_arm_revealed"] is False + assert result["private_reveal_persisted"] is False + assert "arm_id" not in json.dumps(result["commitment"], sort_keys=True) + + before = registry.snapshot(campaign_id) + assert before["campaigns"][0]["state"] == "accepting_commitments" + assert before["campaigns"][0]["commitment_count"] == 2 + assert before["campaigns"][0]["whole_policy_arm_counts"] == {} + + closure = registry.close_campaign(campaign_id) + assert closure["closure"]["frozen_commitment_count"] == 2 + seed_reveal = registry.reveal_seed(campaign_id=campaign_id, seed_capsule=capsule) + assert seed_reveal["reveal"]["seed_material_base64url"] == capsule[ + "seed_material_base64url" + ] + + first_batch = registry.verify_assignment_reveals(campaign_id, batch_size=1) + second_batch = registry.verify_assignment_reveals(campaign_id, batch_size=1) + idempotent_batch = registry.verify_assignment_reveals(campaign_id, batch_size=1) + assert first_batch == { + **first_batch, + "ok": True, + "processed": 1, + "matched": 1, + "mismatched": 0, + "remaining": 1, + "complete": False, + "execution_enabled": False, + "activation_available": False, + } + assert second_batch["processed"] == second_batch["matched"] == 1 + assert second_batch["remaining"] == 0 + assert second_batch["complete"] is True + assert idempotent_batch["processed"] == 0 + assert idempotent_batch["complete"] is True + + after = registry.snapshot(campaign_id) + campaign = after["campaigns"][0] + assert campaign["state"] == "reveal_verification_complete" + assert campaign["matched_assignment_count"] == 2 + assert campaign["mismatched_assignment_count"] == 0 + assert sum(campaign["whole_policy_arm_counts"].values()) == 2 + assert after["event_chain"]["ok"] is True + assert after["event_chain"]["verified_events"] == 7 + assert after["schema_integrity"]["ok"] is True + assert after["verification_level"] == "local_append_only_chain_without_external_anchor" + assert after["authenticity_proof_available"] is False + assert after["trusted_timestamp_available"] is False + assert after["execution_enabled"] is False + assert after["activation_available"] is False + assert after["automatic_promotion_allowed"] is False + + manifest_arms = { + row["arm_id"]: row + for row in package["assignment_manifest"]["manifest"]["whole_policy_arms"] + } + with sqlite3.connect(sealed["db_path"]) as connection: + rows = connection.execute( + """ + SELECT c.assignment_reveal_commitment, r.assignment_reveal_hash, r.reveal_json + FROM shadow_assignment_commitments c + JOIN shadow_assignment_reveals r USING (commitment_hash) + ORDER BY c.commitment_hash + """ + ).fetchall() + assert len(rows) == 2 + for committed_reveal_hash, actual_reveal_hash, reveal_json in rows: + reveal = json.loads(reveal_json) + assignment = reveal["assignment"] + assert committed_reveal_hash == actual_reveal_hash == reveal["assignment_reveal_hash"] + assert assignment["arm_id"] in manifest_arms + assert assignment["arm_hash"] == manifest_arms[assignment["arm_id"]]["arm_hash"] + assert assignment["arm_probability_bps"] == 5000 + assert assignment["assignment_algorithm"] == SHADOW_ASSIGNMENT_ALGORITHM + assert assignment["boundaries"]["shadow_only"] is True + assert assignment["boundaries"]["ledger_eligible"] is False + + for artifact in (package, capsule, first, second, closure, seed_reveal, after): + _assert_no_floats(artifact) + json.dumps(artifact, sort_keys=True, allow_nan=False) + + +def test_read_only_snapshot_does_not_mutate_registry_and_rejects_writes(tmp_path): + sealed = _sealed(tmp_path) + sealed["registry"].append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="read-only-cluster", + ) + before_database = sealed["db_path"].read_bytes() + + reader = RouteShadowAssignmentRegistry(sealed["db_path"], read_only=True) + snapshot = reader.snapshot(sealed["campaign_id"]) + assert snapshot["ok"] is True + assert snapshot["campaigns"][0]["commitment_count"] == 1 + with pytest.raises(RouteShadowRegistryError, match="opened read-only"): + reader.close_campaign(sealed["campaign_id"]) + + # SQLite may create or retire WAL coordination sidecars while servicing a + # read-only connection, but the durable registry file and logical rows are + # unchanged and the URI mode makes SQL writes impossible. + assert sealed["db_path"].read_bytes() == before_database + with sqlite3.connect(sealed["db_path"]) as connection: + assert connection.execute( + "SELECT COUNT(*) FROM shadow_campaign_closures" + ).fetchone()[0] == 0 + assert connection.execute( + "SELECT COUNT(*) FROM shadow_assignment_commitments" + ).fetchone()[0] == 1 + + +def test_manifest_has_two_explicit_whole_policy_arms_not_route_actions(tmp_path): + package = _sealed(tmp_path)["package"] + manifest = package["assignment_manifest"]["manifest"] + design = package["design_binding"] + arms = manifest["whole_policy_arms"] + + assert manifest["assignment_algorithm"] == SHADOW_ASSIGNMENT_ALGORITHM + assert [row["arm_id"] for row in arms] == [ + "incumbent_source_policy", + "candidate_target_policy", + ] + assert [row["allocation_bps"] for row in arms] == [5000, 5000] + assert sum(row["allocation_bps"] for row in arms) == 10_000 + assert [row["policy_binding"]["binding_type"] for row in arms] == [ + "frozen_source_policy_class", + "frozen_target_policy_class", + ] + assert all( + len(row["policy_binding"]["policy_class_hash"]) == 64 for row in arms + ) + assert arms[0]["policy_binding"]["policy_class_manifest"]["support_strata"] + assert arms[1]["policy_binding"]["policy_class_manifest"]["thresholds"] + assert arms == design["whole_policy_arms"] + assert all(len(row["arm_hash"]) == 64 for row in arms) + arm_json = json.dumps(arms, sort_keys=True) + assert "eligible_actions" not in arm_json + assert "rehearsed_action_probabilities" not in arm_json + + +def test_seed_and_raw_cluster_identity_are_absent_from_db_and_wal_before_reveal(tmp_path): + sealed = _sealed(tmp_path) + raw_cluster = "private-raw-session-identity-never-persisted" + sealed["registry"].append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier=raw_cluster, + ) + + with sqlite3.connect(sealed["db_path"]) as connection: + assert connection.execute("SELECT COUNT(*) FROM shadow_seed_reveals").fetchone()[0] == 0 + stored = connection.execute( + "SELECT commitment_json, cluster_pseudonym FROM shadow_assignment_commitments" + ).fetchone() + assert raw_cluster not in stored[0] + assert raw_cluster != stored[1] + + encoded_seed = base64.urlsafe_b64encode(SEED).decode("ascii").rstrip("=").encode("ascii") + assert encoded_seed.decode("ascii") not in json.dumps(sealed["package"], sort_keys=True) + for path in _database_files(sealed["db_path"]): + contents = path.read_bytes() + assert SEED not in contents + assert encoded_seed not in contents + assert raw_cluster.encode("utf-8") not in contents + + +def test_wrong_seed_or_campaign_capsule_is_rejected(tmp_path): + first = _sealed(tmp_path / "first", bundle=_review_bundle(candidate="a")) + second = _sealed( + tmp_path / "second", + seed=OTHER_SEED, + bundle=_review_bundle(candidate="c", distribution="d"), + ) + + with pytest.raises(ValueError, match="campaign_id does not match"): + first["registry"].append_assignment_commitment( + campaign_id=first["campaign_id"], + seed_capsule=second["capsule"], + cluster_identifier="cluster-one", + ) + + wrong_seed = copy.deepcopy(first["capsule"]) + wrong_seed["seed_material_base64url"] = base64.urlsafe_b64encode(OTHER_SEED).decode( + "ascii" + ).rstrip("=") + with pytest.raises(ValueError, match="does not open"): + prepare_shadow_assignment_commitment( + first["package"], + wrong_seed, + "cluster-one", + ) + + noncanonical = copy.deepcopy(first["capsule"]) + alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" + canonical = noncanonical["seed_material_base64url"] + last_index = alphabet.index(canonical[-1]) + noncanonical["seed_material_base64url"] = canonical[:-1] + alphabet[last_index | 1] + assert base64.urlsafe_b64decode( + noncanonical["seed_material_base64url"] + "=" + ) == SEED + with pytest.raises(ValueError, match="not canonical"): + audit_shadow_seed_capsule(first["package"], noncanonical) + + +def test_assignment_pseudonym_is_derived_internally_and_commitment_is_deterministic(tmp_path): + sealed = _sealed(tmp_path) + caller_supplied_hash_shaped_identifier = "a" * 64 + first = prepare_shadow_assignment_commitment( + sealed["package"], + sealed["capsule"], + caller_supplied_hash_shaped_identifier, + ) + second = prepare_shadow_assignment_commitment( + sealed["package"], + sealed["capsule"], + caller_supplied_hash_shaped_identifier, + ) + + assert first == second + commitment = first["commitment"]["commitment"] + reveal = first["private_reveal"]["assignment"] + assert commitment["cluster_pseudonym"] == reveal["cluster_pseudonym"] + assert commitment["cluster_pseudonym"] != caller_supplied_hash_shaped_identifier + assert commitment["chosen_arm_withheld_until_reveal"] is True + assert "arm_id" not in commitment + arms = { + row["arm_id"]: row + for row in sealed["package"]["assignment_manifest"]["manifest"][ + "whole_policy_arms" + ] + } + assert reveal["arm_hash"] == arms[reveal["arm_id"]]["arm_hash"] + + +def test_illegal_state_transitions_fail_closed(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + campaign_id = sealed["campaign_id"] + + with pytest.raises(RouteShadowRegistryError, match="before closure"): + registry.reveal_seed(campaign_id=campaign_id, seed_capsule=sealed["capsule"]) + with pytest.raises(RouteShadowRegistryError, match="before assignment verification"): + registry.verify_assignment_reveals(campaign_id) + closure = registry.close_campaign(campaign_id) + assert closure["closure"]["frozen_commitment_count"] == 0 + with pytest.raises(RouteShadowRegistryError, match="commitments are closed"): + registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=sealed["capsule"], + cluster_identifier="cluster-two", + ) + with pytest.raises(RouteShadowRegistryError, match="before assignment verification"): + registry.verify_assignment_reveals(campaign_id) + + +def test_seal_commit_close_reveal_and_verify_are_idempotent(tmp_path): + bundle = _review_bundle() + sealed = _sealed(tmp_path, bundle=bundle) + registry = sealed["registry"] + campaign_id = sealed["campaign_id"] + + same_seal = registry.seal_campaign(bundle, SEED) + assert same_seal["created"] is False + assert same_seal["private_seed_capsule"] is None + assert same_seal["seed_capsule_returned_once"] is False + with pytest.raises(ShadowRegistryConflictError, match="different seed commitment"): + registry.seal_campaign(bundle, OTHER_SEED) + + first = registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=sealed["capsule"], + cluster_identifier="cluster-one", + ) + same = registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=sealed["capsule"], + cluster_identifier="cluster-one", + ) + assert first["created"] is True + assert same["created"] is False + assert same["commitment"] == first["commitment"] + + closure = registry.close_campaign(campaign_id) + assert registry.close_campaign(campaign_id) == closure + reveal = registry.reveal_seed(campaign_id=campaign_id, seed_capsule=sealed["capsule"]) + assert registry.reveal_seed(campaign_id=campaign_id, seed_capsule=sealed["capsule"]) == reveal + verified = registry.verify_assignment_reveals(campaign_id) + assert verified["processed"] == verified["matched"] == 1 + assert registry.verify_assignment_reveals(campaign_id)["processed"] == 0 + assert registry.snapshot(campaign_id)["event_chain"]["verified_events"] == 5 + + +def test_planned_cluster_ceiling_is_enforced_before_any_extra_event(tmp_path): + sealed = _sealed(tmp_path, bundle=_review_bundle(planned_clusters=2)) + registry = sealed["registry"] + for cluster in ("cluster-one", "cluster-two"): + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier=cluster, + ) + before = registry.snapshot(sealed["campaign_id"]) + with pytest.raises(RouteShadowRegistryError, match="planned cluster ceiling reached"): + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="cluster-three", + ) + after = registry.snapshot(sealed["campaign_id"]) + assert after["campaigns"][0]["commitment_count"] == 2 + assert after["event_chain"] == before["event_chain"] + + +def test_rehashed_internal_package_tampering_and_duplicate_keys_are_rejected(tmp_path): + package = _sealed(tmp_path)["package"] + tampered = copy.deepcopy(package) + tampered["assignment_manifest"]["manifest"]["whole_policy_arms"][0][ + "allocation_bps" + ] = 4999 + _rehash_public_package(tampered) + with pytest.raises(ValueError, match="manifest does not reconstruct"): + audit_shadow_campaign_artifacts(tampered) + + duplicate = copy.deepcopy(package) + encoded = duplicate["origin_review_bundle"] + duplicate["origin_review_bundle"] = encoded.replace( + '{"bundle":', + '{"bundle":null,"bundle":', + 1, + ) + _rehash_public_package(duplicate) + with pytest.raises(ValueError, match="strict canonical JSON string"): + audit_shadow_campaign_artifacts(duplicate) + + +def test_shadow_native_canonicalization_rejects_floats_and_noncanonical_inputs(tmp_path): + with pytest.raises(ValueError, match="floating-point"): + _canonical_json({"probability": 0.5}) + + sealed = _sealed(tmp_path) + with pytest.raises(ValueError, match="cluster identifier must be a string"): + sealed["registry"].append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier={"caller": "supplied-hash"}, + ) + with pytest.raises(ValueError, match="batch_size must be an integer"): + sealed["registry"].verify_assignment_reveals( + sealed["campaign_id"], + batch_size=1.0, + ) + + +def test_artifact_tables_and_event_log_are_append_only(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="cluster-one", + ) + registry.close_campaign(sealed["campaign_id"]) + registry.reveal_seed(campaign_id=sealed["campaign_id"], seed_capsule=sealed["capsule"]) + registry.verify_assignment_reveals(sealed["campaign_id"]) + + tables = [ + "shadow_registry_metadata", + "shadow_campaign_seals", + "shadow_assignment_commitments", + "shadow_campaign_closures", + "shadow_seed_reveals", + "shadow_assignment_reveals", + "shadow_registry_events", + ] + for table in tables: + with sqlite3.connect(sealed["db_path"]) as connection: + with pytest.raises(sqlite3.IntegrityError, match="append-only"): + connection.execute(f"UPDATE {table} SET rowid = rowid") + with sqlite3.connect(sealed["db_path"]) as connection: + with pytest.raises(sqlite3.IntegrityError, match="append-only"): + connection.execute(f"DELETE FROM {table}") + + +def test_schema_indexes_and_cross_campaign_reveal_guard_are_enforced(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="first-campaign-cluster", + ) + registry.close_campaign(sealed["campaign_id"]) + registry.reveal_seed( + campaign_id=sealed["campaign_id"], seed_capsule=sealed["capsule"] + ) + + second_seed = bytes(range(32, 64)) + second = registry.seal_campaign(_review_bundle(planned_clusters=6), second_seed) + second_package = second["public_package"] + second_campaign = second_package["campaign_seal"]["seal"]["campaign_id"] + second_capsule = second["private_seed_capsule"] + registry.close_campaign(second_campaign) + registry.reveal_seed(campaign_id=second_campaign, seed_capsule=second_capsule) + + with sqlite3.connect(sealed["db_path"]) as connection: + index_names = { + row[1] + for row in connection.execute( + "PRAGMA index_list('shadow_assignment_commitments')" + ).fetchall() + } + reveal_index_names = { + row[1] + for row in connection.execute( + "PRAGMA index_list('shadow_assignment_reveals')" + ).fetchall() + } + assert "shadow_commitments_campaign_order_idx" in index_names + assert "shadow_reveals_campaign_order_idx" in reveal_index_names + commitment_hash = connection.execute( + "SELECT commitment_hash FROM shadow_assignment_commitments WHERE campaign_id = ?", + (sealed["campaign_id"],), + ).fetchone()[0] + with pytest.raises(sqlite3.IntegrityError, match="campaign does not match commitment"): + connection.execute( + """ + INSERT INTO shadow_assignment_reveals( + commitment_hash, campaign_id, verification_status, + assignment_reveal_hash, reveal_json, verified_at_us + ) VALUES (?, ?, 'matched', ?, '{}', 1) + """, + (commitment_hash, second_campaign, "a" * 64), + ) + + +def test_event_chain_detects_local_tampering_if_append_only_trigger_is_removed(tmp_path): + sealed = _sealed(tmp_path) + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_events_no_update") + connection.execute( + "UPDATE shadow_registry_events SET event_hash = ? WHERE event_sequence = 1", + ("f" * 64,), + ) + + snapshot = sealed["registry"].snapshot(sealed["campaign_id"]) + assert snapshot["ok"] is False + assert snapshot["schema_integrity"]["ok"] is False + assert snapshot["schema_integrity"]["missing_triggers"] == [ + "shadow_events_no_update" + ] + assert snapshot["event_chain"]["ok"] is False + assert snapshot["event_chain"]["reason"] == "event_hash_mismatch" + assert snapshot["authenticity_proof_available"] is False + assert snapshot["external_transparency_anchor_available"] is False + + +def test_schema_audit_detects_same_name_noop_trigger_replacement(tmp_path): + sealed = _sealed(tmp_path) + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_commitments_no_update") + connection.execute( + """ + CREATE TRIGGER shadow_commitments_no_update + AFTER INSERT ON shadow_registry_metadata + BEGIN + SELECT 1; + END + """ + ) + + snapshot = sealed["registry"].snapshot(sealed["campaign_id"]) + assert snapshot["ok"] is False + assert snapshot["schema_integrity"]["missing_triggers"] == [] + assert snapshot["schema_integrity"]["definitions_ok"] is False + assert snapshot["schema_integrity"]["definition_fingerprint"] != snapshot[ + "schema_integrity" + ]["expected_definition_fingerprint"] + + +def test_mismatched_reveal_is_never_reported_as_verified_or_complete(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="mismatch-cluster", + ) + registry.close_campaign(sealed["campaign_id"]) + registry.reveal_seed( + campaign_id=sealed["campaign_id"], seed_capsule=sealed["capsule"] + ) + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_commitments_no_update") + connection.execute( + """ + UPDATE shadow_assignment_commitments + SET assignment_reveal_commitment = ? + WHERE campaign_id = ? + """, + ("0" * 64, sealed["campaign_id"]), + ) + + first = registry.verify_assignment_reveals(sealed["campaign_id"]) + assert first["ok"] is False + assert first["processing_complete"] is True + assert first["verification_complete"] is False + assert first["complete"] is False + assert first["total_mismatched"] == 1 + + # An idempotent later batch must not forget the persisted mismatch. + second = registry.verify_assignment_reveals(sealed["campaign_id"]) + assert second["processed"] == 0 + assert second["ok"] is False + assert second["complete"] is False + assert second["total_mismatched"] == 1 + + campaign = registry.snapshot(sealed["campaign_id"])["campaigns"][0] + assert campaign["state"] == "reveal_verification_failed" + assert campaign["lifecycle_state"] == "reveal_verification_complete" + assert campaign["processed_reveal_count"] == 1 + assert campaign["verified_assignment_count"] == 0 + assert campaign["matched_assignment_count"] == 0 + assert campaign["mismatched_assignment_count"] == 1 + + +def test_verify_marks_isolated_commitment_json_corruption_as_mismatch(tmp_path): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="corrupt-json-cluster", + ) + registry.close_campaign(sealed["campaign_id"]) + registry.reveal_seed( + campaign_id=sealed["campaign_id"], seed_capsule=sealed["capsule"] + ) + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_commitments_no_update") + connection.execute( + "UPDATE shadow_assignment_commitments SET commitment_json = '{}'" + ) + + result = registry.verify_assignment_reveals(sealed["campaign_id"]) + assert result["ok"] is False + assert result["complete"] is False + assert result["invalid_commitment_artifacts"] == 1 + assert result["mismatched"] == 1 + assert result["campaign_audit_performed"] is True + assert result["campaign_audit_ok"] is False + + +@pytest.mark.parametrize( + ("table", "trigger", "column", "message"), + [ + ( + "shadow_campaign_closures", + "shadow_closures_no_update", + "closure_json", + "invalid campaign closure", + ), + ( + "shadow_seed_reveals", + "shadow_seed_reveals_no_update", + "reveal_json", + "invalid seed reveal", + ), + ], +) +def test_verify_fails_closed_on_closure_or_seed_artifact_corruption( + tmp_path, table, trigger, column, message +): + sealed = _sealed(tmp_path) + registry = sealed["registry"] + registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="preflight-cluster", + ) + registry.close_campaign(sealed["campaign_id"]) + registry.reveal_seed( + campaign_id=sealed["campaign_id"], seed_capsule=sealed["capsule"] + ) + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute(f"DROP TRIGGER {trigger}") + connection.execute(f"UPDATE {table} SET {column} = '{{}}'") + + with pytest.raises(RouteShadowRegistryError, match=message): + registry.verify_assignment_reveals(sealed["campaign_id"]) + with sqlite3.connect(sealed["db_path"]) as connection: + assert connection.execute( + "SELECT COUNT(*) FROM shadow_assignment_reveals" + ).fetchone()[0] == 0 + + +def test_snapshot_reconstructs_artifacts_and_matches_events_to_evidence_rows(tmp_path): + sealed = _sealed(tmp_path / "package") + with sqlite3.connect(sealed["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_seals_no_update") + connection.execute( + "UPDATE shadow_campaign_seals SET public_package_json = '{}'" + ) + + package_snapshot = sealed["registry"].snapshot(sealed["campaign_id"]) + assert package_snapshot["event_chain"]["ok"] is True + assert package_snapshot["campaigns"][0]["artifact_audit_ok"] is False + assert package_snapshot["ok"] is False + + second = _sealed(tmp_path / "orphan-event") + second["registry"].append_assignment_commitment( + campaign_id=second["campaign_id"], + seed_capsule=second["capsule"], + cluster_identifier="cluster-to-remove", + ) + with sqlite3.connect(second["db_path"]) as connection: + connection.execute("DROP TRIGGER shadow_commitments_no_delete") + connection.execute("DELETE FROM shadow_assignment_commitments") + + orphan_snapshot = second["registry"].snapshot(second["campaign_id"]) + assert orphan_snapshot["event_chain"]["ok"] is True + assert orphan_snapshot["event_artifact_consistency"] == { + "ok": False, + "registered_event_artifacts": 2, + "evidence_artifacts": 1, + "reason": "event_evidence_artifact_mismatch", + } + assert orphan_snapshot["ok"] is False + + +def test_duplicate_cluster_commit_is_linearizable_under_concurrency(tmp_path): + sealed = _sealed(tmp_path) + workers = 8 + barrier = Barrier(workers) + + def append_once(_): + registry = RouteShadowAssignmentRegistry(sealed["db_path"]) + barrier.wait() + return registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="same-cluster", + ) + + with ThreadPoolExecutor(max_workers=workers) as pool: + results = list(pool.map(append_once, range(workers))) + + assert sum(result["created"] is True for result in results) == 1 + assert len({result["commitment"]["commitment_hash"] for result in results}) == 1 + snapshot = sealed["registry"].snapshot(sealed["campaign_id"]) + assert snapshot["campaigns"][0]["commitment_count"] == 1 + assert snapshot["event_chain"]["verified_events"] == 2 + + +def test_close_vs_commit_race_freezes_exactly_the_linearized_count(tmp_path): + sealed = _sealed(tmp_path) + sealed["registry"].append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="existing-cluster", + ) + barrier = Barrier(2) + + def append_second(): + registry = RouteShadowAssignmentRegistry(sealed["db_path"]) + barrier.wait() + try: + return ("append", registry.append_assignment_commitment( + campaign_id=sealed["campaign_id"], + seed_capsule=sealed["capsule"], + cluster_identifier="racing-cluster", + )) + except RouteShadowRegistryError as exc: + return ("append_error", str(exc)) + + def close(): + registry = RouteShadowAssignmentRegistry(sealed["db_path"]) + barrier.wait() + return ("close", registry.close_campaign(sealed["campaign_id"])) + + with ThreadPoolExecutor(max_workers=2) as pool: + results = [pool.submit(append_second), pool.submit(close)] + outcomes = [future.result() for future in results] + + snapshot = sealed["registry"].snapshot(sealed["campaign_id"]) + campaign = snapshot["campaigns"][0] + assert campaign["state"] == "commitments_closed" + assert campaign["frozen_commitment_count"] == campaign["commitment_count"] + assert campaign["commitment_count"] in (1, 2) + append_outcome = next(value for kind, value in outcomes if kind.startswith("append")) + if campaign["commitment_count"] == 1: + assert isinstance(append_outcome, str) + assert "commitments are closed" in append_outcome + else: + assert append_outcome["created"] is True + + +def test_unknown_schema_version_is_rejected_without_migration(tmp_path): + db_path = tmp_path / "future.sqlite3" + with sqlite3.connect(db_path) as connection: + connection.execute("PRAGMA user_version = 999") + before = db_path.read_bytes() + with pytest.raises(RouteShadowRegistryError, match="unsupported shadow registry schema 999"): + RouteShadowAssignmentRegistry(db_path) + assert db_path.read_bytes() == before + + +def test_unrelated_unversioned_database_is_rejected_without_pollution(tmp_path): + db_path = tmp_path / "unrelated.sqlite3" + with sqlite3.connect(db_path) as connection: + connection.execute("CREATE TABLE unrelated_records(value TEXT NOT NULL)") + connection.execute("INSERT INTO unrelated_records(value) VALUES ('keep-me')") + before = db_path.read_bytes() + + with pytest.raises( + RouteShadowRegistryError, + match="refusing to initialize a non-empty unversioned SQLite database", + ): + RouteShadowAssignmentRegistry(db_path) + + assert db_path.read_bytes() == before + with sqlite3.connect(db_path) as connection: + assert connection.execute("PRAGMA user_version").fetchone()[0] == 0 + assert connection.execute("SELECT value FROM unrelated_records").fetchall() == [ + ("keep-me",) + ] + assert connection.execute( + "SELECT COUNT(*) FROM sqlite_master WHERE name LIKE 'shadow_%'" + ).fetchone()[0] == 0 diff --git a/test_route_policy_study_cli.py b/test_route_policy_study_cli.py new file mode 100644 index 000000000..295b3ddc4 --- /dev/null +++ b/test_route_policy_study_cli.py @@ -0,0 +1,89 @@ +import json + +from source.route_policy_study_cli import main + + +def test_cli_example_exports_same_read_only_charter_to_stdout(capsys): + assert main(["--example", "--compact"]) == 0 + captured = capsys.readouterr() + plan = json.loads(captured.out) + assert captured.err == "" + assert plan["study"]["study_id"] == "auto-route-adjacent-explorer-v1" + assert plan["charter"]["causal_boundaries"]["execution_enabled"] is False + assert plan["charter"]["causal_boundaries"]["ledger_write_performed"] is False + assert plan["charter"]["causal_boundaries"]["automatic_promotion_allowed"] is False + assert plan["charter"]["probability_design"]["assignment_performed"] is False + assert plan["charter"]["source_contract"] == { + "policy_id": "auto-route-v2", + "policy_version": "2.0.0", + "feature_schema_version": "route-context-v1", + "support_schema_version": "route-support-v1", + "candidate_set_hash": "a" * 64, + "distribution_hash": "b" * 64, + "outcome_contract_schema_version": "route-outcome-contract-v1", + } + exact = plan["charter"]["traffic_scenario"]["observed_label_scenario"][ + "exact_simultaneous_target" + ] + assert exact["method"] == "exact_joint_multinomial_tail_inversion_two_alternates" + assert exact["minimum_routes_for_target_on_every_alternate_action"] == 1971 + + +def test_cli_rejects_prompt_bearing_or_unknown_input_fields(tmp_path, capsys): + input_path = tmp_path / "support.json" + input_path.write_text( + json.dumps( + { + "baseline_mode": "off", + "post_filter_candidates": [], + "post_filter_exclusions": [], + "prompt": "must not enter a prompt-free charter", + } + ), + encoding="utf-8", + ) + assert main(["--input", str(input_path)]) == 2 + captured = capsys.readouterr() + assert captured.out == "" + assert "unsupported or non-prompt-free fields: prompt" in captured.err + + +def test_cli_writes_explicit_output_without_assigning(tmp_path, capsys): + output_path = tmp_path / "study.json" + assert main(["--example", "--output", str(output_path)]) == 0 + assert capsys.readouterr().out == "" + plan = json.loads(output_path.read_text(encoding="utf-8")) + assert len(plan["design_hash"]) == 64 + assert plan["charter"]["probability_design"]["assignment_performed"] is False + assert plan["charter"]["traffic_scenario"]["observed_label_scenario"][ + "not_power_analysis" + ] is True + + +def test_cli_requires_bound_source_contract(tmp_path, capsys): + input_path = tmp_path / "support.json" + input_path.write_text( + json.dumps( + { + "baseline_mode": "off", + "post_filter_candidates": [], + "post_filter_exclusions": [], + } + ), + encoding="utf-8", + ) + assert main(["--input", str(input_path)]) == 2 + captured = capsys.readouterr() + assert captured.out == "" + assert "input is missing required fields: source_contract" in captured.err + + +def test_cli_accepts_windows_utf8_bom_inputs(tmp_path, capsys): + from source.route_policy_study_cli import _example_payload + + input_path = tmp_path / "support.json" + input_path.write_text(json.dumps(_example_payload()), encoding="utf-8-sig") + assert main(["--input", str(input_path), "--compact"]) == 0 + captured = capsys.readouterr() + assert json.loads(captured.out)["study"]["study_id"] == "auto-route-adjacent-explorer-v1" + assert captured.err == "" From 00705575424c21f83fb83636863f18722f73a9bd Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:37:46 +0100 Subject: [PATCH 2/6] feat: add v51 adaptive inference and stability diagnostics --- launch_v51_ultra_chat_demo.ps1 | 73 + runtime_python/chat_pipeline.py | 131 +- runtime_python/model_variants.py | 8532 +++++++++++++++++- source/benchmark_cognitive_leap_ultra_v51.py | 253 + source/benchmark_v51_prediction_stability.py | 229 + source/materialize_v51_chat_demo.py | 206 + source/model_variants.py | 1142 ++- source/sync_runtime_model_variants.py | 76 + test_cognitive_leap_ultra_expert.py | 201 + test_runtime_model_variants_parity.py | 40 + test_v51_materialize_demo.py | 22 + 11 files changed, 10577 insertions(+), 328 deletions(-) create mode 100644 launch_v51_ultra_chat_demo.ps1 create mode 100644 source/benchmark_cognitive_leap_ultra_v51.py create mode 100644 source/benchmark_v51_prediction_stability.py create mode 100644 source/materialize_v51_chat_demo.py create mode 100644 source/sync_runtime_model_variants.py create mode 100644 test_cognitive_leap_ultra_expert.py create mode 100644 test_runtime_model_variants_parity.py create mode 100644 test_v51_materialize_demo.py diff --git a/launch_v51_ultra_chat_demo.ps1 b/launch_v51_ultra_chat_demo.ps1 new file mode 100644 index 000000000..46b2fcefe --- /dev/null +++ b/launch_v51_ultra_chat_demo.ps1 @@ -0,0 +1,73 @@ +param( + [int]$Port = 8017, + [string]$HostName = "127.0.0.1", + [string]$ArtifactDir = "output\benchmark_v51_cognitive_leap_ultra_latest", + [switch]$NoOpen +) + +$ErrorActionPreference = "Stop" +$Root = Split-Path -Parent $MyInvocation.MyCommand.Path +$ArtifactRoot = if ([System.IO.Path]::IsPathRooted($ArtifactDir)) { + $ArtifactDir +} else { + Join-Path $Root $ArtifactDir +} +$Weights = Join-Path $ArtifactRoot "cognitive_leap_ultra_v51_trained.pth" +$Metrics = Join-Path $ArtifactRoot "benchmark_results.json" +$Meta = Join-Path $ArtifactRoot "chat_demo_meta.json" +$Stdout = Join-Path $ArtifactRoot "chat_web_stdout.log" +$Stderr = Join-Path $ArtifactRoot "chat_web_stderr.log" +$PidFile = Join-Path $ArtifactRoot "chat_web.pid" + +function Quote-ProcessArg { + param([Parameter(Mandatory=$true)][string]$Value) + if ($Value -match '[\s"]') { + return '"' + ($Value -replace '"', '\"') + '"' + } + return $Value +} + +if (-not (Test-Path -LiteralPath $Weights)) { + throw "Missing v51 checkpoint: $Weights. Run source\benchmark_cognitive_leap_ultra_v51.py first." +} +if (-not (Test-Path -LiteralPath $Metrics)) { + throw "Missing v51 benchmark metrics: $Metrics. Run source\benchmark_cognitive_leap_ultra_v51.py first." +} + +Push-Location $Root +try { + python source\materialize_v51_chat_demo.py --weights $Weights --metrics $Metrics --meta $Meta --check + + $existing = Get-NetTCPConnection -LocalPort $Port -State Listen -ErrorAction SilentlyContinue + if ($existing) { + Write-Host "Port $Port is already listening. Reusing http://$HostName`:$Port" + if (-not $NoOpen) { + Start-Process "http://$HostName`:$Port" + } + return + } + + $args = @( + "source\chat_web_app.py", + "--weights", $Weights, + "--meta", $Meta, + "--model_size", "auto", + "--autoload", + "--host", $HostName, + "--port", [string]$Port, + "--reasoning_cycles", "3", + "--adaptive_compute", + "--adaptive_exit_entropy", "0.2", + "--prediction_stability_patience", "2", + "--prediction_stability_tol", "0.005" + ) | ForEach-Object { Quote-ProcessArg $_ } + $process = Start-Process -FilePath "python" -ArgumentList $args -WorkingDirectory $Root -WindowStyle Hidden -PassThru -RedirectStandardOutput $Stdout -RedirectStandardError $Stderr + Set-Content -LiteralPath $PidFile -Value $process.Id + Write-Host "Started v51 ultra chat demo: http://$HostName`:$Port (pid $($process.Id))" + if (-not $NoOpen) { + Start-Process "http://$HostName`:$Port" + } +} +finally { + Pop-Location +} diff --git a/runtime_python/chat_pipeline.py b/runtime_python/chat_pipeline.py index 57ebb1410..28fc15fb2 100644 --- a/runtime_python/chat_pipeline.py +++ b/runtime_python/chat_pipeline.py @@ -22,6 +22,7 @@ "context_mix_v1", "context_mix_v2_mm", "context_mix_v3", + "context_mix_v4", ) TOKEN_RE = re.compile(r"[A-Za-z0-9_']+|[^\w\s]") SENTENCE_RE = re.compile(r"[^.!?]+[.!?]?") @@ -91,6 +92,10 @@ # Optimized for reuse SPLIT_ROLE_RE = re.compile(r"^(user|assistant|system|tool|memory \d+ user|memory \d+ assistant)\s*:\s*(.*)$", re.I) ACTION_REQUEST_RE = re.compile(r"\b(please|can you|could you|help|build|fix|make|optimi[sz]e)\b", re.I) +CONTEXT_V3_IMPERATIVE_RE = re.compile( + r"\b(please|can you|could you|help|make|add|improve|continue|explain|rewrite|optimi[sz]e)\b", + re.I, +) COREF_RE = re.compile(r"\b(it|that|this|they|them|he|she|those|these|same|previous|above)\b", re.I) CONTINUE_RE = re.compile(r"\b(continue|go on|more|deeper|expand|another|next)\b", re.I) CAMEL_OR_WORD_RE = re.compile(r"[A-Z]+(?=[A-Z][a-z]|\d|$)|[A-Z]?[a-z]+|\d+") @@ -827,7 +832,7 @@ def featurize_context_v3(context_text: str, dim: int = FEAT_DIM, max_lines: int latest_user = content if "?" in content: question_count += 1 - if ACTION_REQUEST_RE.search(lower): + if CONTEXT_V3_IMPERATIVE_RE.search(lower): imperative_count += 1 elif role_key in {"assistant", "memory_assistant"}: assistant_turns += 1 @@ -1346,18 +1351,130 @@ def featurize_context_mix_v3(context_text: str, dim: int = FEAT_DIM) -> torch.Te return acc +def featurize_context_mix_v4(context_text: str, dim: int = FEAT_DIM) -> torch.Tensor: + """Frontier control-aware encoder layered on context_mix_v3.""" + raw = (context_text or "").strip() + if not raw: + return torch.zeros(dim, dtype=torch.float32) + + base = featurize_context_mix_v3(raw, dim=dim) + if dim != FEAT_DIM: + return base + + acc = 0.78 * base + 0.16 * featurize_context_v5(raw, dim=dim) + reasoning_budget = "" + creativity_level = "" + knowledge_recency = "" + source_quality = "" + task_mode = "" + research_tags: List[str] = [] + latest_user = "" + + for line in [ln.strip() for ln in raw.splitlines() if ln.strip()]: + role, content = _split_role_line(line) + low = content.lower() + if role == "system": + if low.startswith("reasoning_budget="): + reasoning_budget = content.split("=", 1)[1].strip().lower() + elif low.startswith("creativity_level="): + creativity_level = content.split("=", 1)[1].strip().lower() + elif low.startswith("knowledge_recency="): + knowledge_recency = content.split("=", 1)[1].strip().lower() + elif low.startswith("source_quality="): + source_quality = content.split("=", 1)[1].strip().lower() + elif low.startswith("task_mode="): + task_mode = content.split("=", 1)[1].strip().lower() + elif low.startswith("research_tags="): + research_tags = [ + token.strip().lower() + for token in content.split("=", 1)[1].split(",") + if token.strip() + ] + continue + if role == "user" and content: + latest_user = content + + if reasoning_budget: + acc += 0.07 * featurize_text(f"[reasoning_budget] {reasoning_budget}", dim=dim) + if reasoning_budget == "deep": + acc += 0.06 * featurize_text("[deep_reasoning] decompose verify compare synthesize", dim=dim) + elif reasoning_budget == "medium": + acc += 0.04 * featurize_text("[medium_reasoning] explain compare answer", dim=dim) + elif reasoning_budget == "short": + acc += 0.04 * featurize_text("[short_reasoning] direct answer concise", dim=dim) + + if creativity_level: + acc += 0.07 * featurize_text(f"[creativity_level] {creativity_level}", dim=dim) + if creativity_level == "vivid": + acc += 0.05 * featurize_text("[creative_style] vivid metaphor analogy scene twist", dim=dim) + elif creativity_level == "balanced": + acc += 0.04 * featurize_text("[creative_style] polished engaging but grounded", dim=dim) + elif creativity_level == "precise": + acc += 0.04 * featurize_text("[creative_style] precise grounded technical", dim=dim) + + if knowledge_recency: + acc += 0.07 * featurize_text(f"[knowledge_recency] {knowledge_recency}", dim=dim) + if knowledge_recency == "latest": + acc += 0.06 * featurize_text("[latest_info] current recent update newest paper official release", dim=dim) + elif knowledge_recency == "recent": + acc += 0.04 * featurize_text("[recent_info] recent official information", dim=dim) + elif knowledge_recency == "evergreen": + acc += 0.03 * featurize_text("[evergreen_info] stable concept fundamentals", dim=dim) + + if source_quality: + acc += 0.06 * featurize_text(f"[source_quality] {source_quality}", dim=dim) + if source_quality == "research": + acc += 0.05 * featurize_text("[research_style] empirical result ablation benchmark tradeoff", dim=dim) + elif source_quality == "official": + acc += 0.04 * featurize_text("[official_style] documentation exact behavior supported guidance", dim=dim) + elif source_quality == "synthetic": + acc += 0.03 * featurize_text("[synthetic_style] generated candidate draft", dim=dim) + + if task_mode: + acc += 0.07 * featurize_text(f"[task_mode] {task_mode}", dim=dim) + if task_mode == "knowledge": + acc += 0.04 * featurize_text("[knowledge_task] facts concepts research docs", dim=dim) + elif task_mode == "creative": + acc += 0.04 * featurize_text("[creative_task] brainstorm story concept", dim=dim) + elif task_mode == "reasoning": + acc += 0.04 * featurize_text("[reasoning_task] derive debug evaluate", dim=dim) + elif task_mode == "coding": + acc += 0.04 * featurize_text("[coding_task] code api bug patch", dim=dim) + + for tag in research_tags[:8]: + acc += 0.03 * featurize_text(f"[research_tag] {tag}", dim=dim) + + latest_low = latest_user.lower() + if ( + "latest" in latest_low + or "newest" in latest_low + or "research paper" in latest_low + or "arxiv" in latest_low + ): + acc += 0.05 * featurize_text("[research_latest_request] latest paper current results", dim=dim) + if "creative" in latest_low and ("paper" in latest_low or "research" in latest_low): + acc += 0.04 * featurize_text("[creative_research_blend] analogy explain vividly but correctly", dim=dim) + + norm = torch.norm(acc, p=2) + if norm > 0: + acc = acc / norm + return acc + + def resolve_feature_mode(feature_mode: str, smarter_auto: bool = False) -> str: mode = str(feature_mode or "legacy").strip().lower() if mode not in VALID_FEATURE_MODES: mode = "legacy" - if smarter_auto and mode in {"legacy", "context_v2"}: - return "context_mix_v3" + if smarter_auto and mode in {"legacy", "context_v2", "context_mix_v3"}: + return "context_mix_v4" return mode def text_to_model_input(text: str, feature_mode: str = "legacy") -> torch.Tensor: # ChampionNet expects (B, T, 128) mode = resolve_feature_mode(feature_mode, smarter_auto=False) + if mode == "context_mix_v4": + return featurize_context_mix_v4(text).view(1, 1, FEAT_DIM) if mode == "context_mix_v2_mm": return featurize_context_mix_v2_mm(text).view(1, 1, FEAT_DIM) if mode == "context_mix_v3": @@ -1789,7 +1906,9 @@ def build_training_tensors( feature_mode: str = "legacy", ) -> Tuple[torch.Tensor, torch.Tensor]: mode = resolve_feature_mode(feature_mode, smarter_auto=False) - if mode == "context_mix_v2_mm": + if mode == "context_mix_v4": + feats = [featurize_context_mix_v4(ex.context) for ex in examples] + elif mode == "context_mix_v2_mm": feats = [featurize_context_mix_v2_mm(ex.context) for ex in examples] elif mode == "context_mix_v3": feats = [featurize_context_mix_v3(ex.context) for ex in examples] @@ -1831,7 +1950,9 @@ def build_bucket_metadata( return {"buckets": {}, "label_priors": priors} mode = resolve_feature_mode(feature_mode, smarter_auto=False) - if mode == "context_mix_v2_mm": + if mode == "context_mix_v4": + ctx_featurizer = featurize_context_mix_v4 + elif mode == "context_mix_v2_mm": ctx_featurizer = featurize_context_mix_v2_mm elif mode == "context_mix_v3": ctx_featurizer = featurize_context_mix_v3 diff --git a/runtime_python/model_variants.py b/runtime_python/model_variants.py index 8a28b2192..8fde3df37 100644 --- a/runtime_python/model_variants.py +++ b/runtime_python/model_variants.py @@ -1,30 +1,8520 @@ -"""Runtime proxy for shared ChampionNet model variant definitions. +"""Model variant definitions for the Supermix_27 ChampionNet architecture. -This runtime bundle should use the same architecture registry as ``source/`` -so expert variants and checkpoint detection stay in sync. +This module defines all classifier head variants (from simple ExpandedClassifierHead +through the advanced MoE heads) and their backbone wrappers. It also provides +utility functions for model construction, weight loading with backward +compatibility, and checkpoint introspection. + +Classifier Head Evolution: + ExpandedClassifierHead → Single adapter branch + ExpandedClassifierHeadXL → Dual adapters + router + ExpandedClassifierHeadXXL → Tri-branch + two-stage routing + ExpandedClassifierHeadXXXL → Quad-branch + three-stage routing + ExpandedClassifierHeadUltra → Five branches + domain-expert calibration + ExpandedClassifierHeadMega → Six branches + cross-attention fusion + GatedExpertClassifierHead → MoE with Noisy Top-K Gating + HierarchicalMoEClassifierHead → Two-level hierarchical routing + DeepExpertClassifierHead → Aux-loss-free dynamic bias load balancing + ExpertChoiceClassifierHead → Expert Choice (experts pick tokens) + SmarterExpertClassifierHead → Sigma Gating + LoRA adapters """ -from importlib.util import module_from_spec as _module_from_spec, spec_from_file_location as _spec_from_file_location -from pathlib import Path as _Path +import math +from typing import List, Optional, Sequence, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from run import ChampionNet, GatedFFN -_SOURCE_MODEL_VARIANTS = _Path(__file__).resolve().parents[1] / "source" / "model_variants.py" -_SPEC = _spec_from_file_location("_shared_source_model_variants", _SOURCE_MODEL_VARIANTS) -if _SPEC is None or _SPEC.loader is None: - raise ImportError(f"Could not load shared model variants from {_SOURCE_MODEL_VARIANTS}") -_MODULE = _module_from_spec(_SPEC) -_SPEC.loader.exec_module(_MODULE) +SUPPORTED_MODEL_SIZES: Tuple[str, ...] = ( + "base", + "large", + "xlarge", + "xxlarge", + "xxxlarge", + "ultralarge", + "megalarge", + "ultra_expert", + "hierarchical_expert", + "deep_expert", + "expert_choice", + "smarter_expert", + "thought_expert", + "recursive_expert", + "reflexive_expert", + "metacognitive_expert", + "tree_of_thought_expert", + "consensus_expert", + "deliberative_expert", + "omniscient_expert", + "neurogenesis_expert", + "cognitive_expert", + "cognitive_leap_expert", + "cognitive_leap_ultra_expert", + "transcendent_expert", + "omniversal_expert", + "fractal_expert", + "liquid_spiking_expert", + "active_inference_expert", + "holographic_state_space_expert", + "paper_fusion_expert", + "frontier_expert", + "frontier_collective_expert", + "frontier_verifier_expert", + "test_time_diff_expert", + "titan_dreamer_expert", +) +EXPANSION_DIM_MODEL_SIZES = frozenset({"large", "xlarge", "xxlarge", "xxxlarge", "ultralarge", "megalarge"}) +EXTRA_EXPANSION_DIM_MODEL_SIZES = frozenset({"xlarge", "xxlarge", "xxxlarge", "ultralarge", "megalarge"}) +THIRD_EXPANSION_DIM_MODEL_SIZES = frozenset({"xxlarge", "xxxlarge", "ultralarge", "megalarge"}) +FOURTH_EXPANSION_DIM_MODEL_SIZES = frozenset({"xxxlarge", "ultralarge", "megalarge"}) +FIFTH_EXPANSION_DIM_MODEL_SIZES = frozenset({"ultralarge", "megalarge"}) +SIXTH_EXPANSION_DIM_MODEL_SIZES = frozenset({"megalarge"}) -__doc__ = getattr(_MODULE, "__doc__", __doc__) -for _name, _value in vars(_MODULE).items(): - if _name.startswith("__") and _name not in {"__all__", "__doc__"}: - continue - globals()[_name] = _value +class ExpandedClassifierHead(nn.Module): + """ + Larger classifier head that keeps base checkpoint compatibility: + - retains `weight` and `bias` keys used by the original final linear + - adds an extra adapter branch (256 -> expansion_dim -> 10) + - starts near original behavior with alpha initialized to 0 + """ -__all__ = getattr( - _MODULE, - "__all__", - tuple(name for name in globals() if not name.startswith("_")), -) + def __init__(self, in_dim: int = 256, out_dim: int = 10, expansion_dim: int = 512, dropout: float = 0.1): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.dropout = nn.Dropout(dropout) + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.normal_(self.adapter_up.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down.weight) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + adapter_logits = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + return base_logits + self.alpha * adapter_logits + + +class ExpandedClassifierHeadXL(nn.Module): + """ + Wider classifier head for extra capacity and richer routing: + - keeps base-compatible keys (`weight`, `bias`, `adapter_up/down`, `alpha`) + - adds a second wider adapter branch + learned token-wise router + - remains warm-start compatible from base/large checkpoints + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + expansion_dim: int = 768, + extra_expansion_dim: int = 1024, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Reuse large-head key names so xlarge can warm-start from large checkpoints. + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_b = nn.Linear(in_dim, extra_expansion_dim, bias=False) + self.adapter_down_b = nn.Linear(extra_expansion_dim, out_dim, bias=False) + self.router = nn.Linear(in_dim, 2, bias=True) + self.beta = nn.Parameter(torch.tensor(0.0)) + + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.normal_(self.adapter_up.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down.weight) + nn.init.normal_(self.adapter_up_b.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_b.weight) + nn.init.zeros_(self.router.weight) + nn.init.zeros_(self.router.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + a1 = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + a2 = self.adapter_down_b(self.dropout(F.gelu(self.adapter_up_b(x)))) + + route = torch.softmax(self.router(x), dim=-1) + mix = route[..., :1] * a1 + route[..., 1:2] * a2 + + return base_logits + self.alpha * a1 + self.beta * mix + + +class ExpandedClassifierHeadXXL(nn.Module): + """ + Maximum-capacity classifier head with tri-branch routing: + - retains all xlarge-compatible keys for warm-starting + - adds a third branch and second-stage router for richer fusion + - remains near-base behavior at init via alpha/beta/gamma initialized to 0 + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Keep xlarge key names to preserve compatibility. + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_b = nn.Linear(in_dim, extra_expansion_dim, bias=False) + self.adapter_down_b = nn.Linear(extra_expansion_dim, out_dim, bias=False) + self.router = nn.Linear(in_dim, 2, bias=True) + self.beta = nn.Parameter(torch.tensor(0.0)) + + # New xxl-only branch. + self.adapter_up_c = nn.Linear(in_dim, third_expansion_dim, bias=False) + self.adapter_down_c = nn.Linear(third_expansion_dim, out_dim, bias=False) + self.router2 = nn.Linear(in_dim, 3, bias=True) + self.gamma = nn.Parameter(torch.tensor(0.0)) + + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.normal_(self.adapter_up.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down.weight) + nn.init.normal_(self.adapter_up_b.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_b.weight) + nn.init.zeros_(self.router.weight) + nn.init.zeros_(self.router.bias) + + nn.init.normal_(self.adapter_up_c.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_c.weight) + nn.init.zeros_(self.router2.weight) + nn.init.zeros_(self.router2.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + a1 = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + a2 = self.adapter_down_b(self.dropout(F.gelu(self.adapter_up_b(x)))) + a3 = self.adapter_down_c(self.dropout(F.mish(self.adapter_up_c(x)))) + + route_ab = torch.softmax(self.router(x), dim=-1) + mix_ab = route_ab[..., :1] * a1 + route_ab[..., 1:2] * a2 + + route_all = torch.softmax(self.router2(x), dim=-1) + mix_all = route_all[..., :1] * a1 + route_all[..., 1:2] * a2 + route_all[..., 2:3] * a3 + + return base_logits + self.alpha * a1 + self.beta * mix_ab + self.gamma * mix_all + + +class ExpandedClassifierHeadXXXL(nn.Module): + """ + Highest-capacity classifier head with quad-branch routing: + - preserves xxlarge-compatible keys for warm-starting from xxlarge checkpoints + - adds a fourth branch plus a third routing stage + - initializes new scaling params at 0 to remain stable at startup + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Keep xxlarge key names to preserve compatibility. + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_b = nn.Linear(in_dim, extra_expansion_dim, bias=False) + self.adapter_down_b = nn.Linear(extra_expansion_dim, out_dim, bias=False) + self.router = nn.Linear(in_dim, 2, bias=True) + self.beta = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_c = nn.Linear(in_dim, third_expansion_dim, bias=False) + self.adapter_down_c = nn.Linear(third_expansion_dim, out_dim, bias=False) + self.router2 = nn.Linear(in_dim, 3, bias=True) + self.gamma = nn.Parameter(torch.tensor(0.0)) + + # New xxxlarge-only branch. + self.adapter_up_d = nn.Linear(in_dim, fourth_expansion_dim, bias=False) + self.adapter_down_d = nn.Linear(fourth_expansion_dim, out_dim, bias=False) + self.router3 = nn.Linear(in_dim, 4, bias=True) + self.delta = nn.Parameter(torch.tensor(0.0)) + + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.normal_(self.adapter_up.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down.weight) + nn.init.normal_(self.adapter_up_b.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_b.weight) + nn.init.zeros_(self.router.weight) + nn.init.zeros_(self.router.bias) + + nn.init.normal_(self.adapter_up_c.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_c.weight) + nn.init.zeros_(self.router2.weight) + nn.init.zeros_(self.router2.bias) + + nn.init.normal_(self.adapter_up_d.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_d.weight) + nn.init.zeros_(self.router3.weight) + nn.init.zeros_(self.router3.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + a1 = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + a2 = self.adapter_down_b(self.dropout(F.gelu(self.adapter_up_b(x)))) + a3 = self.adapter_down_c(self.dropout(F.mish(self.adapter_up_c(x)))) + a4 = self.adapter_down_d(self.dropout(F.relu(self.adapter_up_d(x)))) + + route_ab = torch.softmax(self.router(x), dim=-1) + mix_ab = route_ab[..., :1] * a1 + route_ab[..., 1:2] * a2 + + route_abc = torch.softmax(self.router2(x), dim=-1) + mix_abc = route_abc[..., :1] * a1 + route_abc[..., 1:2] * a2 + route_abc[..., 2:3] * a3 + + route_abcd = torch.softmax(self.router3(x), dim=-1) + mix_abcd = ( + route_abcd[..., :1] * a1 + + route_abcd[..., 1:2] * a2 + + route_abcd[..., 2:3] * a3 + + route_abcd[..., 3:4] * a4 + ) + + return base_logits + self.alpha * a1 + self.beta * mix_ab + self.gamma * mix_abc + self.delta * mix_abcd + + +class ExpandedClassifierHeadUltra(nn.Module): + """ + Maximum-capacity classifier head with five routed branches: + - keeps xxxlarge-compatible keys for warm-starting + - adds a fifth branch + fourth routing stage + - adds domain-expert calibration (science/code/writing/math friendly) for smarter adaptation + - initializes new scale parameters at 0 for stable adaptation + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + fifth_expansion_dim: int = 6144, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Keep xxxlarge keys for compatibility. + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_b = nn.Linear(in_dim, extra_expansion_dim, bias=False) + self.adapter_down_b = nn.Linear(extra_expansion_dim, out_dim, bias=False) + self.router = nn.Linear(in_dim, 2, bias=True) + self.beta = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_c = nn.Linear(in_dim, third_expansion_dim, bias=False) + self.adapter_down_c = nn.Linear(third_expansion_dim, out_dim, bias=False) + self.router2 = nn.Linear(in_dim, 3, bias=True) + self.gamma = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_d = nn.Linear(in_dim, fourth_expansion_dim, bias=False) + self.adapter_down_d = nn.Linear(fourth_expansion_dim, out_dim, bias=False) + self.router3 = nn.Linear(in_dim, 4, bias=True) + self.delta = nn.Parameter(torch.tensor(0.0)) + + # New ultralarge-only branch. + self.adapter_up_e = nn.Linear(in_dim, fifth_expansion_dim, bias=False) + self.adapter_down_e = nn.Linear(fifth_expansion_dim, out_dim, bias=False) + self.router4 = nn.Linear(in_dim, 5, bias=True) + self.epsilon = nn.Parameter(torch.tensor(0.0)) + + # Architecture upgrade: low-cost domain-expert correction + bounded calibration. + self.pre_norm = nn.LayerNorm(in_dim) + self.domain_router = nn.Linear(in_dim, 4, bias=True) + self.domain_experts = nn.Linear(in_dim, out_dim * 4, bias=False) + self.calib_gate = nn.Linear(in_dim, out_dim, bias=True) + self.zeta = nn.Parameter(torch.tensor(0.0)) + self.theta = nn.Parameter(torch.tensor(0.0)) + + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.normal_(self.adapter_up.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down.weight) + nn.init.normal_(self.adapter_up_b.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_b.weight) + nn.init.zeros_(self.router.weight) + nn.init.zeros_(self.router.bias) + + nn.init.normal_(self.adapter_up_c.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_c.weight) + nn.init.zeros_(self.router2.weight) + nn.init.zeros_(self.router2.bias) + + nn.init.normal_(self.adapter_up_d.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_d.weight) + nn.init.zeros_(self.router3.weight) + nn.init.zeros_(self.router3.bias) + + nn.init.normal_(self.adapter_up_e.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.adapter_down_e.weight) + nn.init.zeros_(self.router4.weight) + nn.init.zeros_(self.router4.bias) + + nn.init.ones_(self.pre_norm.weight) + nn.init.zeros_(self.pre_norm.bias) + nn.init.zeros_(self.domain_router.weight) + nn.init.zeros_(self.domain_router.bias) + nn.init.normal_(self.domain_experts.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.calib_gate.weight) + nn.init.zeros_(self.calib_gate.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + a1 = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + a2 = self.adapter_down_b(self.dropout(F.gelu(self.adapter_up_b(x)))) + a3 = self.adapter_down_c(self.dropout(F.mish(self.adapter_up_c(x)))) + a4 = self.adapter_down_d(self.dropout(F.relu(self.adapter_up_d(x)))) + a5 = self.adapter_down_e(self.dropout(F.selu(self.adapter_up_e(x)))) + + route_ab = torch.softmax(self.router(x), dim=-1) + mix_ab = route_ab[..., :1] * a1 + route_ab[..., 1:2] * a2 + + route_abc = torch.softmax(self.router2(x), dim=-1) + mix_abc = route_abc[..., :1] * a1 + route_abc[..., 1:2] * a2 + route_abc[..., 2:3] * a3 + + route_abcd = torch.softmax(self.router3(x), dim=-1) + mix_abcd = ( + route_abcd[..., :1] * a1 + + route_abcd[..., 1:2] * a2 + + route_abcd[..., 2:3] * a3 + + route_abcd[..., 3:4] * a4 + ) + + route_abcde = torch.softmax(self.router4(x), dim=-1) + mix_abcde = ( + route_abcde[..., :1] * a1 + + route_abcde[..., 1:2] * a2 + + route_abcde[..., 2:3] * a3 + + route_abcde[..., 3:4] * a4 + + route_abcde[..., 4:5] * a5 + ) + + # Domain-expert calibration branch improves robustness across diverse domains. + h = self.pre_norm(x) + dom_logits = self.domain_experts(self.dropout(F.silu(h))) + dom_logits = dom_logits.view(*dom_logits.shape[:-1], 4, base_logits.shape[-1]) + dom_w = torch.softmax(self.domain_router(h), dim=-1).unsqueeze(-1) + dom_mix = (dom_w * dom_logits).sum(dim=-2) + calib = torch.tanh(self.calib_gate(h)) + + return ( + base_logits + + self.alpha * a1 + + self.beta * mix_ab + + self.gamma * mix_abc + + self.delta * mix_abcd + + self.epsilon * mix_abcde + + self.zeta * dom_mix + + self.theta * calib + ) + + +class CrossAttentionFusion(nn.Module): + """ + Lightweight cross-attention module that lets adapter branches attend to each other. + Input: (B, T, N_branches, out_dim) tensor of stacked branch outputs. + Output: (B, T, out_dim) fused representation. + """ + def __init__(self, out_dim: int = 10, n_branches: int = 6, n_heads: int = 2): + super().__init__() + self.n_heads = n_heads + self.head_dim = max(1, out_dim // n_heads) + inner = self.n_heads * self.head_dim + self.q_proj = nn.Linear(out_dim, inner, bias=False) + self.k_proj = nn.Linear(out_dim, inner, bias=False) + self.v_proj = nn.Linear(out_dim, inner, bias=False) + self.out_proj = nn.Linear(inner, out_dim, bias=False) + self.scale = self.head_dim ** -0.5 + self.fusion_weight = nn.Linear(out_dim, 1, bias=True) + + def forward(self, branch_stack): + # branch_stack: (B, T, N, D) + B, T, N, D = branch_stack.shape + flat = branch_stack.view(B * T, N, D) + q = self.q_proj(flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale + attn = torch.softmax(attn, dim=-1) + out = torch.matmul(attn, v) # (B*T, heads, N, head_dim) + out = out.transpose(1, 2).contiguous().view(B * T, N, self.n_heads * self.head_dim) + out = self.out_proj(out) # (B*T, N, D) + # Weighted sum over branches + w = torch.softmax(self.fusion_weight(out), dim=1) # (B*T, N, 1) + fused = (w * out).sum(dim=1) # (B*T, D) + return fused.view(B, T, D) + + +class ExpandedClassifierHeadMega(nn.Module): + """ + Maximum-capacity classifier head with six routed branches + cross-attention fusion: + - keeps ultralarge-compatible keys for warm-starting + - adds a sixth branch (Mish activation) + fifth routing stage + - adds CrossAttentionFusion for inter-branch communication + - adds a reasoning_gate that amplifies logit differences for sharper routing + - initializes new scale parameters at 0 for stable adaptation + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + fifth_expansion_dim: int = 6144, + sixth_expansion_dim: int = 8192, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Keep ultralarge keys for compatibility. + self.adapter_up = nn.Linear(in_dim, expansion_dim, bias=False) + self.adapter_down = nn.Linear(expansion_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_b = nn.Linear(in_dim, extra_expansion_dim, bias=False) + self.adapter_down_b = nn.Linear(extra_expansion_dim, out_dim, bias=False) + self.router = nn.Linear(in_dim, 2, bias=True) + self.beta = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_c = nn.Linear(in_dim, third_expansion_dim, bias=False) + self.adapter_down_c = nn.Linear(third_expansion_dim, out_dim, bias=False) + self.router2 = nn.Linear(in_dim, 3, bias=True) + self.gamma = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_d = nn.Linear(in_dim, fourth_expansion_dim, bias=False) + self.adapter_down_d = nn.Linear(fourth_expansion_dim, out_dim, bias=False) + self.router3 = nn.Linear(in_dim, 4, bias=True) + self.delta = nn.Parameter(torch.tensor(0.0)) + + self.adapter_up_e = nn.Linear(in_dim, fifth_expansion_dim, bias=False) + self.adapter_down_e = nn.Linear(fifth_expansion_dim, out_dim, bias=False) + self.router4 = nn.Linear(in_dim, 5, bias=True) + self.epsilon = nn.Parameter(torch.tensor(0.0)) + + # Domain-expert calibration (from ultralarge). + self.pre_norm = nn.LayerNorm(in_dim) + self.domain_router = nn.Linear(in_dim, 4, bias=True) + self.domain_experts = nn.Linear(in_dim, out_dim * 4, bias=False) + self.calib_gate = nn.Linear(in_dim, out_dim, bias=True) + self.zeta = nn.Parameter(torch.tensor(0.0)) + self.theta = nn.Parameter(torch.tensor(0.0)) + + # --- Megalarge-only additions --- + # Sixth branch with Mish activation. + self.adapter_up_f = nn.Linear(in_dim, sixth_expansion_dim, bias=False) + self.adapter_down_f = nn.Linear(sixth_expansion_dim, out_dim, bias=False) + self.router5 = nn.Linear(in_dim, 6, bias=True) + self.iota = nn.Parameter(torch.tensor(0.0)) + + # Cross-attention fusion lets branches attend to each other. + self.cross_attn_fusion = CrossAttentionFusion(out_dim=out_dim, n_branches=6, n_heads=2) + self.kappa = nn.Parameter(torch.tensor(0.0)) + + # Reasoning gate amplifies logit differences for sharper routing. + self.reasoning_gate = nn.Sequential( + nn.Linear(in_dim, in_dim // 4, bias=False), + nn.SiLU(), + nn.Linear(in_dim // 4, out_dim, bias=True), + ) + self.lambda_ = nn.Parameter(torch.tensor(0.0)) + + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + for name in ['adapter_up', 'adapter_up_b', 'adapter_up_c', 'adapter_up_d', 'adapter_up_e', 'adapter_up_f']: + nn.init.normal_(getattr(self, name).weight, mean=0.0, std=0.02) + for name in ['adapter_down', 'adapter_down_b', 'adapter_down_c', 'adapter_down_d', 'adapter_down_e', 'adapter_down_f']: + nn.init.zeros_(getattr(self, name).weight) + for name in ['router', 'router2', 'router3', 'router4', 'router5', 'domain_router']: + nn.init.zeros_(getattr(self, name).weight) + nn.init.zeros_(getattr(self, name).bias) + nn.init.ones_(self.pre_norm.weight) + nn.init.zeros_(self.pre_norm.bias) + nn.init.normal_(self.domain_experts.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.calib_gate.weight) + nn.init.zeros_(self.calib_gate.bias) + # Reasoning gate init + for m in self.reasoning_gate: + if hasattr(m, 'weight'): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + if hasattr(m, 'bias') and m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + a1 = self.adapter_down(self.dropout(F.silu(self.adapter_up(x)))) + a2 = self.adapter_down_b(self.dropout(F.gelu(self.adapter_up_b(x)))) + a3 = self.adapter_down_c(self.dropout(F.mish(self.adapter_up_c(x)))) + a4 = self.adapter_down_d(self.dropout(F.relu(self.adapter_up_d(x)))) + a5 = self.adapter_down_e(self.dropout(F.selu(self.adapter_up_e(x)))) + a6 = self.adapter_down_f(self.dropout(F.mish(self.adapter_up_f(x)))) + + route_ab = torch.softmax(self.router(x), dim=-1) + mix_ab = route_ab[..., :1] * a1 + route_ab[..., 1:2] * a2 + + route_abc = torch.softmax(self.router2(x), dim=-1) + mix_abc = route_abc[..., :1] * a1 + route_abc[..., 1:2] * a2 + route_abc[..., 2:3] * a3 + + route_abcd = torch.softmax(self.router3(x), dim=-1) + mix_abcd = ( + route_abcd[..., :1] * a1 + + route_abcd[..., 1:2] * a2 + + route_abcd[..., 2:3] * a3 + + route_abcd[..., 3:4] * a4 + ) + + route_abcde = torch.softmax(self.router4(x), dim=-1) + mix_abcde = ( + route_abcde[..., :1] * a1 + + route_abcde[..., 1:2] * a2 + + route_abcde[..., 2:3] * a3 + + route_abcde[..., 3:4] * a4 + + route_abcde[..., 4:5] * a5 + ) + + route_all = torch.softmax(self.router5(x), dim=-1) + mix_all = ( + route_all[..., :1] * a1 + + route_all[..., 1:2] * a2 + + route_all[..., 2:3] * a3 + + route_all[..., 3:4] * a4 + + route_all[..., 4:5] * a5 + + route_all[..., 5:6] * a6 + ) + + # Domain-expert calibration branch. + h = self.pre_norm(x) + dom_logits = self.domain_experts(self.dropout(F.silu(h))) + dom_logits = dom_logits.view(*dom_logits.shape[:-1], 4, base_logits.shape[-1]) + dom_w = torch.softmax(self.domain_router(h), dim=-1).unsqueeze(-1) + dom_mix = (dom_w * dom_logits).sum(dim=-2) + calib = torch.tanh(self.calib_gate(h)) + + # Cross-attention fusion: stack all 6 branches and let them attend. + branch_stack = torch.stack([a1, a2, a3, a4, a5, a6], dim=-2) # (B,T,6,out_dim) + cross_fused = self.cross_attn_fusion(branch_stack) + + # Reasoning gate: amplify logit differences for sharper decisions. + reason = torch.tanh(self.reasoning_gate(x)) + + return ( + base_logits + + self.alpha * a1 + + self.beta * mix_ab + + self.gamma * mix_abc + + self.delta * mix_abcd + + self.epsilon * mix_abcde + + self.iota * mix_all + + self.zeta * dom_mix + + self.theta * calib + + self.kappa * cross_fused + + self.lambda_ * reason + ) + + +class GatedExpertClassifierHead(nn.Module): + """ + Advanced routing head with specialized experts: + - Top-K gating selects most relevant experts per input + - Experts utilize diverse activation functions and expansion ratios + - Residual calibration ensures logit stability + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 6, + top_k: int = 2, + dropout: float = 0.1 + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.top_k = top_k + self.n_experts = n_experts + + # Experts with varying architectures + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, 1024 + (i * 512), bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(1024 + (i * 512), out_dim, bias=False) for i in range(n_experts) + ]) + self.activations = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + + self.gate = nn.Linear(in_dim, n_experts, bias=False) + self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + for eup in self.experts_up: + nn.init.normal_(eup.weight, mean=0.0, std=0.01) + for edown in self.experts_down: + nn.init.zeros_(edown.weight) + nn.init.zeros_(self.gate.weight) + nn.init.zeros_(self.noise_gate.weight) + nn.init.zeros_(self.calibration.weight) + nn.init.zeros_(self.calibration.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + # Gating logic with Noisy Top-K + clean_logits = self.gate(x) + if self.training: + noise_std = F.softplus(self.noise_gate(x)) + noise = torch.randn_like(clean_logits) * noise_std + gate_logits = clean_logits + noise + else: + gate_logits = clean_logits + + weights = torch.softmax(gate_logits, dim=-1) + top_weights, top_idx = torch.topk(weights, k=self.top_k, dim=-1) + top_weights = top_weights / (top_weights.sum(dim=-1, keepdim=True) + 1e-6) + + # Expert processing + expert_outputs = [] + for i in range(self.n_experts): + out = self.experts_down[i](self.dropout(self.activations[i](self.experts_up[i](x)))) + expert_outputs.append(out) + expert_stack = torch.stack(expert_outputs, dim=-2) # (B, T, N, D) + + # Gather top experts + # We need to broadcast the indices for gather + gather_idx = top_idx.unsqueeze(-1).expand(*top_idx.shape, base_logits.shape[-1]) + selected_experts = torch.gather(expert_stack, -2, gather_idx) # (B, T, K, D) + + expert_logits = (selected_experts * top_weights.unsqueeze(-1)).sum(dim=-2) + calib = self.calibration(x) + + return base_logits + self.alpha * expert_logits + self.theta * calib + + +class HierarchicalMoEClassifierHead(nn.Module): + """ + Next-generation classifier head with hierarchical two-level MoE routing: + - Domain-level gating selects which expert group is most relevant + - Per-domain expert gating selects the best expert(s) within the group + - A shared 'always-on' expert provides a stable baseline signal + - Per-expert residual LoRA adapters add lightweight correction capacity + - Per-expert LayerNorm on outputs improves gradient flow + - Auxiliary load-balancing loss prevents expert collapse during training + Inspired by DeepSeek-MoE, GShard, and ST-MoE research. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_domain_groups: int = 2, + experts_per_group: int = 4, + top_k_domains: int = 1, + top_k_experts: int = 2, + lora_rank: int = 64, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_domain_groups = n_domain_groups + self.experts_per_group = experts_per_group + self.n_experts = n_domain_groups * experts_per_group + self.top_k_domains = top_k_domains + self.top_k_experts = top_k_experts + self._aux_loss = torch.tensor(0.0) + + # --- Shared expert (always active, not gated) --- + shared_dim = 1024 + self.shared_up = nn.Linear(in_dim, shared_dim, bias=False) + self.shared_down = nn.Linear(shared_dim, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(0.0)) + + # --- Domain-level gating --- + self.domain_gate = nn.Linear(in_dim, n_domain_groups, bias=False) + self.domain_noise_gate = nn.Linear(in_dim, n_domain_groups, bias=False) + + # --- Per-domain expert gating --- + self.expert_gates = nn.ModuleList([ + nn.Linear(in_dim, experts_per_group, bias=False) + for _ in range(n_domain_groups) + ]) + self.expert_noise_gates = nn.ModuleList([ + nn.Linear(in_dim, experts_per_group, bias=False) + for _ in range(n_domain_groups) + ]) + + # --- Experts with varying expansion dims & activations --- + # 8 experts total: group 0 (analytical): SiLU, GELU, Mish, ReLU + # group 1 (creative): SELU, Tanh, Softplus, Mish-variant + expansion_dims = [1024, 1536, 2048, 2560, 1024, 1536, 2048, 3072] + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, expansion_dims[i], bias=False) for i in range(self.n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(expansion_dims[i], out_dim, bias=False) for i in range(self.n_experts) + ]) + self.activations = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh, F.softplus, F.mish] + + # --- Per-expert output LayerNorm --- + self.expert_norms = nn.ModuleList([ + nn.LayerNorm(out_dim) for _ in range(self.n_experts) + ]) + + # --- Per-expert residual LoRA adapters --- + self.lora_down = nn.ModuleList([ + nn.Linear(in_dim, lora_rank, bias=False) for _ in range(self.n_experts) + ]) + self.lora_up = nn.ModuleList([ + nn.Linear(lora_rank, out_dim, bias=False) for _ in range(self.n_experts) + ]) + + # --- Aggregation --- + self.alpha = nn.Parameter(torch.tensor(0.0)) # expert mixture scale + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) # calibration scale + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + # Shared expert + nn.init.normal_(self.shared_up.weight, mean=0.0, std=0.01) + nn.init.zeros_(self.shared_down.weight) + nn.init.ones_(self.shared_norm.weight) + nn.init.zeros_(self.shared_norm.bias) + # Domain gates + nn.init.zeros_(self.domain_gate.weight) + nn.init.zeros_(self.domain_noise_gate.weight) + # Per-domain expert gates + for g in self.expert_gates: + nn.init.zeros_(g.weight) + for g in self.expert_noise_gates: + nn.init.zeros_(g.weight) + # Experts + for eup in self.experts_up: + nn.init.normal_(eup.weight, mean=0.0, std=0.01) + for edown in self.experts_down: + nn.init.zeros_(edown.weight) + # Expert norms + for norm in self.expert_norms: + nn.init.ones_(norm.weight) + nn.init.zeros_(norm.bias) + # LoRA adapters + for ld in self.lora_down: + nn.init.normal_(ld.weight, mean=0.0, std=0.01) + for lu in self.lora_up: + nn.init.zeros_(lu.weight) + # Calibration + nn.init.zeros_(self.calibration.weight) + nn.init.zeros_(self.calibration.bias) + + def _compute_aux_loss(self, gate_probs, expert_mask): + """ + Auxiliary load-balancing loss (Switch Transformer style). + Encourages uniform expert utilization. + gate_probs: (B*T, N) softmax probabilities + expert_mask: (B*T, N) binary mask of selected experts + """ + # Fraction of tokens dispatched to each expert + f = expert_mask.float().mean(dim=0) # (N,) + # Average gate probability for each expert + p = gate_probs.mean(dim=0) # (N,) + n = gate_probs.shape[-1] + return n * (f * p).sum() + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] # (B, T) or (B,) + x_flat = x.reshape(-1, x.shape[-1]) # (BT, in_dim) + BT = x_flat.shape[0] + + # --- Shared expert (always active) --- + shared_out = self.shared_down(self.dropout(F.silu(self.shared_up(x)))) + shared_out = self.shared_norm(shared_out) + + # --- Domain-level gating --- + clean_domain = self.domain_gate(x_flat) # (BT, n_groups) + if self.training: + noise_std = F.softplus(self.domain_noise_gate(x_flat)) + domain_logits = clean_domain + torch.randn_like(clean_domain) * noise_std + else: + domain_logits = clean_domain + domain_probs = torch.softmax(domain_logits, dim=-1) # (BT, n_groups) + + # Select top-k domains + top_dom_w, top_dom_idx = torch.topk(domain_probs, k=self.top_k_domains, dim=-1) + top_dom_w = top_dom_w / (top_dom_w.sum(dim=-1, keepdim=True) + 1e-6) + + # --- Process all experts --- + all_expert_outs = [] + for i in range(self.n_experts): + act = self.activations[i] + main_out = self.experts_down[i](self.dropout(act(self.experts_up[i](x_flat)))) + lora_out = self.lora_up[i](self.dropout(F.silu(self.lora_down[i](x_flat)))) + expert_out = self.expert_norms[i](main_out + lora_out) + all_expert_outs.append(expert_out) + expert_stack = torch.stack(all_expert_outs, dim=1) # (BT, n_experts, out_dim) + + # --- Per-domain expert gating & selection --- + aggregated = x_flat.new_zeros(BT, base_logits.shape[-1]) + all_gate_probs_list = [] + all_expert_mask_list = [] + + for d in range(self.n_domain_groups): + # Expert gating within this domain + clean_exp = self.expert_gates[d](x_flat) # (BT, experts_per_group) + if self.training: + noise_std = F.softplus(self.expert_noise_gates[d](x_flat)) + exp_logits = clean_exp + torch.randn_like(clean_exp) * noise_std + else: + exp_logits = clean_exp + exp_probs = torch.softmax(exp_logits, dim=-1) # (BT, experts_per_group) + + top_exp_w, top_exp_idx = torch.topk(exp_probs, k=self.top_k_experts, dim=-1) + top_exp_w = top_exp_w / (top_exp_w.sum(dim=-1, keepdim=True) + 1e-6) + + # Map local expert indices to global + global_idx = top_exp_idx + d * self.experts_per_group # (BT, top_k_experts) + gather_idx = global_idx.unsqueeze(-1).expand(-1, -1, base_logits.shape[-1]) + selected = torch.gather(expert_stack, 1, gather_idx) # (BT, top_k, out_dim) + domain_expert_out = (selected * top_exp_w.unsqueeze(-1)).sum(dim=1) # (BT, out_dim) + + # Weight by domain probability + # Get domain d's weight for each token + # Check if domain d is in the top-k for each token + domain_d_mask = (top_dom_idx == d).any(dim=-1) # (BT,) + # Get the weight for domain d where it was selected + domain_d_weight = x_flat.new_zeros(BT) + for k_idx in range(self.top_k_domains): + mask_k = (top_dom_idx[:, k_idx] == d) + domain_d_weight = domain_d_weight + mask_k.float() * top_dom_w[:, k_idx] + + aggregated = aggregated + domain_d_weight.unsqueeze(-1) * domain_expert_out + + # Collect for aux loss + all_gate_probs_list.append(exp_probs) + exp_mask = torch.zeros(BT, self.experts_per_group, device=x_flat.device) + exp_mask.scatter_(1, top_exp_idx, 1.0) + all_expert_mask_list.append(exp_mask) + + # Reshape aggregated back to original shape + aggregated = aggregated.view(*shape_prefix, -1) + + # --- Auxiliary load-balancing loss --- + if self.training: + # Domain-level balance + dom_mask = torch.zeros(BT, self.n_domain_groups, device=x_flat.device) + dom_mask.scatter_(1, top_dom_idx, 1.0) + aux_domain = self._compute_aux_loss(domain_probs, dom_mask) + # Expert-level balance (per group) + aux_expert = sum( + self._compute_aux_loss(gp, em) + for gp, em in zip(all_gate_probs_list, all_expert_mask_list) + ) / max(1, self.n_domain_groups) + self._aux_loss = aux_domain + aux_expert + else: + self._aux_loss = torch.tensor(0.0, device=x.device) + + calib = self.calibration(x) + + return ( + base_logits + + self.shared_scale * shared_out + + self.alpha * aggregated + + self.theta * calib + ) + + +class DeepExpertClassifierHead(nn.Module): + """ + State-of-the-art MoE classifier head featuring: + - 1 always-active shared expert for universal knowledge. + - N routed experts using Top-K gating. + - Auxiliary-loss-free load balancing (dynamic bias adjustment during training). + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 8, + top_k: int = 2, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.top_k = min(top_k, n_experts) + self.n_experts = n_experts + + # Shared Expert (always active) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # Routed Experts + # Using a mix of activations + varying inner dimensions for diversity + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, 1024 + (i * 256), bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(1024 + (i * 256), out_dim, bias=False) for i in range(n_experts) + ]) + # Cycle activations if n_experts > 6 + base_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.activations = [base_acts[i % len(base_acts)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Routing mechanism + self.gate = nn.Linear(in_dim, n_experts, bias=False) + self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) + + # Aux-free load balancing terms + # These biases are added to the routing logits during training. + # They don't require gradients; we update them manually. + self.register_buffer("expert_bias", torch.zeros(n_experts)) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.experts_down[i].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.shared_down.weight, a=math.sqrt(5)) + nn.init.normal_(self.gate.weight, std=0.02) + nn.init.normal_(self.noise_gate.weight, std=0.02) + nn.init.constant_(self.calibration.weight, 0) + nn.init.constant_(self.calibration.bias, 0) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + BT = x_flat.shape[0] + + # Shared Expert + shared_out = self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + shared_out = self.shared_norm(shared_out) + + # Routed Experts Gating + clean_logits = self.gate(x_flat) + if self.training: + noise_std = F.softplus(self.noise_gate(x_flat)) + noise = torch.randn_like(clean_logits) * noise_std + gate_logits = clean_logits + noise + self.expert_bias + else: + gate_logits = clean_logits + + weights = torch.softmax(gate_logits, dim=-1) + top_weights, top_idx = torch.topk(weights, k=self.top_k, dim=-1) + top_weights = top_weights / (top_weights.sum(dim=-1, keepdim=True) + 1e-6) + + # Process Routed Experts + all_expert_outs = [] + for i in range(self.n_experts): + act = self.activations[i] + out = self.experts_down[i](self.dropout(act(self.experts_up[i](x_flat)))) + out = self.expert_norms[i](out) + all_expert_outs.append(out) + expert_stack = torch.stack(all_expert_outs, dim=1) # (BT, n_experts, out_dim) + + # Gather \u0026 combine chosen experts + gather_idx = top_idx.unsqueeze(-1).expand(-1, -1, base_logits.shape[-1]) + selected_experts = torch.gather(expert_stack, 1, gather_idx) # (BT, K, out_dim) + routed_out = (selected_experts * top_weights.unsqueeze(-1)).sum(dim=1) # (BT, out_dim) + + # Reshape combinations back + routed_out = routed_out.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + calib = self.calibration(x) + + # Update expert bias tracking internally during training if needed (aux-free load balance) + # Usually implemented in the loop, but we can store load stats here + if self.training: + # Fraction of tokens each expert received + with torch.no_grad(): + expert_mask = torch.zeros(BT, self.n_experts, device=x.device) + expert_mask.scatter_(1, top_idx, 1.0) + expert_load = expert_mask.mean(dim=0) # (n_experts,) + target_load = self.top_k / self.n_experts + # Adjust bias: decrease bias for overloaded experts, increase for underloaded + alpha_bias = 0.1 # learning rate for bias + self.expert_bias.add_(alpha_bias * (target_load - expert_load)) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * routed_out + + self.theta * calib + ) + + +class ExpertChoiceClassifierHead(nn.Module): + """ + Expert Choice (EC) Routing Classifier Head: + Instead of tokens selecting the top-K experts, experts select the top-K tokens. + This guarantees 100% load balancing across experts and avoids dropped tokens for active experts. + Tokens can be routed to a variable number of experts based on their relevance. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 8, + capacity_factor: float = 2.0, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + self.n_experts = n_experts + # How many tokens an expert will pick, relative to average token load per expert + self.capacity_factor = capacity_factor + + # Routed Experts + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, 1024 + (i * 256), bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(1024 + (i * 256), out_dim, bias=False) for i in range(n_experts) + ]) + base_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.activations = [base_acts[i % len(base_acts)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Routing mechanism: Token-to-expert affinity + # We need a score for each (token, expert) pair. + self.gate = nn.Linear(in_dim, n_experts, bias=False) + self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.experts_down[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.gate.weight, std=0.02) + nn.init.normal_(self.noise_gate.weight, std=0.02) + nn.init.constant_(self.calibration.weight, 0) + nn.init.constant_(self.calibration.bias, 0) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) # (BT, D) + BT = x_flat.shape[0] + + # Expert token capacity (how many tokens each expert will select) + # Average load per expert is BT / n_experts. We multiply by capacity_factor. + # e.g., if capacity_factor is 2.0, each expert processes 2 * (BT / n) tokens. + expert_capacity = min(BT, max(1, int(math.ceil((BT * self.capacity_factor) / self.n_experts)))) + + # 1. Compute affinity scores + clean_logits = self.gate(x_flat) # (BT, n_experts) + if self.training: + noise_std = F.softplus(self.noise_gate(x_flat)) + noise = torch.randn_like(clean_logits) * noise_std + affinity_logits = clean_logits + noise + else: + affinity_logits = clean_logits + + # In Expert Choice, normalization is over experts first to get routing probabilities + # then experts pick top tokens based on those probabilities. + # This implies experts compete for tokens based on how much the token "wants" that expert. + token_to_expert_probs = torch.softmax(affinity_logits, dim=-1) # (BT, n_experts) + + # Transpose to (n_experts, BT) so each expert can pick top tokens + expert_to_token_probs = token_to_expert_probs.t() # (n_experts, BT) + + expert_outputs = torch.zeros((BT, base_logits.shape[-1]), device=x.device, dtype=x.dtype) + + # We need a scaling factor later. If an expert picks a token, its output is scaled by + # the token_to_expert_prob. + + # 2. Each expert processes its top tokens + for i in range(self.n_experts): + probs_i = expert_to_token_probs[i] # (BT,) + + # Select top-k tokens for this expert + top_probs, top_token_idx = torch.topk(probs_i, k=expert_capacity, dim=0) # (capacity,) + + # Gather tokens + selected_tokens = x_flat[top_token_idx] # (capacity, D) + + # Process through expert i + act = self.activations[i] + out = self.experts_down[i](self.dropout(act(self.experts_up[i](selected_tokens)))) + out = self.expert_norms[i](out) # (capacity, out_dim) + + # Scale by routing probability + out = out * top_probs.unsqueeze(-1) # (capacity, out_dim) + + # Add to the global output buffer using `scatter_add_` or index_add + # out is (capacity, out_dim). top_token_idx is (capacity,). + # expert_outputs is (BT, out_dim). + expert_outputs.index_add_(0, top_token_idx, out) + + expert_outputs = expert_outputs.view(*shape_prefix, -1) + calib = self.calibration(x) + + return ( + base_logits + + self.alpha * expert_outputs + + self.theta * calib + ) + + +class SmarterExpertClassifierHead(nn.Module): + """ + Next-generation MoE classifier head featuring: + - 1 always-active shared anchor expert. + - 8 diversity-driven routed experts. + - Sigma Gating (independent sigmoid scores per expert). + - Aux-loss-free load balancing (DeepSeek-V3 style dynamic bias). + - Lightweight residual LoRA adapters per expert. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 8, + lora_rank: int = 32, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_experts = n_experts + + # 1. Shared Anchor Expert + shared_dim = 2048 + self.shared_up = nn.Linear(in_dim, shared_dim, bias=False) + self.shared_down = nn.Linear(shared_dim, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Routed Experts with diversity + expansion_dims = [1024, 1536, 2048, 2560, 1024, 1536, 2048, 3072] + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, expansion_dims[i % 8], bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(expansion_dims[i % 8], out_dim, bias=False) for i in range(n_experts) + ]) + # Mix of activations + base_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh, F.softplus, F.elu] + self.activations = [base_acts[i % len(base_acts)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # 3. Residual LoRA Adapters for experts + self.lora_down = nn.ModuleList([nn.Linear(in_dim, lora_rank, bias=False) for _ in range(n_experts)]) + self.lora_up = nn.ModuleList([nn.Linear(lora_rank, out_dim, bias=False) for _ in range(n_experts)]) + + # 4. Routing Mechanism (Sigma Gating) + self.gate = nn.Linear(in_dim, n_experts, bias=False) + self.register_buffer("expert_bias", torch.zeros(n_experts)) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.zeros_(self.shared_down.weight) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) + nn.init.zeros_(self.experts_down[i].weight) + nn.init.normal_(self.lora_down[i].weight, std=0.02) + nn.init.zeros_(self.lora_up[i].weight) + + nn.init.normal_(self.gate.weight, std=0.02) + nn.init.zeros_(self.calibration.weight) + nn.init.zeros_(self.calibration.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + BT = x.reshape(-1, x.shape[-1]).shape[0] + x_flat = x.reshape(-1, x.shape[-1]) + + # Shared Anchor processing + shared_out = self.shared_norm(self.shared_down(self.dropout(F.silu(self.shared_up(x_flat))))) + + # Sigma Gating + clean_logits = self.gate(x_flat) + if self.training: + gate_logits = clean_logits + self.expert_bias + else: + gate_logits = clean_logits + + # Sigma activation + gate_scores = torch.sigmoid(gate_logits) + + all_expert_logits = [] + for i in range(self.n_experts): + core = self.experts_down[i](self.dropout(self.activations[i](self.experts_up[i](x_flat)))) + lora = self.lora_up[i](self.lora_down[i](x_flat)) + out = self.expert_norms[i](core + lora) + all_expert_logits.append(out) + + expert_stack = torch.stack(all_expert_logits, dim=1) + routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + if self.training: + with torch.no_grad(): + target_load = 1.0 / self.n_experts + actual_load = gate_scores.mean(dim=0) + self.expert_bias.add_(0.01 * (target_load - actual_load)) + + routed_out = routed_out.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + calib = self.calibration(x) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * routed_out + + self.theta * calib + ) + + +class CrossAttentionFusion(nn.Module): + """ + Multi-head cross-attention fusion for expert representations. + Allows different expert branches to attend to each other's outputs. + """ + def __init__(self, out_dim: int, n_branches: int, n_heads: int = 2): + super().__init__() + self.out_dim = out_dim + self.n_branches = n_branches + self.n_heads = n_heads + self.head_dim = out_dim // n_heads + + self.q_proj = nn.Linear(out_dim, out_dim, bias=False) + self.k_proj = nn.Linear(out_dim, out_dim, bias=False) + self.v_proj = nn.Linear(out_dim, out_dim, bias=False) + self.out_proj = nn.Linear(out_dim, out_dim, bias=False) + + self.fusion_weight = nn.Linear(out_dim, 1) + + def forward(self, x): + # x: (B, T, N, D) + B, T, N, D = x.shape + x_flat = x.view(B * T, N, D) + + q = self.q_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) + + attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) + attn = F.softmax(attn, dim=-1) + + out = (attn @ v).transpose(1, 2).reshape(B * T, N, D) + out = self.out_proj(out) + + # Weighted reduction to single D-dim vector per token + scores = torch.sigmoid(self.fusion_weight(out)) # (B*T, N, 1) + fused = (out * scores).sum(dim=1) / (scores.sum(dim=1) + 1e-6) + + return fused.view(B, T, D) + + +class ReasoningCell(nn.Module): + """ + Sub-module for iterative reasoning steps. + Combines a residual MLP with LayerNorm and dropout. + """ + def __init__(self, dim: int = 256, inner_dim: int = 512, dropout: float = 0.1): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.up = nn.Linear(dim, inner_dim) + self.down = nn.Linear(inner_dim, dim) + self.dropout = nn.Dropout(dropout) + self.act = F.silu + + def forward(self, x): + h = self.norm(x) + h = self.down(self.dropout(self.act(self.up(h)))) + return x + h + + +class ThoughtExpertClassifierHead(nn.Module): + """ + Ultra-advanced MoE classifier head featuring iterative reasoning: + - 1 always-active shared anchor expert. + - 8 diversity-driven routed experts. + - 3-step reasoning loop that refines the feature representation. + - Sigma Gating (independent sigmoid scores) at each step. + - Cross-expert attention fusion integrated into the "thinking" process. + - Lightweight residual LoRA adapters per expert. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 8, + reasoning_steps: int = 3, + lora_rank: int = 32, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_experts = n_experts + self.reasoning_steps = reasoning_steps + + # 1. Shared Anchor Expert + shared_dim = 2048 + self.shared_up = nn.Linear(in_dim, shared_dim, bias=False) + self.shared_down = nn.Linear(shared_dim, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Routed Experts with diversity + expansion_dims = [1024, 1536, 2048, 2560, 1024, 1536, 2048, 3072] + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, expansion_dims[i % 8], bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(expansion_dims[i % 8], out_dim, bias=False) for i in range(n_experts) + ]) + base_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh, F.softplus, F.elu] + self.activations = [base_acts[i % len(base_acts)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # 3. Residual LoRA Adapters for experts + self.lora_down = nn.ModuleList([nn.Linear(in_dim, lora_rank, bias=False) for _ in range(n_experts)]) + self.lora_up = nn.ModuleList([nn.Linear(lora_rank, out_dim, bias=False) for _ in range(n_experts)]) + + # 4. Reasoning Transformation + self.reasoning_cells = nn.ModuleList([ + ReasoningCell(in_dim, inner_dim=512, dropout=dropout) for _ in range(reasoning_steps) + ]) + + # 5. Iterative Routing (Sigma Gating) + self.gates = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps) + ]) + self.register_buffer("expert_bias", torch.zeros(n_experts)) + + # 6. Cross-Expert Attention Fusion + self.cross_attn = CrossAttentionFusion(out_dim=out_dim, n_branches=n_experts, n_heads=2) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.zeros_(self.shared_down.weight) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) + nn.init.zeros_(self.experts_down[i].weight) + nn.init.normal_(self.lora_down[i].weight, std=0.02) + nn.init.zeros_(self.lora_up[i].weight) + + for g in self.gates: + nn.init.normal_(g.weight, std=0.02) + + nn.init.zeros_(self.calibration.weight) + nn.init.zeros_(self.calibration.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + + # Shared Anchor processing + shared_out = self.shared_norm(self.shared_down(self.dropout(F.silu(self.shared_up(x_flat))))) + + # Pre-compute expert candidates (B*T, n_experts, out_dim) + all_expert_logits = [] + for i in range(self.n_experts): + core = self.experts_down[i](self.dropout(self.activations[i](self.experts_up[i](x_flat)))) + lora = self.lora_up[i](self.lora_down[i](x_flat)) + out = self.expert_norms[i](core + lora) + all_expert_logits.append(out) + expert_stack = torch.stack(all_expert_logits, dim=1) + + # Iterative Reasoning Loop + current_features = x_flat + total_routed_out = 0 + + for step in range(self.reasoning_steps): + # 1. Refine features + current_features = self.reasoning_cells[step](current_features) + + # 2. Gate experts based on refined features + gate_logits = self.gates[step](current_features) + if self.training: + gate_logits = gate_logits + self.expert_bias + gate_scores = torch.sigmoid(gate_logits) # (B*T, n_experts) + + # 3. Fuse experts for this step + step_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + # 4. Attention-based cross-expert correction (residual) + # Re-weight experts with attention to find hidden correlations + # Shape: expert_stack is (BT, N, D), we need (BT, 1, N, D) for CrossAttentionFusion + attn_fused = self.cross_attn(expert_stack.unsqueeze(1)).squeeze(1) + + total_routed_out = total_routed_out + step_out + 0.1 * attn_fused + + if self.training: + with torch.no_grad(): + target_load = 1.0 / self.n_experts + actual_load = gate_scores.mean(dim=0) + self.expert_bias.add_(0.01 * (target_load - actual_load)) + + routed_out = total_routed_out / self.reasoning_steps + routed_out = routed_out.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + calib = self.calibration(x) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * routed_out + + self.theta * calib + ) + + +class RecursiveThoughtExpertHead(nn.Module): + """ + Generation v14: RecursiveThoughtExpert (Recursive reasoning + ACE + Multi-Head Routing) + Features: + - Adaptive Reasoning Depth (ACE): Early exit logic based on confidence. + - Multi-Head Sigma Gating (2 heads): Richer routing subspace. + - Hierarchical Shared Experts: Global (2048) + Local (512) experts. + - Cross-Expert Attention Fusion (Residual). + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 8, + reasoning_steps: int = 3, + n_heads: int = 2, + lora_rank: int = 32, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_experts = n_experts + self.reasoning_steps = reasoning_steps + self.n_heads = n_heads + + # 1. Hierarchical Shared Experts + # Global + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # Local (New in v14) + self.local_up = nn.Linear(in_dim, 512, bias=False) + self.local_down = nn.Linear(512, out_dim, bias=False) + self.local_norm = nn.LayerNorm(out_dim) + self.local_scale = nn.Parameter(torch.tensor(0.5)) + + # 2. Routed Experts with diversity + expansion_dims = [1024, 1536, 2048, 2560, 1024, 1536, 2048, 3072] + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, expansion_dims[i % 8], bias=False) for i in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(expansion_dims[i % 8], out_dim, bias=False) for i in range(n_experts) + ]) + base_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh, F.softplus, F.elu] + self.activations = [base_acts[i % len(base_acts)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # 3. Residual LoRA Adapters + self.lora_down = nn.ModuleList([nn.Linear(in_dim, lora_rank, bias=False) for _ in range(n_experts)]) + self.lora_up = nn.ModuleList([nn.Linear(lora_rank, out_dim, bias=False) for _ in range(n_experts)]) + + # 4. Reasoning & ACE + self.reasoning_cells = nn.ModuleList([ + ReasoningCell(in_dim, inner_dim=512, dropout=dropout) for _ in range(reasoning_steps) + ]) + self.exit_gates = nn.ModuleList([ + nn.Linear(in_dim, 1) for _ in range(reasoning_steps) + ]) + + # 5. Multi-Head Sigma Gating (Flattened for registration) + self.gates = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps * n_heads) + ]) + self.register_buffer("expert_bias", torch.zeros(n_experts)) + + # 6. Cross-Expert Attention Fusion + self.cross_attn = CrossAttentionFusion(out_dim=out_dim, n_branches=n_experts, n_heads=2) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.calibration = nn.Linear(in_dim, out_dim, bias=True) + self.theta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.zeros_(self.shared_down.weight) + nn.init.kaiming_uniform_(self.local_up.weight, a=math.sqrt(5)) + nn.init.zeros_(self.local_down.weight) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.experts_down[i].weight, std=0.01) # Small initial signal + nn.init.normal_(self.lora_down[i].weight, std=0.02) + nn.init.zeros_(self.lora_up[i].weight) + + for g in self.gates: + nn.init.normal_(g.weight, std=0.02) + + for e in self.exit_gates: + nn.init.constant_(e.bias, -3.0) # Start with low exit prob + + nn.init.zeros_(self.calibration.weight) + nn.init.zeros_(self.calibration.bias) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + + # 1. Global Shared Expert + shared_out = self.shared_norm(self.shared_down(self.dropout(F.silu(self.shared_up(x_flat))))) + # 2. Local Shared Expert + local_out = self.local_norm(self.local_down(self.dropout(F.gelu(self.local_up(x_flat))))) + + # Pre-compute expert candidates + all_expert_logits = [] + for i in range(self.n_experts): + core = self.experts_down[i](self.dropout(self.activations[i](self.experts_up[i](x_flat)))) + lora = self.lora_up[i](self.lora_down[i](x_flat)) + out = self.expert_norms[i](core + lora) + all_expert_logits.append(out) + expert_stack = torch.stack(all_expert_logits, dim=1) + + # Iterative Reasoning Loop with ACE + current_features = x_flat + total_routed_out = torch.zeros_like(expert_stack[:, 0, :]) + cumulative_exit_prob = torch.zeros(x_flat.shape[0], 1, device=x.device) + depth_count = 0 + + for step in range(self.reasoning_steps): + depth_count += 1 + # A. Refine features + current_features = self.reasoning_cells[step](current_features) + + # B. Multi-Head Sigma Gating + head_logits = [] + for h in range(self.n_heads): + idx = step * self.n_heads + h + head_logits.append(self.gates[idx](current_features)) + # Average head logits (Sigma Gating) + gate_logits = torch.stack(head_logits, dim=0).mean(dim=0) + + if self.training: + gate_logits = gate_logits + self.expert_bias + gate_scores = torch.sigmoid(gate_logits) # (B*T, n_experts) + + # C. ACE: Early Exit check + exit_logit = self.exit_gates[step](current_features) + exit_prob = torch.sigmoid(exit_logit) # (B*T, 1) + + # D. Fuse experts for this step + step_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + # E. Attention-based cross-expert correction + attn_fused = self.cross_attn(expert_stack.unsqueeze(1)).squeeze(1) + + # Weighted contribution based on current depth + total_routed_out = total_routed_out + (1.0 - cumulative_exit_prob) * (step_out + 0.1 * attn_fused) + + # Update cumulative exit probability + cumulative_exit_prob = cumulative_exit_prob + (1.0 - cumulative_exit_prob) * exit_prob + + if self.training: + with torch.no_grad(): + target_load = 1.0 / self.n_experts + actual_load = gate_scores.mean(dim=0) + self.expert_bias.add_(0.01 * (target_load - actual_load)) + + # Early exit condition + if not self.training and cumulative_exit_prob.mean() > 0.9: + break + + routed_out = total_routed_out / depth_count + routed_out = routed_out.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + local_out = local_out.view(*shape_prefix, -1) + calib = self.calibration(x) + + return ( + base_logits + + self.shared_scale * shared_out + + self.local_scale * local_out + + (self.alpha + 1e-4) * routed_out + + self.theta * calib + ) + + +class ChampionNetRecursiveExpert(nn.Module): + """ + Backbone-compatible model with RecursiveThoughtExpertHead. + """ + def __init__(self, n_experts: int = 8, reasoning_steps: int = 3, lora_rank: int = 32, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(RecursiveThoughtExpertHead( + 256, 10, + n_experts=n_experts, + reasoning_steps=reasoning_steps, + lora_rank=lora_rank, + dropout=dropout + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + + + +class ReflexiveThoughtExpertHead(nn.Module): + """ + Generation v15: ReflexiveThoughtExpert (Self-Reflection MoE) + Features: + - Initial Pass: Generates baseline logits using MoE routing. + - Critique Pass: Takes original features + initial logits to self-correct using specialized critique experts. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 4, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_experts = n_experts + + # 1. Initial Pass Experts + self.initial_up = nn.ModuleList([nn.Linear(in_dim, 1024, bias=False) for _ in range(n_experts)]) + self.initial_down = nn.ModuleList([nn.Linear(1024, out_dim, bias=False) for _ in range(n_experts)]) + self.initial_gate = nn.Linear(in_dim, n_experts, bias=False) + self.initial_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # 2. Critique Pass Experts (Input: features + intermediate logits) + critique_dim = in_dim + out_dim + self.critique_up = nn.ModuleList([nn.Linear(critique_dim, 1024, bias=False) for _ in range(n_experts)]) + self.critique_down = nn.ModuleList([nn.Linear(1024, out_dim, bias=False) for _ in range(n_experts)]) + self.critique_gate = nn.Linear(critique_dim, n_experts, bias=False) + self.critique_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.beta = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.zeros_(self.initial_gate.weight) + nn.init.zeros_(self.critique_gate.weight) + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.initial_up[i].weight, a=5**0.5) + nn.init.normal_(self.initial_down[i].weight, std=0.01) + nn.init.kaiming_uniform_(self.critique_up[i].weight, a=5**0.5) + nn.init.normal_(self.critique_down[i].weight, std=0.01) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + # --- Initial Pass --- + gate_init = torch.softmax(self.initial_gate(x), dim=-1) # (B, T, E) + init_outs = [] + for i in range(self.n_experts): + h = self.dropout(F.silu(self.initial_up[i](x))) + init_outs.append(self.initial_norms[i](self.initial_down[i](h))) + init_stack = torch.stack(init_outs, dim=-2) # (B, T, E, D) + init_fused = (init_stack * gate_init.unsqueeze(-1)).sum(dim=-2) + + # Intermediate prediction + inter_logits = base_logits + self.alpha * init_fused + + # --- Critique Pass --- + crit_in = torch.cat([x, inter_logits], dim=-1) + gate_crit = torch.softmax(self.critique_gate(crit_in), dim=-1) + crit_outs = [] + for i in range(self.n_experts): + h = self.dropout(F.gelu(self.critique_up[i](crit_in))) + crit_outs.append(self.critique_norms[i](self.critique_down[i](h))) + crit_stack = torch.stack(crit_outs, dim=-2) + crit_fused = (crit_stack * gate_crit.unsqueeze(-1)).sum(dim=-2) + + # Final prediction + return inter_logits + self.beta * crit_fused + + +class ChampionNetReflexiveExpert(nn.Module): + """ + Backbone-compatible model with ReflexiveThoughtExpertHead. + """ + def __init__(self, n_experts: int = 4, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ReflexiveThoughtExpertHead(256, 10, n_experts=n_experts, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class MetaCognitiveExpertHead(nn.Module): + """ + Generation v16: MetaCognitiveExpert (Iterative Reflection + PonderNet Halting) + Combines the best ideas from v14 (recursive reasoning, shared experts, ACE) + and v15 (self-reflection critique pass) into a unified "think → critique → refine" + loop with learned adaptive halting. + + Features: + - Shared Expert: Always-on global expert for stable signal. + - ReasoningCells: Per-step feature refinement with residual MLP. + - Proposal MoE: Generates predictions via routed experts at each step. + - Critique MoE: Inspects [features || proposal_logits] to output corrections. + - PonderNet Halting: Learned halt probability per step for adaptive depth. + - Dynamic Bias Load Balancing: Buffer biases maintain expert utilization. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_proposal_experts: int = 6, + n_critique_experts: int = 4, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + self.n_proposal_experts = n_proposal_experts + self.n_critique_experts = n_critique_experts + self.reasoning_steps = reasoning_steps + + # 1. Shared Expert (always-on) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Reasoning Cells (per step) + self.reasoning_cells = nn.ModuleList([ + ReasoningCell(in_dim, inner_dim=512, dropout=dropout) + for _ in range(reasoning_steps) + ]) + + # 3. Proposal MoE (per step routing) + prop_dims = [1024, 1536, 2048, 2560, 1024, 1536] + prop_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.proposal_up = nn.ModuleList([ + nn.Linear(in_dim, prop_dims[i % len(prop_dims)], bias=False) + for i in range(n_proposal_experts) + ]) + self.proposal_down = nn.ModuleList([ + nn.Linear(prop_dims[i % len(prop_dims)], out_dim, bias=False) + for i in range(n_proposal_experts) + ]) + self.proposal_acts = [prop_acts[i % len(prop_acts)] for i in range(n_proposal_experts)] + self.proposal_norms = nn.ModuleList([ + nn.LayerNorm(out_dim) for _ in range(n_proposal_experts) + ]) + self.proposal_gates = nn.ModuleList([ + nn.Linear(in_dim, n_proposal_experts, bias=False) + for _ in range(reasoning_steps) + ]) + self.register_buffer("proposal_bias", torch.zeros(n_proposal_experts)) + + # 4. Critique MoE (input: features + proposal logits) + critique_dim = in_dim + out_dim + self.critique_up = nn.ModuleList([ + nn.Linear(critique_dim, 1024, bias=False) + for _ in range(n_critique_experts) + ]) + self.critique_down = nn.ModuleList([ + nn.Linear(1024, out_dim, bias=False) + for _ in range(n_critique_experts) + ]) + self.critique_norms = nn.ModuleList([ + nn.LayerNorm(out_dim) for _ in range(n_critique_experts) + ]) + self.critique_gates = nn.ModuleList([ + nn.Linear(critique_dim, n_critique_experts, bias=False) + for _ in range(reasoning_steps) + ]) + + # 5. Halting Gates (PonderNet-style) + self.halt_gates = nn.ModuleList([ + nn.Linear(in_dim, 1) for _ in range(reasoning_steps) + ]) + + # Aggregation + self.alpha = nn.Parameter(torch.tensor(0.0)) # proposal scale + self.beta = nn.Parameter(torch.tensor(0.0)) # critique scale + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.zeros_(self.shared_down.weight) + + for i in range(self.n_proposal_experts): + nn.init.kaiming_uniform_(self.proposal_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.proposal_down[i].weight, std=0.01) + for i in range(self.n_critique_experts): + nn.init.kaiming_uniform_(self.critique_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.critique_down[i].weight, std=0.01) + + for g in self.proposal_gates: + nn.init.normal_(g.weight, std=0.02) + for g in self.critique_gates: + nn.init.normal_(g.weight, std=0.02) + for h in self.halt_gates: + nn.init.constant_(h.bias, -3.0) # start with low halt probability + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + + # 1. Shared Expert (always-on) + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # 2. Pre-compute proposal expert outputs + prop_outs = [] + for i in range(self.n_proposal_experts): + h = self.dropout(self.proposal_acts[i](self.proposal_up[i](x_flat))) + prop_outs.append(self.proposal_norms[i](self.proposal_down[i](h))) + prop_stack = torch.stack(prop_outs, dim=1) # (BT, E_prop, D) + + # 3. Iterative Reflect Loop + current_features = x_flat + total_proposal = torch.zeros_like(base_logits.reshape(-1, base_logits.shape[-1])) + total_critique = torch.zeros_like(total_proposal) + cumulative_halt = torch.zeros(x_flat.shape[0], 1, device=x.device) + depth_count = 0 + + for step in range(self.reasoning_steps): + depth_count += 1 + + # A. Refine features + current_features = self.reasoning_cells[step](current_features) + + # B. Proposal MoE routing (Sigma Gating) + gate_logits = self.proposal_gates[step](current_features) + if self.training: + gate_logits = gate_logits + self.proposal_bias + gate_scores = torch.sigmoid(gate_logits) # (BT, E_prop) + step_proposal = (prop_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + # C. Critique MoE: inspect [features || proposal] + crit_in = torch.cat([current_features, step_proposal], dim=-1) + crit_gate = torch.softmax( + self.critique_gates[step](crit_in), dim=-1 + ) # (BT, E_crit) + crit_outs = [] + for i in range(self.n_critique_experts): + h = self.dropout(F.gelu(self.critique_up[i](crit_in))) + crit_outs.append(self.critique_norms[i](self.critique_down[i](h))) + crit_stack = torch.stack(crit_outs, dim=1) # (BT, E_crit, D) + step_critique = (crit_stack * crit_gate.unsqueeze(-1)).sum(dim=1) + + # D. Halting gate (PonderNet-style) + halt_prob = torch.sigmoid(self.halt_gates[step](current_features)) + + # E. Accumulate weighted by remaining probability mass + remaining = 1.0 - cumulative_halt + total_proposal = total_proposal + remaining * step_proposal + total_critique = total_critique + remaining * step_critique + cumulative_halt = cumulative_halt + remaining * halt_prob + + # F. Dynamic bias load balancing + if self.training: + with torch.no_grad(): + target_load = 1.0 / self.n_proposal_experts + actual_load = gate_scores.mean(dim=0) + self.proposal_bias.add_(0.01 * (target_load - actual_load)) + + # G. Early exit during inference + if not self.training and cumulative_halt.mean() > 0.9: + break + + # Normalize by effective depth + total_proposal = total_proposal / max(1, depth_count) + total_critique = total_critique / max(1, depth_count) + + # Reshape back + total_proposal = total_proposal.view(*shape_prefix, -1) + total_critique = total_critique.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * total_proposal + + (self.beta + 1e-4) * total_critique + ) + + +class ChampionNetMetaCognitiveExpert(nn.Module): + """ + Backbone-compatible model with MetaCognitiveExpertHead (v16). + """ + def __init__( + self, + n_proposal_experts: int = 6, + n_critique_experts: int = 4, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(MetaCognitiveExpertHead( + 256, 10, + n_proposal_experts=n_proposal_experts, + n_critique_experts=n_critique_experts, + reasoning_steps=reasoning_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TreeOfThoughtExpertHead(nn.Module): + """ + Generation v17: Tree-of-Thought Expert (Latent Beam Search) + Embeds true test-time compute scaling into the forward pass. + Instead of a single trajectory, maintains a beam of k states. + At each step, m Action Experts propose modifications, creating k*m candidates. + A Value Network scores them, and the top-k are kept. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_action_experts: int = 6, + beam_size: int = 4, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_action_experts = n_action_experts + self.beam_size = beam_size + self.reasoning_steps = reasoning_steps + + # Base projection + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # 1. Global Shared Expert (always-on) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Value Network (Scorer) + # Evaluates the "goodness" of a latent state. + self.value_net = nn.Sequential( + nn.Linear(in_dim, 512), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(512, 1) + ) + + # 3. Action MoE (Proposer) + # We use a set of experts that propose feature modifications + self.action_up = nn.ModuleList([ + nn.Linear(in_dim, 1024, bias=False) for _ in range(n_action_experts) + ]) + self.action_down = nn.ModuleList([ + nn.Linear(1024, in_dim, bias=False) for _ in range(n_action_experts) + ]) + self.action_norms = nn.ModuleList([ + nn.LayerNorm(in_dim) for _ in range(n_action_experts) + ]) + + # Final projection from feature to out_dim + self.final_up = nn.Linear(in_dim, 1024, bias=False) + self.final_down = nn.Linear(1024, out_dim, bias=False) + self.final_norm = nn.LayerNorm(out_dim) + + self.alpha = nn.Parameter(torch.tensor(0.0)) # tree scale + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + for p in self.value_net.parameters(): + if p.dim() > 1: + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + else: + nn.init.zeros_(p) + + for i in range(self.n_action_experts): + nn.init.kaiming_uniform_(self.action_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.action_down[i].weight, std=0.01) + + nn.init.kaiming_uniform_(self.final_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.final_down.weight, std=0.01) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # 1. Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # 2. Tree-of-Thought Search + # Initial beam state: K=1 + beam_states = x_flat.unsqueeze(1) # (N, 1, D) + + for step in range(self.reasoning_steps): + K = beam_states.shape[1] + + # A. Expansion: Each action expert proposes a modification to every state in the beam + nk_states = beam_states.reshape(N * K, self.in_dim) + + candidates = [] + for i in range(self.n_action_experts): + h = self.dropout(F.gelu(self.action_up[i](nk_states))) + delta = self.action_norms[i](self.action_down[i](h)) + new_state = nk_states + delta + candidates.append(new_state) + + # Stack candidates + cand_stack = torch.stack(candidates, dim=1) + cand_stack = cand_stack.reshape(N, K * self.n_action_experts, self.in_dim) + + # B. Evaluation (Value Network) + flat_cands = cand_stack.reshape(N * K * self.n_action_experts, self.in_dim) + scores = self.value_net(flat_cands).squeeze(-1) # (N * K * M) + scores = scores.reshape(N, K * self.n_action_experts) # (N, K*M) + + # C. Selection (Beam Search) + K_next = min(self.beam_size, K * self.n_action_experts) + + topk_scores, topk_indices = torch.topk(scores, k=K_next, dim=1) # (N, K_next) + + expanded_indices = topk_indices.unsqueeze(-1).expand(-1, -1, self.in_dim) + beam_states = torch.gather(cand_stack, 1, expanded_indices) # (N, K_next, D) + + # 3. Aggregation across the final beam + final_K = beam_states.shape[1] + flat_final = beam_states.reshape(N * final_K, self.in_dim) + + h_final = self.dropout(F.silu(self.final_up(flat_final))) + final_preds = self.final_norm(self.final_down(h_final)) # (N * final_K, out_dim) + final_preds = final_preds.reshape(N, final_K, self.out_dim) # (N, K, out_dim) + + final_scores = self.value_net(flat_final).squeeze(-1).reshape(N, final_K) # (N, K) + agg_weights = torch.softmax(final_scores, dim=1) # (N, K) + + tree_fused = (final_preds * agg_weights.unsqueeze(-1)).sum(dim=1) + + base_logits = base_logits.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + tree_fused = tree_fused.view(*shape_prefix, -1) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * tree_fused + ) + + +class ChampionNetTreeOfThoughtExpert(nn.Module): + """ + Backbone-compatible model with TreeOfThoughtExpertHead (v17). + """ + def __init__( + self, + n_action_experts: int = 6, + beam_size: int = 4, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(TreeOfThoughtExpertHead( + 256, 10, + n_action_experts=n_action_experts, + beam_size=beam_size, + reasoning_steps=reasoning_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ConsensusExpertHead(nn.Module): + """ + Generation v18: Consensus Expert (Architectural Diversity + Learned Arbitration) + Uses three fundamentally different computational paradigms as independent + reasoning pathways, then a Consensus Network arbitrates. + + Pathways: + 1. MLP Pathway — Deep feed-forward MoE experts (fast pattern matching) + 2. Attention Pathway — Self-attention over learned expert embeddings (relational reasoning) + 3. Convolutional Pathway — 1D convolution over the feature vector (local pattern detection) + + A Consensus Network evaluates inter-pathway agreement and produces the final answer. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_mlp_experts: int = 6, + n_attn_heads: int = 4, + conv_channels: int = 64, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_mlp_experts = n_mlp_experts + + # Base projection + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # --- Pathway 1: MLP MoE --- + mlp_dims = [1024, 1536, 2048, 1024, 1536, 2048] + mlp_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.mlp_up = nn.ModuleList([ + nn.Linear(in_dim, mlp_dims[i % len(mlp_dims)], bias=False) + for i in range(n_mlp_experts) + ]) + self.mlp_down = nn.ModuleList([ + nn.Linear(mlp_dims[i % len(mlp_dims)], out_dim, bias=False) + for i in range(n_mlp_experts) + ]) + self.mlp_acts = [mlp_acts[i % len(mlp_acts)] for i in range(n_mlp_experts)] + self.mlp_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_mlp_experts)]) + self.mlp_gate = nn.Linear(in_dim, n_mlp_experts, bias=False) + + # --- Pathway 2: Self-Attention --- + self.n_attn_heads = n_attn_heads + self.attn_embed = nn.Parameter(torch.randn(8, in_dim) * 0.02) + self.attn_q = nn.Linear(in_dim, in_dim, bias=False) + self.attn_k = nn.Linear(in_dim, in_dim, bias=False) + self.attn_v = nn.Linear(in_dim, in_dim, bias=False) + self.attn_proj = nn.Linear(in_dim, out_dim, bias=False) + self.attn_norm = nn.LayerNorm(out_dim) + self.attn_scale = (in_dim // n_attn_heads) ** -0.5 + + # --- Pathway 3: 1D Convolution --- + self.conv_channels = conv_channels + self.conv1 = nn.Conv1d(1, conv_channels, kernel_size=5, padding=2) + self.conv2 = nn.Conv1d(conv_channels, conv_channels, kernel_size=3, padding=1) + self.conv_pool = nn.AdaptiveAvgPool1d(1) + self.conv_proj = nn.Linear(conv_channels, out_dim, bias=False) + self.conv_norm = nn.LayerNorm(out_dim) + + # --- Consensus Network (Arbiter) --- + consensus_in = 3 * out_dim + self.consensus = nn.Sequential( + nn.Linear(consensus_in, 128), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(128, 64), + nn.GELU(), + nn.Linear(64, 3), + ) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + for i in range(self.n_mlp_experts): + nn.init.kaiming_uniform_(self.mlp_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.mlp_down[i].weight, std=0.01) + nn.init.normal_(self.mlp_gate.weight, std=0.02) + + nn.init.kaiming_uniform_(self.attn_q.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.attn_k.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.attn_v.weight, a=math.sqrt(5)) + nn.init.normal_(self.attn_proj.weight, std=0.01) + + nn.init.normal_(self.conv_proj.weight, std=0.01) + + for p in self.consensus.parameters(): + if p.dim() > 1: + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + else: + nn.init.zeros_(p) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # --- Pathway 1: MLP MoE --- + gate_scores = torch.sigmoid(self.mlp_gate(x_flat)) + mlp_outs = [] + for i in range(self.n_mlp_experts): + h = self.dropout(self.mlp_acts[i](self.mlp_up[i](x_flat))) + mlp_outs.append(self.mlp_norms[i](self.mlp_down[i](h))) + mlp_stack = torch.stack(mlp_outs, dim=1) + mlp_pred = (mlp_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + # --- Pathway 2: Self-Attention --- + embeds = self.attn_embed.unsqueeze(0).expand(N, -1, -1) + q = self.attn_q(x_flat).unsqueeze(1) + k = self.attn_k(embeds) + v = self.attn_v(embeds) + attn_weights = torch.softmax( + (q * k).sum(dim=-1, keepdim=True) * self.attn_scale, dim=1 + ) + attn_out = (attn_weights * v).sum(dim=1) + attn_pred = self.attn_norm(self.attn_proj(attn_out)) + + # --- Pathway 3: 1D Convolution --- + x_conv = x_flat.unsqueeze(1) + h_conv = F.gelu(self.conv1(x_conv)) + h_conv = F.gelu(self.conv2(h_conv)) + h_conv = self.conv_pool(h_conv).squeeze(-1) + conv_pred = self.conv_norm(self.conv_proj(h_conv)) + + # --- Consensus Network --- + pathway_preds = torch.stack([mlp_pred, attn_pred, conv_pred], dim=1) + consensus_in = torch.cat([mlp_pred, attn_pred, conv_pred], dim=-1) + consensus_weights = torch.softmax(self.consensus(consensus_in), dim=-1) + consensus_fused = (pathway_preds * consensus_weights.unsqueeze(-1)).sum(dim=1) + + consensus_fused = consensus_fused.view(*shape_prefix, -1) + base_logits = base_logits.view(*shape_prefix, -1) + + return base_logits + (self.alpha + 1e-4) * consensus_fused + + +class ChampionNetConsensusExpert(nn.Module): + """ + Backbone-compatible model with ConsensusExpertHead (v18). + """ + def __init__( + self, + n_mlp_experts: int = 6, + n_attn_heads: int = 4, + conv_channels: int = 64, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ConsensusExpertHead( + 256, 10, + n_mlp_experts=n_mlp_experts, + n_attn_heads=n_attn_heads, + conv_channels=conv_channels, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class DeliberativeAlignmentExpertHead(nn.Module): + """ + Generation v19: Deliberative Alignment Expert + (Multi-Draft Deliberation with Working Memory) + + The most advanced reasoning head. Simultaneously generates K independent + 'draft' reasoning chains, lets them share insights via cross-draft + attention, consults a working memory bank, and uses PonderNet-style + adaptive halting to allocate more compute to hard inputs. + + Key innovations: + 1. Working Memory Bank — learned key-value slots for persistent scratchpad + 2. Multi-Draft Generation — K parallel reasoning perspectives + 3. Cross-Draft Attention — drafts attend to each others states + 4. Iterative Refinement + Adaptive Depth — PonderNet halting + 5. Consistency-Weighted Aggregation — prefer internally consistent conclusions + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 6, + n_drafts: int = 3, + n_mem_slots: int = 16, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + self.n_drafts = n_drafts + self.n_mem_slots = n_mem_slots + self.reasoning_steps = reasoning_steps + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # 1. Global Shared Expert (always-on) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Working Memory Bank + self.memory_keys = nn.Parameter(torch.randn(n_mem_slots, in_dim) * 0.02) + self.memory_values = nn.Parameter(torch.randn(n_mem_slots, in_dim) * 0.02) + self.mem_read_query = nn.Linear(in_dim, in_dim, bias=False) + self.mem_write_query = nn.Linear(in_dim, in_dim, bias=False) + self.mem_write_gate = nn.Linear(in_dim, n_mem_slots, bias=True) + self.mem_scale = (in_dim) ** -0.5 + + # 3. Per-Draft Reasoning Cells (iterative refinement) + self.draft_reasoning_cells = nn.ModuleList() + for _ in range(reasoning_steps): + step_cells = nn.ModuleList([ + nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, 512), + nn.SiLU(), + nn.Linear(512, in_dim), + ) + for _ in range(n_drafts) + ]) + self.draft_reasoning_cells.append(step_cells) + + # 4. Expert MoE (shared across drafts) + expert_dims = [1024, 1536, 2048, 1024, 1536, 2048] + expert_acts_fns = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.expert_up = nn.ModuleList([ + nn.Linear(in_dim, expert_dims[i % len(expert_dims)], bias=False) + for i in range(n_experts) + ]) + self.expert_down = nn.ModuleList([ + nn.Linear(expert_dims[i % len(expert_dims)], out_dim, bias=False) + for i in range(n_experts) + ]) + self.expert_acts = [expert_acts_fns[i % len(expert_acts_fns)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Per-step Sigma Gates (one per draft per step) + self.draft_gates = nn.ModuleList([ + nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) + for _ in range(n_drafts) + ]) + for _ in range(reasoning_steps) + ]) + self.register_buffer( + "expert_bias", torch.zeros(n_experts), persistent=False + ) + + # 5. Cross-Draft Attention (lightweight) + self.cross_draft_q = nn.ModuleList([ + nn.Linear(out_dim, out_dim, bias=False) for _ in range(reasoning_steps) + ]) + self.cross_draft_k = nn.ModuleList([ + nn.Linear(out_dim, out_dim, bias=False) for _ in range(reasoning_steps) + ]) + self.cross_draft_v = nn.ModuleList([ + nn.Linear(out_dim, out_dim, bias=False) for _ in range(reasoning_steps) + ]) + self.cross_draft_norm = nn.ModuleList([ + nn.LayerNorm(out_dim) for _ in range(reasoning_steps) + ]) + self.cross_draft_scale = (out_dim) ** -0.5 + + # 6. PonderNet-style Halting Gates + self.halt_gates = nn.ModuleList([ + nn.Linear(in_dim, 1, bias=True) for _ in range(reasoning_steps) + ]) + + # 7. Consistency Network (Aggregation Arbiter) + # Takes all K draft predictions, scores internal consistency, + # and produces K confidence weights for final fusion. + consistency_in = n_drafts * out_dim + self.consistency_net = nn.Sequential( + nn.Linear(consistency_in, 128), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(128, 64), + nn.GELU(), + nn.Linear(64, n_drafts), + ) + + self.alpha = nn.Parameter(torch.tensor(0.0)) # deliberation scale + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.mem_read_query.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.mem_write_query.weight, a=math.sqrt(5)) + nn.init.normal_(self.mem_write_gate.weight, std=0.02) + nn.init.constant_(self.mem_write_gate.bias, 0.0) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.expert_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.expert_down[i].weight, std=0.01) + + for step_gates in self.draft_gates: + for g in step_gates: + nn.init.normal_(g.weight, std=0.02) + + for step in range(self.reasoning_steps): + nn.init.kaiming_uniform_(self.cross_draft_q[step].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.cross_draft_k[step].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.cross_draft_v[step].weight, a=math.sqrt(5)) + + for h in self.halt_gates: + nn.init.constant_(h.bias, -3.0) # start with low halt probability + + for p in self.consistency_net.parameters(): + if p.dim() > 1: + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + else: + nn.init.zeros_(p) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # 1. Shared Expert (always-on) + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # 2. Pre-compute expert outputs (shared across drafts) + expert_outs = [] + for i in range(self.n_experts): + h = self.dropout(self.expert_acts[i](self.expert_up[i](x_flat))) + expert_outs.append(self.expert_norms[i](self.expert_down[i](h))) + expert_stack = torch.stack(expert_outs, dim=1) # (N, E, D_out) + + # 3. Initialize draft states + draft_features = [x_flat.clone() for _ in range(self.n_drafts)] + draft_preds = [torch.zeros(N, self.out_dim, device=x.device) + for _ in range(self.n_drafts)] + + cumulative_halt = torch.zeros(N, 1, device=x.device) + total_deliberated = torch.zeros(N, self.out_dim, device=x.device) + depth_count = 0 + + for step in range(self.reasoning_steps): + depth_count += 1 + + # A. Memory Read: each draft queries the working memory + for d in range(self.n_drafts): + read_q = self.mem_read_query(draft_features[d]) # (N, D) + attn_scores = torch.matmul( + read_q, self.memory_keys.t() # (N, M) + ) * self.mem_scale + attn_weights = torch.softmax(attn_scores, dim=-1) # (N, M) + mem_read = torch.matmul(attn_weights, self.memory_values) # (N, D) + draft_features[d] = draft_features[d] + mem_read + + # B. Per-draft reasoning cell refinement + for d in range(self.n_drafts): + residual = draft_features[d] + draft_features[d] = residual + self.draft_reasoning_cells[step][d](residual) + + # C. Per-draft expert MoE routing (Sigma Gating) + step_draft_preds = [] + for d in range(self.n_drafts): + gate_logits = self.draft_gates[step][d](draft_features[d]) + if self.training: + gate_logits = gate_logits + self.expert_bias + gate_scores = torch.sigmoid(gate_logits) # (N, E) + routed = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) # (N, D_out) + step_draft_preds.append(routed) + + # D. Cross-Draft Attention: drafts attend to each other + pred_stack = torch.stack(step_draft_preds, dim=1) # (N, K, D_out) + q = self.cross_draft_q[step](pred_stack) # (N, K, D_out) + k = self.cross_draft_k[step](pred_stack) # (N, K, D_out) + v = self.cross_draft_v[step](pred_stack) # (N, K, D_out) + attn = torch.bmm(q, k.transpose(1, 2)) * self.cross_draft_scale # (N, K, K) + attn = torch.softmax(attn, dim=-1) + attn = self.dropout(attn) + cross_out = torch.bmm(attn, v) # (N, K, D_out) + refined_preds = self.cross_draft_norm[step](pred_stack + cross_out) # (N, K, D_out) + + for d in range(self.n_drafts): + draft_preds[d] = refined_preds[:, d, :] + + # E. Memory Write: update memory with new insights + # Average draft features as write signal + write_signal = torch.stack(draft_features, dim=0).mean(dim=0) # (N, D) + write_q = self.mem_write_query(write_signal) # (N, D) + write_scores = torch.sigmoid(self.mem_write_gate(write_q)) # (N, M) + # Differentiable soft memory update: weighted retrieval from write signal + mem_retrieval = torch.matmul(write_scores, self.memory_values) # (N, D) + # Inject memory write contribution into draft features for next step + for d in range(self.n_drafts): + draft_features[d] = draft_features[d] + 0.1 * mem_retrieval + + # F. Halting gate (PonderNet-style) + # Use average draft features for halt decision + avg_features = torch.stack(draft_features, dim=0).mean(dim=0) + halt_prob = torch.sigmoid(self.halt_gates[step](avg_features)) # (N, 1) + + # G. Accumulate: weighted by remaining probability mass + remaining = 1.0 - cumulative_halt + + # Consistency-aware per-step aggregation across drafts + all_draft_preds = torch.stack(draft_preds, dim=1) # (N, K, D_out) + consistency_in = all_draft_preds.reshape(N, -1) # (N, K*D_out) + consistency_weights = torch.softmax( + self.consistency_net(consistency_in), dim=-1 + ) # (N, K) + step_fused = (all_draft_preds * consistency_weights.unsqueeze(-1)).sum(dim=1) # (N, D_out) + + total_deliberated = total_deliberated + remaining * step_fused + cumulative_halt = cumulative_halt + remaining * halt_prob + + # H. Dynamic bias load balancing (across all drafts) + if self.training: + with torch.no_grad(): + all_gates = torch.stack( + [torch.sigmoid(self.draft_gates[step][d](draft_features[d])) + for d in range(self.n_drafts)], dim=0 + ).mean(dim=0) # (N, E) + target_load = 1.0 / self.n_experts + actual_load = all_gates.mean(dim=0) + self.expert_bias.add_(0.01 * (target_load - actual_load)) + + # I. Early exit during inference + if not self.training and cumulative_halt.mean() > 0.9: + break + + # Normalize by effective depth + total_deliberated = total_deliberated / max(1, depth_count) + + # Reshape back + base_logits = base_logits.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + total_deliberated = total_deliberated.view(*shape_prefix, -1) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * total_deliberated + ) + + +class ChampionNetDeliberativeExpert(nn.Module): + """ + Backbone-compatible model with DeliberativeAlignmentExpertHead (v19). + """ + def __init__( + self, + n_experts: int = 6, + n_drafts: int = 3, + n_mem_slots: int = 16, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(DeliberativeAlignmentExpertHead( + 256, 10, + n_experts=n_experts, + n_drafts=n_drafts, + n_mem_slots=n_mem_slots, + reasoning_steps=reasoning_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class OmniscientSynergyExpertHead(nn.Module): + """ + Generation v20: Omniscient Synergy Expert + (Latent Knowledge Core + Graph-of-Thought + Stochastic MoE) + + The most advanced intelligence architecture, moving beyond parallel drafts + to an interactive Graph of Thoughts grounded in a Latent Knowledge Core. + + Key Innovations: + 1. Latent Knowledge Core (LKC) — queries a massive continuous associative memory + 2. Graph-of-Thought (GoT) — thoughts exchange information via Graph Attention (GAT) + 3. Stochastic Bayesian Routing — samples expert routes via reparameterization (mu/var) + 4. Meta-Cognitive Self-Correction — global critique network broadcasts correction + 5. Dynamic Node Halting — per-thought confidence evaluation + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 6, + n_nodes: int = 4, + n_knowledge_slots: int = 64, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + self.n_nodes = n_nodes + self.n_knowledge_slots = n_knowledge_slots + self.reasoning_steps = reasoning_steps + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # 1. Global Shared Expert (always-on backbone) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 2. Latent Knowledge Core (Implicit RAG) + self.knowledge_core_keys = nn.Parameter(torch.randn(n_knowledge_slots, in_dim) * 0.02) + self.knowledge_core_values = nn.Parameter(torch.randn(n_knowledge_slots, in_dim) * 0.02) + self.knowledge_query = nn.Linear(in_dim, in_dim, bias=False) + self.knowledge_scale = in_dim ** -0.5 + + # 3. Graph-of-Thought (GoT) Attention (Node-to-Node communication) + self.got_q = nn.ModuleList([nn.Linear(in_dim, in_dim, bias=False) for _ in range(reasoning_steps)]) + self.got_k = nn.ModuleList([nn.Linear(in_dim, in_dim, bias=False) for _ in range(reasoning_steps)]) + self.got_v = nn.ModuleList([nn.Linear(in_dim, in_dim, bias=False) for _ in range(reasoning_steps)]) + self.got_norm = nn.ModuleList([nn.LayerNorm(in_dim) for _ in range(reasoning_steps)]) + + # 4. Stochastic Bayesian MoE Routing + expert_dims = [1024, 1536, 2048, 1024, 1536, 2048] + expert_acts_fns = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.expert_up = nn.ModuleList([ + nn.Linear(in_dim, expert_dims[i % len(expert_dims)], bias=False) + for i in range(n_experts) + ]) + self.expert_down = nn.ModuleList([ + nn.Linear(expert_dims[i % len(expert_dims)], out_dim, bias=False) + for i in range(n_experts) + ]) + self.expert_acts = [expert_acts_fns[i % len(expert_acts_fns)] for i in range(n_experts)] + self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Bayesian Routers (predicts mean and log-variance for expert weighting) + self.route_mu = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps) + ]) + self.route_logvar = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=True) for _ in range(reasoning_steps) + ]) + + # 5. Meta-Cognitive Critique (Global Review & Broadcast) + self.critique_net = nn.ModuleList([ + nn.Sequential( + nn.Linear(n_nodes * in_dim, 512), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(512, in_dim) + ) for _ in range(reasoning_steps) + ]) + + # 6. Dynamic Depth (Per-Node Halting) + self.node_halt_gates = nn.ModuleList([ + nn.Linear(in_dim, 1, bias=True) for _ in range(reasoning_steps) + ]) + + # Final Node Fusion + self.node_fusion_q = nn.Linear(out_dim, out_dim, bias=False) + self.node_fusion_k = nn.Linear(out_dim, out_dim, bias=False) + self.node_fusion_v = nn.Linear(out_dim, out_dim, bias=False) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # Divergence temperature for Bayesian routing during inference + self.register_buffer("temperature", torch.tensor(0.2)) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.knowledge_query.weight, a=math.sqrt(5)) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.expert_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.expert_down[i].weight, std=0.01) + + for step in range(self.reasoning_steps): + nn.init.kaiming_uniform_(self.got_q[step].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.got_k[step].weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.got_k[step].weight, a=math.sqrt(5)) + + nn.init.normal_(self.route_mu[step].weight, std=0.02) + nn.init.normal_(self.route_logvar[step].weight, std=0.02) + nn.init.constant_(self.route_logvar[step].bias, -2.0) # start with low variance + + for p in self.critique_net[step].parameters(): + if p.dim() > 1: nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + else: nn.init.zeros_(p) + + nn.init.constant_(self.node_halt_gates[step].bias, -2.0) + + nn.init.kaiming_uniform_(self.node_fusion_q.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.node_fusion_k.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.node_fusion_v.weight, a=math.sqrt(5)) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # 1. Global Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # 2. Latent Knowledge Core Retrieval + # Query the massive internal memory before reasoning begins + q = self.knowledge_query(x_flat) # (N, D) + attn = torch.matmul(q, self.knowledge_core_keys.t()) * self.knowledge_scale # (N, M) + attn = torch.softmax(attn, dim=-1) # (N, M) + knowledge_context = torch.matmul(attn, self.knowledge_core_values) # (N, D) + + # 3. Initialize Graph-of-Thought (N nodes) + # Seed nodes with input, knowledge, and diverse noise to split perspectives + nodes = [] + for i in range(self.n_nodes): + noise = torch.randn_like(x_flat) * 0.05 + nodes.append(x_flat + knowledge_context + noise) + + node_stack = torch.stack(nodes, dim=1) # (N, V, D) + node_cumulative_halt = torch.zeros(N, self.n_nodes, 1, device=x.device) + total_node_preds = torch.zeros(N, self.n_nodes, self.out_dim, device=x.device) + + # Pre-compute expert outputs to save compute + expert_outs = [] + for i in range(self.n_experts): + h = self.dropout(self.expert_acts[i](self.expert_up[i](x_flat))) + expert_outs.append(self.expert_norms[i](self.expert_down[i](h))) + expert_stack = torch.stack(expert_outs, dim=1) # (N, E, D_out) + + kl_loss = 0.0 + + for step in range(self.reasoning_steps): + + # A. Graph Attention: Nodes exchange information + q_got = self.got_q[step](node_stack) # (N, V, D) + k_got = self.got_k[step](node_stack) # (N, V, D) + v_got = self.got_v[step](node_stack) # (N, V, D) + + got_attn = torch.bmm(q_got, k_got.transpose(1, 2)) * self.knowledge_scale # (N, V, V) + got_attn = self.dropout(torch.softmax(got_attn, dim=-1)) + got_updates = torch.bmm(got_attn, v_got) # (N, V, D) + + node_stack = self.got_norm[step](node_stack + got_updates) + + # B. Meta-Cognitive Critique + # Global view of the graph spots logic errors, broadcasts correction + global_state = node_stack.reshape(N, -1) # (N, V*D) + global_correction = self.critique_net[step](global_state) # (N, D) + + # Apply correction to all active nodes + node_stack = node_stack + global_correction.unsqueeze(1) + + # C. Stochastic Bayesian MoE Routing & Halting + step_node_preds = [] + step_halt_probs = [] + + for v in range(self.n_nodes): + node_feat = node_stack[:, v, :] # (N, D) + + # Halt eval + halt_prob = torch.sigmoid(self.node_halt_gates[step](node_feat)) # (N, 1) + step_halt_probs.append(halt_prob) + + # Bayesian Routing + mu = self.route_mu[step](node_feat) # (N, E) + logvar = self.route_logvar[step](node_feat) # (N, E) + var = torch.exp(logvar) + + if self.training: + # Reparameterization trick: sample routing weights + eps = torch.randn_like(mu) + sample = mu + eps * torch.sqrt(var) + # KL divergence from standard normal prior + kl_loss = kl_loss + 0.5 * torch.sum(var + mu**2 - 1 - logvar, dim=1).mean() + else: + # In eval, add controllable random divergence + eps = torch.randn_like(mu) + sample = mu + eps * torch.sqrt(var) * self.temperature + + gate_scores = torch.sigmoid(sample) # (N, E) + + # Execute experts and fusion + routed = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) # (N, D_out) + step_node_preds.append(routed) + + # D. Accumulate Predictions + step_preds_stack = torch.stack(step_node_preds, dim=1) # (N, V, D_out) + step_halts_stack = torch.stack(step_halt_probs, dim=1) # (N, V, 1) + + remaining = 1.0 - node_cumulative_halt + total_node_preds = total_node_preds + remaining * step_preds_stack + node_cumulative_halt = node_cumulative_halt + remaining * step_halts_stack + + if not self.training and node_cumulative_halt.mean() > 0.95: + break + + # Optional: return KL loss to be added to total loss via an attribute + self.last_kl_loss = kl_loss * 0.001 + + # Final Node Fusion (Attentional pooling of Graph) + fusion_q = self.node_fusion_q(total_node_preds).mean(dim=1, keepdim=True) # (N, 1, D_out) + fusion_k = self.node_fusion_k(total_node_preds) # (N, V, D_out) + fusion_v = self.node_fusion_v(total_node_preds) # (N, V, D_out) + + pool_attn = torch.bmm(fusion_q, fusion_k.transpose(1, 2)) * (self.out_dim**-0.5) # (N, 1, V) + pool_attn = torch.softmax(pool_attn, dim=-1) + fused_got_out = torch.bmm(pool_attn, fusion_v).squeeze(1) # (N, D_out) + + # Reshape back + base_logits = base_logits.view(*shape_prefix, -1) + shared_out = shared_out.view(*shape_prefix, -1) + fused_got_out = fused_got_out.view(*shape_prefix, -1) + + return ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * fused_got_out + ) + + +class ChampionNetOmniscientExpert(nn.Module): + """ + Backbone-compatible model with OmniscientSynergyExpertHead (v20). + """ + def __init__( + self, + n_experts: int = 6, + n_nodes: int = 4, + n_knowledge_slots: int = 64, + reasoning_steps: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(OmniscientSynergyExpertHead( + 256, 10, + n_experts=n_experts, + n_nodes=n_nodes, + n_knowledge_slots=n_knowledge_slots, + reasoning_steps=reasoning_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class NeurogenesisExpertHead(nn.Module): + """ + Generation v21: Neurogenesis Expert + (Hierarchical Abstraction + Adversarial Self-Play) + + Processes input at multiple abstraction levels, uses adversarial debate + between a Proposer and Adversary to produce robust conclusions, and + dynamically synthesizes emergent expert combinations. + + Key Innovations: + 1. Hierarchical Abstraction Pyramid — 3 levels (concrete/abstract/meta) + 2. Adversarial Self-Play — Proposer vs Adversary MoE debate + 3. Cross-Level Knowledge Banks — per-level memory with distillation + 4. Emergent Expert Synthesis — dynamic expert weight interpolation + 5. Recursive Confidence Calibration — self-assessed certainty triggers extra rounds + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_experts: int = 6, + n_levels: int = 3, + n_knowledge_slots: int = 32, + debate_rounds: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + self.n_levels = n_levels + self.debate_rounds = debate_rounds + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 1. Hierarchical Abstraction — Up/Down projections between levels + self.level_up_projs = nn.ModuleList([ + nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, in_dim), nn.GELU()) + for _ in range(n_levels - 1) + ]) + self.level_down_projs = nn.ModuleList([ + nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, in_dim), nn.GELU()) + for _ in range(n_levels - 1) + ]) + + # 2. Cross-Level Knowledge Banks (one per level) + self.knowledge_keys = nn.ParameterList([ + nn.Parameter(torch.randn(n_knowledge_slots, in_dim) * 0.02) + for _ in range(n_levels) + ]) + self.knowledge_values = nn.ParameterList([ + nn.Parameter(torch.randn(n_knowledge_slots, in_dim) * 0.02) + for _ in range(n_levels) + ]) + self.knowledge_queries = nn.ModuleList([ + nn.Linear(in_dim, in_dim, bias=False) for _ in range(n_levels) + ]) + self.knowledge_scale = in_dim ** -0.5 + + # Distillation bridges (transfer knowledge between adjacent levels) + self.distill_bridges = nn.ModuleList([ + nn.Linear(in_dim, in_dim, bias=False) for _ in range(n_levels - 1) + ]) + + # 3. Proposer MoE (diverse activations) + expert_dims = [1024, 1536, 2048, 1024, 1536, 2048] + act_fns = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + self.proposer_up = nn.ModuleList([ + nn.Linear(in_dim, expert_dims[i % len(expert_dims)], bias=False) + for i in range(n_experts) + ]) + self.proposer_down = nn.ModuleList([ + nn.Linear(expert_dims[i % len(expert_dims)], out_dim, bias=False) + for i in range(n_experts) + ]) + self.proposer_acts = [act_fns[i % len(act_fns)] for i in range(n_experts)] + self.proposer_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Per-round proposer gates (Sigma Gating) + self.proposer_gates = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) for _ in range(debate_rounds) + ]) + + # 4. Adversary MoE (separate experts for critique) + self.adversary_up = nn.ModuleList([ + nn.Linear(in_dim, expert_dims[i % len(expert_dims)], bias=False) + for i in range(n_experts) + ]) + self.adversary_down = nn.ModuleList([ + nn.Linear(expert_dims[i % len(expert_dims)], out_dim, bias=False) + for i in range(n_experts) + ]) + self.adversary_acts = [act_fns[(i + 3) % len(act_fns)] for i in range(n_experts)] + self.adversary_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) + + # Per-round adversary gates + self.adversary_gates = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=False) for _ in range(debate_rounds) + ]) + + # 5. Emergent Expert Synthesis — interpolation weights + self.synthesis_net = nn.ModuleList([ + nn.Sequential( + nn.Linear(in_dim, 64), nn.GELU(), nn.Linear(64, n_experts) + ) for _ in range(debate_rounds) + ]) + + # 6. Resolution Network (merges proposer + adversary) + self.resolution = nn.ModuleList([ + nn.Sequential( + nn.Linear(out_dim * 2, 128), nn.GELU(), + nn.Dropout(dropout), nn.Linear(128, out_dim) + ) for _ in range(debate_rounds) + ]) + + # 7. Confidence Calibration + self.confidence_gates = nn.ModuleList([ + nn.Linear(in_dim, 1, bias=True) for _ in range(debate_rounds) + ]) + + # 8. Level fusion (aggregate across abstraction levels) + self.level_fusion = nn.Sequential( + nn.Linear(n_levels * out_dim, 128), nn.GELU(), + nn.Linear(128, out_dim) + ) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + for i in range(self.n_experts): + nn.init.kaiming_uniform_(self.proposer_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.proposer_down[i].weight, std=0.01) + nn.init.kaiming_uniform_(self.adversary_up[i].weight, a=math.sqrt(5)) + nn.init.normal_(self.adversary_down[i].weight, std=0.01) + + for g in self.proposer_gates: + nn.init.normal_(g.weight, std=0.02) + for g in self.adversary_gates: + nn.init.normal_(g.weight, std=0.02) + + for c in self.confidence_gates: + nn.init.constant_(c.bias, -2.0) + + for q in self.knowledge_queries: + nn.init.kaiming_uniform_(q.weight, a=math.sqrt(5)) + for b in self.distill_bridges: + nn.init.kaiming_uniform_(b.weight, a=math.sqrt(5)) + + for net in [self.level_fusion, *self.resolution, *self.synthesis_net]: + for p in net.parameters(): + if p.dim() > 1: + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + else: + nn.init.zeros_(p) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # ═══ PHASE 1: Hierarchical Abstraction Pyramid ═══ + # Build multi-level representations + level_features = [x_flat] # Level 0 = concrete + h = x_flat + for i in range(self.n_levels - 1): + h = h + self.level_up_projs[i](h) # Residual up-projection + level_features.append(h) + + # Query cross-level knowledge banks + for lvl in range(self.n_levels): + q = self.knowledge_queries[lvl](level_features[lvl]) + attn = torch.matmul(q, self.knowledge_keys[lvl].t()) * self.knowledge_scale + attn = torch.softmax(attn, dim=-1) + knowledge = torch.matmul(attn, self.knowledge_values[lvl]) + level_features[lvl] = level_features[lvl] + knowledge + + # Cross-level distillation (top-down grounding) + for i in range(self.n_levels - 2, -1, -1): + distilled = self.distill_bridges[i](level_features[i + 1]) + level_features[i] = level_features[i] + 0.1 * distilled + + # Merge all levels into a unified reasoning state + unified = torch.stack(level_features, dim=0).mean(dim=0) # (N, D) + + # Pre-compute proposer expert outputs + proposer_outs = [] + for i in range(self.n_experts): + h = self.dropout(self.proposer_acts[i](self.proposer_up[i](unified))) + proposer_outs.append(self.proposer_norms[i](self.proposer_down[i](h))) + proposer_stack = torch.stack(proposer_outs, dim=1) # (N, E, D_out) + + # Pre-compute adversary expert outputs + adversary_outs = [] + for i in range(self.n_experts): + h = self.dropout(self.adversary_acts[i](self.adversary_up[i](unified))) + adversary_outs.append(self.adversary_norms[i](self.adversary_down[i](h))) + adversary_stack = torch.stack(adversary_outs, dim=1) # (N, E, D_out) + + # ═══ PHASE 2: Adversarial Debate Loop ═══ + cumulative_confidence = torch.zeros(N, 1, device=x.device) + total_resolved = torch.zeros(N, self.out_dim, device=x.device) + debate_state = unified # Evolves through debate + + for rnd in range(self.debate_rounds): + + # A. Emergent Expert Synthesis — dynamic interpolation weights + synth_weights = torch.softmax( + self.synthesis_net[rnd](debate_state), dim=-1 + ) # (N, E) + + # B. Proposer: route through experts with synthesis + prop_gates = torch.sigmoid(self.proposer_gates[rnd](debate_state)) + combined_prop_gates = prop_gates * synth_weights # Emergent combination + proposal = (proposer_stack * combined_prop_gates.unsqueeze(-1)).sum(dim=1) + + # C. Adversary: critique the proposal + adv_gates = torch.sigmoid(self.adversary_gates[rnd](debate_state)) + critique = (adversary_stack * adv_gates.unsqueeze(-1)).sum(dim=1) + + # D. Resolution: merge proposer + adversary insights + combined = torch.cat([proposal, critique], dim=-1) # (N, 2*D_out) + resolved = self.resolution[rnd](combined) # (N, D_out) + + # E. Confidence calibration + confidence = torch.sigmoid(self.confidence_gates[rnd](debate_state)) # (N, 1) + + # F. Accumulate with remaining probability mass + remaining = 1.0 - cumulative_confidence + total_resolved = total_resolved + remaining * resolved + cumulative_confidence = cumulative_confidence + remaining * confidence + + # G. Early exit if highly confident (inference only) + if not self.training and cumulative_confidence.mean() > 0.9: + break + + # ═══ PHASE 3: Final Assembly ═══ + # Also produce per-level predictions for richer signal + level_preds = [] + for lvl in range(self.n_levels): + level_preds.append(F.linear(level_features[lvl], self.weight, self.bias)) + level_concat = torch.cat(level_preds, dim=-1) # (N, L*D_out) + level_fused = self.level_fusion(level_concat) # (N, D_out) + + # Combine everything + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * (total_resolved + 0.1 * level_fused) + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetNeurogenesisExpert(nn.Module): + """ + Backbone-compatible model with NeurogenesisExpertHead (v21). + """ + def __init__( + self, + n_experts: int = 6, + n_levels: int = 3, + n_knowledge_slots: int = 32, + debate_rounds: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(NeurogenesisExpertHead( + 256, 10, + n_experts=n_experts, + n_levels=n_levels, + n_knowledge_slots=n_knowledge_slots, + debate_rounds=debate_rounds, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + + + +class CognitiveSingularityExpertHead(nn.Module): + """ + Generation v22: Cognitive Singularity Expert + (Hypernetwork + Latent Tree-Search + Episodic Memory + Orthogonal Superposition) + + Pushes reasoning to its limits using fully dynamic computation: + 1. Hypernetwork-Generated Experts — dynamically generates weights based on input + 2. Latent Tree-Search — Value Network rolls out and evaluates paths like AlphaGo + 3. Episodic Neural Memory — differentiable read/write memory matrix for context + 4. Orthogonal Superposition — forces generation of mutually exclusive hypotheses + 5. Causal Logic Refinement — self-attention enforces logical sequence flow + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_hypotheses: int = 4, + hyper_dim: int = 64, # Custom expert hidden dimension + mem_slots: int = 64, + mem_dim: int = 128, + search_depth: int = 3, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_hypotheses = n_hypotheses + self.hyper_dim = hyper_dim + self.mem_slots = mem_slots + self.mem_dim = mem_dim + self.search_depth = search_depth + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # 1. Episodic Neural Memory (Fully Differentiable) + self.memory_matrix = nn.Parameter(torch.randn(mem_slots, mem_dim) * 0.02) + # Read heads + self.mem_read_q = nn.Linear(in_dim, mem_dim, bias=False) + self.mem_read_out = nn.Linear(mem_dim, in_dim, bias=False) + # Write heads + self.mem_write_gate = nn.Linear(in_dim, mem_slots, bias=True) + self.mem_write_data = nn.Linear(in_dim, mem_dim, bias=False) + self.mem_write_proj = nn.Linear(mem_dim, out_dim, bias=False) + + # 2. Hypernetwork (Generates custom expert weights on-the-fly) + # Input: state + memory context (2 * in_dim) + # Output: Weights for Linear(in_dim, hyper_dim) and Linear(hyper_dim, out_dim) + self.hypernet_input_dim = in_dim * 2 + self.hypernet_out_size = (in_dim * hyper_dim) + (hyper_dim * out_dim) + self.hypernet = nn.Sequential( + nn.Linear(self.hypernet_input_dim, 512), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(512, self.hypernet_out_size) + ) + + # 3. Orthogonal Superposition + # Forces the model to span K disjoint regions of thought-space + self.path_projs = nn.ModuleList([ + nn.Linear(in_dim, in_dim, bias=False) for _ in range(n_hypotheses) + ]) + + # 4. Latent Tree-Search Value Network + # Evaluates the "goodness" or logical consistency of a hypothesis + self.value_net = nn.Sequential( + nn.Linear(out_dim, 64), + nn.GELU(), + nn.Linear(64, 1) + ) + + # 5. Causal Logic Refinement + self.causal_q = nn.Linear(out_dim, out_dim, bias=False) + self.causal_k = nn.Linear(out_dim, out_dim, bias=False) + self.causal_v = nn.Linear(out_dim, out_dim, bias=False) + self.causal_norm = nn.LayerNorm(out_dim) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # Track orthogonal penalty loss for debugging/regularization + self.register_buffer("last_ortho_loss", torch.tensor(0.0)) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.mem_read_q.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.mem_read_out.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.mem_write_gate.weight, a=math.sqrt(5)) + nn.init.constant_(self.mem_write_gate.bias, -2.0) + nn.init.kaiming_uniform_(self.mem_write_data.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.mem_write_proj.weight, a=math.sqrt(5)) + + # Hypernetwork init + for layer in self.hypernet: + if isinstance(layer, nn.Linear): + nn.init.normal_(layer.weight, std=0.01) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + + # Output of hypernetwork should be initialized small to prevent explosion + nn.init.normal_(self.hypernet[-1].weight, std=0.001) + + for proj in self.path_projs: + nn.init.kaiming_uniform_(proj.weight, a=math.sqrt(5)) + + for layer in self.value_net: + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + + nn.init.kaiming_uniform_(self.causal_q.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.causal_k.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.causal_v.weight, a=math.sqrt(5)) + + def compute_orthogonality_loss(self, hypotheses): + # hypotheses: (N, K, D) + normed = F.normalize(hypotheses, p=2, dim=-1) # (N, K, D) + sim_matrix = torch.bmm(normed, normed.transpose(1, 2)) # (N, K, K) + # Mask out diagonal (self-similarity is 1) + eye = torch.eye(sim_matrix.size(1), device=sim_matrix.device).unsqueeze(0) + sim_matrix = sim_matrix * (1.0 - eye) + # Repulsion loss: minimize absolute cosine similarity between different hypotheses + ortho_loss = sim_matrix.abs().sum() / (self.n_hypotheses * (self.n_hypotheses - 1)) + return ortho_loss + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # 1. Global Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # 2. Episodic Memory Read + read_q = self.mem_read_q(x_flat) # (N, mem_dim) + attn = torch.matmul(read_q, self.memory_matrix.t()) / math.sqrt(self.mem_dim) # (N, M) + attn = torch.softmax(attn, dim=-1) # (N, M) + mem_retrieval = torch.matmul(attn, self.memory_matrix) # (N, mem_dim) + mem_context = self.mem_read_out(mem_retrieval) # (N, in_dim) + + # 3. Hypernetwork Weight Generation + # Context is x_flat + episodic memory recall + hyper_input = torch.cat([x_flat, mem_context], dim=-1) # (N, 2*in_dim) + generated_weights = self.hypernet(hyper_input) # (N, W) + + w1_size = self.in_dim * self.hyper_dim + w2_size = self.hyper_dim * self.out_dim + + w1_flat = generated_weights[:, :w1_size] + w2_flat = generated_weights[:, w1_size:] + + # Reshape generated weights: (N, hyper_dim, in_dim) and (N, out_dim, hyper_dim) + W1 = w1_flat.view(N, self.hyper_dim, self.in_dim) + W2 = w2_flat.view(N, self.out_dim, self.hyper_dim) + + # 4. Orthogonal Superposition & Dynamic Execution + # Generate K divergent hypotheses and run them through the custom dynamic expert + hypotheses = [] + custom_outs = [] + for k in range(self.n_hypotheses): + h_in = self.path_projs[k](x_flat) # (N, D) + hypotheses.append(h_in) + + # Execute dynamically generated expert layer 1: (N, hyper_dim) + h_hidden = self.dropout(F.gelu(torch.bmm(W1, h_in.unsqueeze(-1)).squeeze(-1))) + + # Execute dynamically generated expert layer 2: (N, out_dim) + h_out = torch.bmm(W2, h_hidden.unsqueeze(-1)).squeeze(-1) + custom_outs.append(h_out) + + hypotheses_stack = torch.stack(hypotheses, dim=1) # (N, K, D) + custom_outs_stack = torch.stack(custom_outs, dim=1) # (N, K, D_out) + + # Record orthogonality penalty + if self.training: + self.last_ortho_loss = self.compute_orthogonality_loss(hypotheses_stack) + + # 5. Latent Tree-Search Value Network + # Evaluate expected reward of each hypothesis + values = self.value_net(custom_outs_stack).squeeze(-1) # (N, K) + # Select best hypotheses (softmax converts to probabilities) + path_probs = torch.softmax(values, dim=-1) # (N, K) + + # 6. Causal Logic Refinement + # Enforce that the top hypotheses are logically coherent sequences + causal_in = custom_outs_stack # (N, K, D_out) + q = self.causal_q(causal_in) + k = self.causal_k(causal_in) + v = self.causal_v(causal_in) + + # Masked attention (causal path flow) + causal_attn = torch.bmm(q, k.transpose(1, 2)) / math.sqrt(self.out_dim) + mask = torch.tril(torch.ones(self.n_hypotheses, self.n_hypotheses, device=x.device)) + causal_attn = causal_attn.masked_fill(mask == 0, float('-inf')) + causal_attn = self.dropout(torch.softmax(causal_attn, dim=-1)) + + refined_hypotheses = self.causal_norm(causal_in + torch.bmm(causal_attn, v)) # (N, K, D_out) + + # 7. Episodic Memory Write + # Write the selected (highest value) path logic back to memory for future steps + best_path_idx = path_probs.argmax(dim=-1) # (N) + best_h_in = hypotheses_stack[torch.arange(N), best_path_idx] # (N, D) + + write_gates = torch.sigmoid(self.mem_write_gate(best_h_in)) # (N, M) + write_data = self.mem_write_data(best_h_in) # (N, mem_dim) + + # Differentiable update to parameter matrix + # This operates as a residual accumulation gradient update during training + # For true recurrent state changes across tokens, an external hidden state is needed, + # but the Parameter acts as a static "long term episodic registry" + mem_update = torch.matmul(write_gates.t(), write_data) / (write_gates.sum(dim=0, keepdim=True).t() + 1e-6) + + # Add to predictions + # Weigh refined hypotheses by their tree-search value + fused_cognitive = (refined_hypotheses * path_probs.unsqueeze(-1)).sum(dim=1) # (N, D_out) + + # Add a residual from the memory write operation so write heads receive gradients + fused_write = self.mem_write_proj(write_data * write_gates.mean(dim=-1, keepdim=True)) + fused_cognitive = fused_cognitive + 0.1 * fused_write + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * fused_cognitive + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetCognitiveExpert(nn.Module): + """ + Backbone-compatible model with CognitiveSingularityExpertHead (v22). + """ + def __init__( + self, + n_hypotheses: int = 4, + hyper_dim: int = 64, + mem_slots: int = 64, + mem_dim: int = 128, + search_depth: int = 3, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(CognitiveSingularityExpertHead( + 256, 10, + n_hypotheses=n_hypotheses, + hyper_dim=hyper_dim, + mem_slots=mem_slots, + mem_dim=mem_dim, + search_depth=search_depth, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TranscendentArchitectExpertHead(nn.Module): + """ + Generation v23: Transcendent Architect Expert + (Diffusion Refinement + World Model + Compositional Decomposition + + Self-Modifying Attention + Evolution Layer) + + Instead of computing answers in a single pass, v23 sculpts its answer + from noise using iterative denoising, simulates consequences via a + learned World Model, and evolves candidate solutions. + + Key Innovations: + 1. Diffusion-Based Iterative Refinement — denoise answer from noise over T steps + 2. World Model Simulation — predict future states to score candidate quality + 3. Compositional Decomposition — break input into K sub-problems, solve & compose + 4. Self-Modifying Attention — meta-controller rewires Q/K/V mid-computation + 5. Gradient-Free Evolution — tournament selection + crossover in latent space + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_subproblems: int = 4, + pop_size: int = 6, + denoise_steps: int = 4, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_subproblems = n_subproblems + self.pop_size = pop_size + self.denoise_steps = denoise_steps + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # ═══ 1. Compositional Decomposition ═══ + # Learned decomposition heads: split input into K sub-representations + self.decompose_heads = nn.ModuleList([ + nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU() + ) for _ in range(n_subproblems) + ]) + + # ═══ 2. Diffusion-Based Iterative Refinement ═══ + # Per-step denoisers (each step gets its own denoiser) + self.denoiser = nn.ModuleList([ + nn.Sequential( + nn.Linear(out_dim + in_dim, 256), # noisy_z + context + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(256, out_dim) + ) for _ in range(denoise_steps) + ]) + + # Timestep embeddings (learnable) + self.time_embed = nn.Parameter(torch.randn(denoise_steps, out_dim) * 0.02) + + # Noise schedule (learnable betas) + self.log_betas = nn.Parameter(torch.linspace(-3, -1, denoise_steps)) + + # ═══ 3. World Model Simulation ═══ + # Transition function: predicts next state from current + self.world_transition = nn.Sequential( + nn.Linear(out_dim, 128), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(128, out_dim) + ) + # Value function: scores how "good" a state is + self.world_value = nn.Sequential( + nn.Linear(out_dim, 64), + nn.GELU(), + nn.Linear(64, 1) + ) + + # ═══ 4. Self-Modifying Attention ═══ + # Base Q/K/V projections + self.attn_q = nn.Linear(out_dim, out_dim, bias=False) + self.attn_k = nn.Linear(out_dim, out_dim, bias=False) + self.attn_v = nn.Linear(out_dim, out_dim, bias=False) + self.attn_norm = nn.LayerNorm(out_dim) + + # Meta-controller: modifies Q/K/V weights based on intermediate results + self.meta_q_mod = nn.Linear(in_dim, out_dim * out_dim, bias=False) + self.meta_k_mod = nn.Linear(in_dim, out_dim * out_dim, bias=False) + + # ═══ 5. Evolution Layer ═══ + # Fitness projection (maps sub-solutions to fitness scores) + self.fitness_proj = nn.Linear(out_dim, 1) + # Crossover network (merges two parents into offspring) + self.crossover_net = nn.Sequential( + nn.Linear(out_dim * 2, 128), + nn.GELU(), + nn.Linear(128, out_dim) + ) + + # Final composition: merge sub-problem solutions + self.composer = nn.Sequential( + nn.Linear(n_subproblems * out_dim, 256), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(256, out_dim) + ) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + for head in self.decompose_heads: + for m in head.modules(): + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + for step_net in self.denoiser: + for m in step_net.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.01) + if m.bias is not None: + nn.init.zeros_(m.bias) + + for net in [self.world_transition, self.world_value, self.crossover_net, self.composer]: + for m in net.modules(): + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + nn.init.kaiming_uniform_(self.attn_q.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.attn_k.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.attn_v.weight, a=math.sqrt(5)) + + nn.init.normal_(self.meta_q_mod.weight, std=0.001) + nn.init.normal_(self.meta_k_mod.weight, std=0.001) + + nn.init.kaiming_uniform_(self.fitness_proj.weight, a=math.sqrt(5)) + nn.init.zeros_(self.fitness_proj.bias) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # ═══ PHASE 1: Compositional Decomposition ═══ + sub_contexts = [] + for k in range(self.n_subproblems): + sub_ctx = self.decompose_heads[k](x_flat) # (N, D) + sub_contexts.append(sub_ctx) + + # ═══ PHASE 2: Diffusion Refinement (per sub-problem) ═══ + betas = torch.sigmoid(self.log_betas) # (T,) in [0, 1] + sub_solutions = [] + + for k in range(self.n_subproblems): + ctx = sub_contexts[k] # (N, D_in) + + # Start from noise + z = torch.randn(N, self.out_dim, device=x.device) * 0.1 + + # Iteratively denoise + for t in range(self.denoise_steps): + beta_t = betas[t] + t_emb = self.time_embed[t].unsqueeze(0) # (1, D_out) + + # Condition on context + noisy state + timestep + z_conditioned = z + t_emb + denoiser_input = torch.cat([z_conditioned, ctx], dim=-1) # (N, D_out + D_in) + noise_pred = self.denoiser[t](denoiser_input) # (N, D_out) + + # Denoise step: remove predicted noise + z = z - beta_t * noise_pred + + sub_solutions.append(z) # (N, D_out) + + # ═══ PHASE 3: World Model Simulation ═══ + # Score each sub-solution by simulating one step into the future + world_scores = [] + for sol in sub_solutions: + future_state = self.world_transition(sol) # (N, D_out) + value = self.world_value(future_state).squeeze(-1) # (N,) + world_scores.append(value) + world_scores = torch.stack(world_scores, dim=1) # (N, K) + world_probs = torch.softmax(world_scores, dim=-1) # (N, K) + + # ═══ PHASE 4: Evolution Layer ═══ + # Create population from sub-solutions + pop = torch.stack(sub_solutions, dim=1) # (N, K, D_out) + + # Compute fitness scores + fitness = self.fitness_proj(pop).squeeze(-1) # (N, K) + + # Soft tournament selection: differentiable weighted parents + fitness_weights = torch.softmax(fitness * 5.0, dim=-1) # (N, K) temperature-sharpened + parent_a = (pop * fitness_weights.unsqueeze(-1)).sum(dim=1) # (N, D_out) + + # Second parent: use inverse fitness for diversity + inv_weights = torch.softmax(-fitness * 5.0, dim=-1) + parent_b = (pop * inv_weights.unsqueeze(-1)).sum(dim=1) # (N, D_out) + + # Crossover: merge parents into evolved offspring + offspring = self.crossover_net(torch.cat([parent_a, parent_b], dim=-1)) # (N, D_out) + + # ═══ PHASE 5: Self-Modifying Attention ═══ + # Use the offspring + world model insights to modify attention + # Meta-controller generates attention weight modifications + meta_dq = self.meta_q_mod(x_flat).view(N, self.out_dim, self.out_dim) # (N, D, D) + meta_dk = self.meta_k_mod(x_flat).view(N, self.out_dim, self.out_dim) # (N, D, D) + + # Apply self-modifying attention over sub-solutions + offspring + all_candidates = torch.stack(sub_solutions + [offspring], dim=1) # (N, K+1, D_out) + K1 = all_candidates.shape[1] + + # Base Q/K/V + q = self.attn_q(all_candidates) # (N, K+1, D_out) + k_proj = self.attn_k(all_candidates) + v = self.attn_v(all_candidates) + + # Apply meta-modifications (batched matrix multiply for Q and K) + q = q + torch.bmm(q, meta_dq) * 0.1 # Residual modification + k_proj = k_proj + torch.bmm(k_proj, meta_dk) * 0.1 + + # Compute attention + attn_scores = torch.bmm(q, k_proj.transpose(1, 2)) / math.sqrt(self.out_dim) + attn_weights = self.dropout(torch.softmax(attn_scores, dim=-1)) + attn_out = self.attn_norm(all_candidates + torch.bmm(attn_weights, v)) # (N, K+1, D_out) + + # ═══ PHASE 6: Final Composition ═══ + # Weight attended candidates by world model probabilities + offspring bonus + # Extend world_probs with offspring weight + offspring_score = self.world_value( + self.world_transition(offspring) + ).squeeze(-1).unsqueeze(-1) # (N, 1) + all_scores = torch.cat([world_scores, offspring_score], dim=-1) # (N, K+1) + all_probs = torch.softmax(all_scores, dim=-1) # (N, K+1) + + # Weighted sum of attended candidates + fused = (attn_out * all_probs.unsqueeze(-1)).sum(dim=1) # (N, D_out) + + # Also compose sub-solutions directly for a rich signal + composed = self.composer(torch.cat(sub_solutions, dim=-1)) # (N, D_out) + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * (fused + 0.1 * composed) + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetTranscendentExpert(nn.Module): + """ + Backbone-compatible model with TranscendentArchitectExpertHead (v23). + """ + def __init__( + self, + n_subproblems: int = 4, + pop_size: int = 6, + denoise_steps: int = 4, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(TranscendentArchitectExpertHead( + 256, 10, + n_subproblems=n_subproblems, + pop_size=pop_size, + denoise_steps=denoise_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ODEFunc(nn.Module): + """ + Continuous-Time Neural ODE Function for the Omniversal Quantum Simulation. + Defines the derivative dz/dt for the neural integrator. + """ + def __init__(self, dim: int, dropout: float = 0.1): + super().__init__() + # Project Phase + Magnitude + Time + self.net = nn.Sequential( + nn.Linear(dim * 2 + 1, dim * 2), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim * 2, dim * 2) + ) + + # Entanglement Routing approximations + self.q_proj = nn.Linear(dim * 2, dim) + self.k_proj = nn.Linear(dim * 2, dim) + self.v_proj = nn.Linear(dim * 2, dim * 2) + + def forward(self, t: torch.Tensor, z: torch.Tensor): + # z is (N, dim*2) representing complex state [Real, Imag] + N, D2 = z.shape + D = D2 // 2 + + # 1. Neural Mechanics Step + t_tensor = t.expand(N, 1) + z_t = torch.cat([z, t_tensor], dim=-1) + dz_dt_base = self.net(z_t) # (N, D*2) + + # 2. Entanglement Routing + # Instead of dot product similarity, we use a cheap pseudo-Kronecker entanglement + q = self.q_proj(z) # (N, D) + k = self.k_proj(z) # (N, D) + v = self.v_proj(z) # (N, D*2) + + # Entangle Q and K via element-wise multiplication (diagonal of Kronecker) + # then non-linear mix + entangled = torch.relu(q * k).sum(dim=-1, keepdim=True) / math.sqrt(D) # (N, 1) + routing_update = torch.sigmoid(entangled) * v # (N, D*2) + + return dz_dt_base + 0.1 * routing_update + + +class OmniversalQuantumExpertHead(nn.Module): + """ + Generation v24: Omniversal Quantum Simulation Expert + (Complex Superposition + Neural ODE + Holographic HRR + Entanglement) + + Abandons discrete real-valued layers for Continuous-Time Complex integration. + Allows mutual contradiction handling via wave interference before Measurement Collapse. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + ode_steps: int = 5, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.ode_steps = ode_steps # Discretization steps for Euler integrator + self.complex_dim = in_dim # Operate in complex space of size in_dim + + # Base projection (backward-compatible residual) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Real Expert (Standard MLP fallback) + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # ═══ 1. Quantum Lift (Measurement Prep) ═══ + # Lifts Real-valued input to Complex-valued Wavefunction (Real, Imag) + self.real_lift = nn.Linear(in_dim, self.complex_dim) + self.imag_lift = nn.Linear(in_dim, self.complex_dim) + + # ═══ 2. Holographic Knowledge Binding (HRR Context) ═══ + # Global holographic memory trace (learnable contextual superposition) + self.holographic_memory = nn.Parameter(torch.randn(self.complex_dim * 2)) + + # ═══ 3. Continuous-Time Neural ODE ═══ + # Defines the continuous derivative of the complex thought process + self.ode_func = ODEFunc(dim=self.complex_dim, dropout=dropout) + + # ═══ 4. Measurement Collapse ═══ + # Collapses the final Complex state back into distinct Real-valued logic (logits) + self.measurement_proj = nn.Sequential( + nn.Linear(self.complex_dim * 2, 256), + nn.GELU(), + nn.Linear(256, out_dim) + ) + + # Measurement uncertainty modulator (Phase variance modulates confidence) + self.uncertainty_gate = nn.Linear(self.complex_dim, 1) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.real_lift.weight, a=math.sqrt(5)) + nn.init.zeros_(self.real_lift.bias) + nn.init.kaiming_uniform_(self.imag_lift.weight, a=math.sqrt(5)) + nn.init.zeros_(self.imag_lift.bias) + + for m in self.ode_func.modules(): + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + for m in self.measurement_proj: + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + nn.init.kaiming_uniform_(self.uncertainty_gate.weight, a=math.sqrt(5)) + nn.init.zeros_(self.uncertainty_gate.bias) + + def circular_convolution(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """ + Holographic Reduced Representation (HRR) Binding via Fast Fourier Transform. + Binds two vectors in superposition such that x * y = z, where z contains features of both. + """ + # x and y are (N, D) + X = torch.fft.rfft(x, dim=-1) + Y = torch.fft.rfft(y, dim=-1) + Z = X * Y # Element-wise complex multiplication in frequency domain + return torch.fft.irfft(Z, n=x.shape[-1], dim=-1) + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Shared Real Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # ═══ PHASE 1: Quantum Lift (Superposition Initialization) ═══ + # Create Complex Wavefunction: Ψ = Real + i*Imag + z_real = self.real_lift(x_flat) # (N, D) + z_imag = self.imag_lift(x_flat) # (N, D) + + # Pack into (N, D*2) for compatibility with real-valued PyTorch ops + z0 = torch.cat([z_real, z_imag], dim=-1) # (N, D*2) + + # ═══ PHASE 2: Holographic Knowledge Binding ═══ + # Bind the current state with the global holographic memory trace + memory_expanded = self.holographic_memory.unsqueeze(0).expand(N, -1) # (N, D*2) + z_bound = self.circular_convolution(z0, memory_expanded) # (N, D*2) + + # Add bound knowledge into superposition + z_t = z0 + 0.1 * z_bound + + # ═══ PHASE 3: Continuous-Time Neural ODE Integration ═══ + # Solve the Initial Value Problem: z(1) = z(0) + ∫ f(z(t), t) dt + # using a simple Euler integrator (differentiable) + dt = 1.0 / self.ode_steps + t = torch.zeros(1, device=x.device) + + for _ in range(self.ode_steps): + # Compute derivative + dz_dt = self.ode_func(t, z_t) + # Step forward in continuous time + z_t = z_t + dz_dt * dt + t = t + dt + + # z_t is now the final complex state at t=1.0 + + # ═══ PHASE 4: Measurement Collapse ═══ + # Extract Phase and Magnitude to measure uncertainty + final_real, final_imag = z_t.chunk(2, dim=-1) # (N, D), (N, D) + + # Phase = atan2(Imag, Real). Represents conceptual angle/ambiguity. + phase = torch.atan2(final_imag, final_real + 1e-8) # (N, D) + + # Uncertainty is based on phase complexity. High variance = high uncertainty. + collapse_confidence = torch.sigmoid(self.uncertainty_gate(phase)) # (N, 1) + + # Collapse the Complex wavefunction back into Real-valued Action (Logits) + collapsed_logits = self.measurement_proj(z_t) # (N, out_dim) + + # Apply confidence: if uncertainty is high, the model's absolute amplitude drops + final_quantum_logic = collapsed_logits * collapse_confidence + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * final_quantum_logic + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetOmniversalExpert(nn.Module): + """ + Backbone-compatible model with OmniversalQuantumExpertHead (v24). + """ + def __init__( + self, + ode_steps: int = 5, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(OmniversalQuantumExpertHead( + 256, 10, + ode_steps=ode_steps, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class CellularAutomataGrid(nn.Module): + """ + Neural Cellular Automata grid. Computes localized cell evolution instead of dense layers. + """ + def __init__(self, dim: int, grid_size: int = 4, ca_steps: int = 3, dropout: float = 0.1): + super().__init__() + self.dim = dim + self.grid_size = grid_size + self.ca_steps = ca_steps + + # State update network (operates on 3x3 local neighborhood) + self.update_net = nn.Sequential( + nn.Conv2d(dim, dim * 2, kernel_size=3, padding=1), + nn.GELU(), + nn.Dropout2d(dropout), + nn.Conv2d(dim * 2, dim, kernel_size=1) + ) + self.stochastic_gate = nn.Conv2d(dim, 1, kernel_size=1) + + def forward(self, x: torch.Tensor): + # x: (N, dim) + N = x.shape[0] + # Reshape flat vector into a 2D grid + grid = x.view(N, self.dim, self.grid_size, self.grid_size) + + for _ in range(self.ca_steps): + # NCA Update: compute delta + dx = self.update_net(grid) + + # Stochastic update gate (allows cells to preserve state or change) + update_prob = torch.sigmoid(self.stochastic_gate(grid)) + # In training, we use soft updates. In true NCA this would be stochastic mask. + grid = grid + update_prob * dx + + return grid.view(N, -1) + + +class FractalGenesisExpertHead(nn.Module): + """ + Generation v25: Fractal Genesis Expert + (Hyperbolic Space + Fractal Graph + Cellular Automata + Self-Assembly) + + Transforms reasoning from a flat vector sequence into a dynamically + spawning hyperbolic computational fractal tree. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + fractal_depth: int = 2, + atoms_count: int = 16, + grid_size: int = 4, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.fractal_depth = fractal_depth + self.atoms_count = atoms_count + self.grid_size = grid_size + + # Base Euclidean projection (residual fallback) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Euclidean Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # ═══ 1. Hyperbolic Projection ═══ + # Map Euclidean features to the Poincaré Disk manifold parameters + self.hyperbolic_proj = nn.Linear(in_dim, in_dim) + self.curvature = nn.Parameter(torch.tensor(1.0)) # Learned negative curvature c + + # ═══ 2. Semantic Self-Assembly ═══ + # Pool of 'Logic Atoms' that dynamically construct the fractal node circuits + self.logic_atoms = nn.Parameter(torch.randn(atoms_count, in_dim)) + self.assembly_net = nn.Sequential( + nn.Linear(in_dim * 2, in_dim), + nn.GELU(), + nn.Linear(in_dim, in_dim) + ) + + # ═══ 3. Fractal Graph Recursion & CA ═══ + # Each level of the fractal tree has its own CA grid evolver + # and a mechanism to spawn K=2 child nodes. + self.fractal_evolvers = nn.ModuleList([ + CellularAutomataGrid( + dim=in_dim // (grid_size**2), + grid_size=grid_size, + dropout=dropout + ) for _ in range(fractal_depth) + ]) + + # Node spawn projections: splits a parent node into 2 children + self.node_spawners = nn.ModuleList([ + nn.Linear(in_dim, in_dim * 2) for _ in range(fractal_depth - 1) + ]) + + # ═══ 4. Resonance Collapse Aggregation ═══ + # Collapses 2 children back into a parent via phase resonance + self.resonance_combiners = nn.ModuleList([ + nn.Linear(in_dim * 2, in_dim) for _ in range(fractal_depth - 1) + ]) + + self.final_projection = nn.Sequential( + nn.Linear(in_dim, 256), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(256, out_dim) + ) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.hyperbolic_proj.weight, a=math.sqrt(5)) + nn.init.zeros_(self.hyperbolic_proj.bias) + + std = 1.0 / math.sqrt(self.in_dim) + nn.init.uniform_(self.logic_atoms, -std, std) + + for m in self.assembly_net: + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + for m in self.node_spawners: + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + nn.init.zeros_(m.bias) + + for m in self.resonance_combiners: + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + nn.init.zeros_(m.bias) + + for m in self.final_projection: + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def exp_map_zero(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + """Hyperbolic exponential map from origin to Poincaré disk.""" + sqrt_c = torch.clamp(c, min=1e-5).sqrt() + norm_x = torch.clamp(torch.norm(x, dim=-1, keepdim=True), min=1e-5) + return torch.tanh(sqrt_c * norm_x) * x / (sqrt_c * norm_x) + + def poincare_dist(self, x: torch.Tensor, y: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + """Hyperbolic distance between vectors x and y in Poincaré disk.""" + sqrt_c = torch.clamp(c, min=1e-5).sqrt() + norm_x_sq = torch.norm(x, dim=-1, keepdim=True)**2 + norm_y_sq = torch.norm(y, dim=-1, keepdim=True)**2 + # Mobius addition squared norm + num = torch.norm(x - y, dim=-1, keepdim=True)**2 + den = (1 - c * norm_x_sq) * (1 - c * norm_y_sq) + arg = 1 + 2 * c * num / torch.clamp(den, min=1e-5) + # Using arcosh trick + dist = (2 / sqrt_c) * torch.log(arg + torch.sqrt(arg**2 - 1 + 1e-6)) + return dist + + def assemble_node(self, node_state: torch.Tensor) -> torch.Tensor: + """Dynamically binds logic atoms to the node state using hyperbolic distance.""" + # Map node to hyperbolic space + c = torch.clamp(F.softplus(self.curvature), min=1e-4) # Ensure c > 0 + h_node = self.exp_map_zero(self.hyperbolic_proj(node_state), c) # (N, D) + h_atoms = self.exp_map_zero(self.logic_atoms, c) # (A, D) + + # Compute pairwise hyperbolic distances: (N, 1, D) and (1, A, D) -> (N, A) + dists = self.poincare_dist(h_node.unsqueeze(1), h_atoms.unsqueeze(0), c).squeeze(-1) + + # Attention over atoms (closer = higher weight) + atom_weights = torch.softmax(-dists, dim=-1) # (N, A) + + # Assemble bounded circuit state + bound_atoms = torch.matmul(atom_weights, self.logic_atoms) # (N, D) + + # Feed back to Euclidean space for processing + assembled = self.assembly_net(torch.cat([node_state, bound_atoms], dim=-1)) # (N, D) + return node_state + 0.1 * assembled + + def recursive_fractal_pass(self, node_state: torch.Tensor, depth: int) -> torch.Tensor: + """ + Recursively spawns children, evolves them via CA, and collapses them via resonance. + Depth 0 = deepest leaf nodes. + depth = self.fractal_depth - 1 is the Root node. + """ + # 1. Self-Assembly: bind logic atoms to current node state + assembled_node = self.assemble_node(node_state) + + # 2. Base case (leaf nodes) + if depth == 0: + # Evolve via Cellular Automata and return + return self.fractal_evolvers[0](assembled_node) + + # 3. Recursive step (internal nodes) + # Spawn two children from current node + children_raw = self.node_spawners[depth-1](assembled_node) # (N, D*2) + child_left, child_right = children_raw.chunk(2, dim=-1) # (N, D), (N, D) + + # Recurse on children + evolved_left = self.recursive_fractal_pass(child_left, depth - 1) + evolved_right = self.recursive_fractal_pass(child_right, depth - 1) + + # 4. Resonance Collapse + # Combine children back into parent. Instead of simple addition, we use + # a resonance combiner to simulate constructive/destructive logic interference + combined = self.resonance_combiners[depth-1](torch.cat([evolved_left, evolved_right], dim=-1)) + + # Finally evolve the combined parent node via its CA grid + final_node = self.fractal_evolvers[depth](combined + assembled_node) + return final_node + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Shared Euclidean Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # ═══ FRACTAL GENESIS ENGINE ═══ + # Launch the recursive fractal graph processing from the root + root_state = x_flat + final_euclidean_state = self.recursive_fractal_pass(root_state, depth=self.fractal_depth - 1) + + # Project back to logits + fractal_logic = self.final_projection(final_euclidean_state) + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * fractal_logic + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetFractalExpert(nn.Module): + """ + Backbone-compatible model with FractalGenesisExpertHead (v25). + """ + def __init__( + self, + fractal_depth: int = 2, + atoms_count: int = 16, + grid_size: int = 4, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(FractalGenesisExpertHead( + 256, 10, + fractal_depth=fractal_depth, + atoms_count=atoms_count, + grid_size=grid_size, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class SurrogateHeaviside(torch.autograd.Function): + """ + Surrogate gradient for the non-differentiable spiking Heaviside step function. + Uses a fast sigmoid approximation for the backward pass. + """ + @staticmethod + def forward(ctx, input, threshold, alpha): + ctx.save_for_backward(input, threshold, alpha) + # Forward pass: Strict binary spike (Heaviside step) + return (input >= threshold).float() + + @staticmethod + def backward(ctx, grad_output): + input, threshold, alpha = ctx.saved_tensors + # Backward pass: Fast sigmoid surrogate gradient + sgax = (input - threshold) * alpha + S_dtps = torch.exp(-torch.abs(sgax)) / (1 + torch.exp(-torch.abs(sgax)))**2 + + grad_input = grad_output * alpha * S_dtps + grad_threshold = (grad_output * (-alpha) * S_dtps).sum() + grad_alpha = (grad_output * (input - threshold) * S_dtps).sum() + + return grad_input, grad_threshold, grad_alpha + +def surrogate_spike(input, threshold, alpha=torch.tensor(2.0)): + return SurrogateHeaviside.apply(input, threshold, alpha) + +class LIFNeuron(nn.Module): + """ + Leaky Integrate-and-Fire (LIF) Spiking Neuron. + Maintains a continuous membrane potential that decays over time and spikes when crossing a threshold. + """ + def __init__(self, in_features: int, decay: float = 0.9, threshold: float = 1.0): + super().__init__() + self.in_features = in_features + self.decay = decay + self.threshold = nn.Parameter(torch.tensor(threshold)) + self.alpha = nn.Parameter(torch.tensor(2.0)) # Surrogate gradient sharpness + + def forward(self, x: torch.Tensor, v_mem: torch.Tensor): + # Update membrane potential: decay previous + integrate new input + v_mem = v_mem * self.decay + x + + # Fire spike if v_mem > threshold + spike = surrogate_spike(v_mem, self.threshold, self.alpha) + + # Soft reset: subtract threshold from membrane if fired + v_mem = v_mem - spike * self.threshold + + return spike, v_mem + +class LiquidSynapseMPS(nn.Module): + """ + Combines Tensor Networks (1D Matrix Product State approximation) + with Liquid Neural Network synaptic dynamics (State-dependent differential scaling). + """ + def __init__(self, dim: int, rank: int = 16): + super().__init__() + self.dim = dim + self.rank = rank + + # Matrix Product State (MPS) core tensors (highly compressed representation of dense weights) + # Instead of dim x dim, we use (dim x rank) and (rank x dim) to emulate a tensor train + self.core_1 = nn.Parameter(torch.randn(dim, rank) / math.sqrt(dim)) + self.core_2 = nn.Parameter(torch.randn(rank, dim) / math.sqrt(rank)) + + # Liquid Synaptic Time Constants (LNN mechanics) + self.tau = nn.Parameter(torch.ones(dim)) + # Non-linear input processor for the liquid equation + self.liquid_gate = nn.Linear(dim, dim) + + def forward(self, x: torch.Tensor, liquid_state: torch.Tensor, dt: float = 0.1): + # 1. Tensor Network (MPS) transformation + # x is (..., dim) + hidden = torch.matmul(x, self.core_1) # (..., rank) + mps_transform = torch.matmul(hidden, self.core_2) # (..., dim) + + # 2. Liquid Synapse update: dx/dt = -x/tau + S(x, I) + sensory_input = torch.sigmoid(self.liquid_gate(mps_transform)) + + # Differential equation Euler step + d_state = -liquid_state / torch.clamp(F.softplus(self.tau), min=1e-3) + sensory_input + new_state = liquid_state + d_state * dt + + return new_state + + +class LiquidSpikingTensorExpertHead(nn.Module): + """ + Generation v26: Liquid Spiking Tensor Expert + (Spiking Neural Networks + Liquid Synapses + Test-Time Training + Tensor Networks) + + A bio-plausible, highly compressed architectured adapting dynamically during inference. + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + spiking_steps: int = 4, + mps_rank: int = 16, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.spiking_steps = spiking_steps + + # Base Euclidean projection + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Euclidean Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # ═══ 1. Test-Time Training (TTT) Fast-Weight Module ═══ + # Internal LoRA adapter that updates *during the forward pass* + self.ttt_lora_A = nn.Parameter(torch.randn(in_dim, 8) / math.sqrt(in_dim)) + self.ttt_lora_B = nn.Parameter(torch.zeros(8, out_dim)) + # Self-supervised auxiliary task (Reconstruct input from bottleneck) + self.ttt_decoder = nn.Linear(8, in_dim) + + # ═══ 2. Liquid Synapse / Tensor Network Core ═══ + self.liquid_synapse = LiquidSynapseMPS(dim=in_dim, rank=mps_rank) + + # ═══ 3. Spiking Neural Mechanics (LIF Neurons) ═══ + self.lif_neuron = LIFNeuron(in_features=in_dim) + + # ═══ 4. Neuro-Symbolic Thresholding ═══ + self.spike_decoder = nn.Linear(in_dim, out_dim) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.normal_(self.ttt_lora_B, std=0.01) + self.lif_neuron.threshold.data.fill_(0.1) + + nn.init.kaiming_uniform_(self.ttt_decoder.weight, a=math.sqrt(5)) + nn.init.zeros_(self.ttt_decoder.bias) + + nn.init.kaiming_uniform_(self.spike_decoder.weight, a=math.sqrt(5)) + nn.init.zeros_(self.spike_decoder.bias) + + def test_time_training_step(self, x: torch.Tensor) -> torch.Tensor: + """ + Runs one step of gradient descent on the internal fast-weights (LoRA matrices) + using a self-supervised autoencoding objective, strictly separated from global backprop. + """ + # Detach x from the global computation graph so inner loop grads don't leak + x_inner = x.detach().requires_grad_(True) + + # Self-supervised forward: Encode using shared representation, decode using aux task + latent = torch.matmul(x_inner, self.ttt_lora_A) + reconstruction = self.ttt_decoder(latent) + + # Loss: Mean Squared Error representation + loss = F.mse_loss(reconstruction, x_inner) + + # Calculate gradients for the Lora A encoder and aux decoder + grads = torch.autograd.grad( + loss, + [self.ttt_lora_A, self.ttt_decoder.weight, self.ttt_decoder.bias], + create_graph=self.training, + retain_graph=True + ) + + # TTT SGD Step (lr = 0.01). Update only the shared encoder fast-weights. + fast_lora_A = self.ttt_lora_A - 0.01 * grads[0] + + # Compute TTT contribution for this batch using adapted encoder and frozen task-decoder + ttt_output = torch.matmul(torch.matmul(x, fast_lora_A), self.ttt_lora_B) # (N, out_dim) + return ttt_output + + def forward(self, x): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Shared Euclidean Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + # ═══ 1. Test-Time Training Update ═══ + ttt_logits = self.test_time_training_step(x_flat) + + # ═══ 2. Spiking Liquid Tensor Engine ═══ + # Initialize Membrane Voltage and Liquid State for the Spiking Loop + v_mem = torch.zeros_like(x_flat) + liquid_state = torch.zeros_like(x_flat) + spike_train_accum = torch.zeros_like(x_flat) + + for _ in range(self.spiking_steps): + # Liquid Synapse (MPS) update + liquid_state = self.liquid_synapse(x_flat, liquid_state) + + # Feed continuously changing liquid state into LIF Neuron + spike, v_mem = self.lif_neuron(liquid_state, v_mem) + + # Accumulate discrete spikes + spike_train_accum = spike_train_accum + spike + + # ═══ 4. Neuro-Symbolic Thresholding ═══ + # Average firing rate decoded back into continuous real logits + firing_rate = spike_train_accum / self.spiking_steps + spiking_logic = self.spike_decoder(firing_rate) + + output = ( + base_logits + + self.shared_scale * shared_out + + ttt_logits + + (self.alpha + 1e-4) * spiking_logic + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetLiquidSpikingExpert(nn.Module): + """ + Backbone-compatible model with LiquidSpikingTensorExpertHead (v26). + """ + def __init__( + self, + spiking_steps: int = 4, + mps_rank: int = 16, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(LiquidSpikingTensorExpertHead( + 256, 10, + spiking_steps=spiking_steps, + mps_rank=mps_rank, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ActiveInferenceExpertHead(nn.Module): + """ + v27 Active Inference Expert Head based on the Free Energy Principle. + It introduces a generative World Model and minimizes Variational Free Energy + at test time before routing to experts. + """ + def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 6, + inference_steps: int = 3, lr: float = 0.1, dropout: float = 0.1): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + self.inference_steps = inference_steps + self.lr = lr + + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + self.world_model = nn.Linear(in_dim, in_dim, bias=True) + self.world_prior = nn.Linear(in_dim, in_dim, bias=False) + self.precision_gate = nn.Linear(in_dim, n_experts * 2, bias=True) + + self.experts_up = nn.ModuleList([nn.Linear(in_dim, 1024, bias=False) for _ in range(n_experts)]) + self.experts_down = nn.ModuleList([nn.Linear(1024, out_dim, bias=False) for _ in range(n_experts)]) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.zeros_(self.bias) + nn.init.normal_(self.world_model.weight, std=0.01) + nn.init.zeros_(self.world_model.bias) + nn.init.normal_(self.world_prior.weight, std=0.01) + nn.init.normal_(self.precision_gate.weight, std=0.01) + nn.init.zeros_(self.precision_gate.bias) + for e in self.experts_up: + nn.init.normal_(e.weight, std=0.01) + for e in self.experts_down: + nn.init.zeros_(e.weight) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + original_shape = x.shape + x_flat = x.view(-1, self.in_dim) + + # Test-time inner loop optimization + with torch.enable_grad(): + z = x_flat.clone().detach().requires_grad_(True) + for _ in range(self.inference_steps): + predicted_state = F.silu(self.world_model(z)) + surprise = F.mse_loss(predicted_state, x_flat, reduction='sum') + complexity = F.mse_loss(z, self.world_prior(x_flat), reduction='sum') + free_energy = surprise + 0.1 * complexity + + grad_z = torch.autograd.grad(free_energy, z, create_graph=self.training)[0] + z = z - self.lr * grad_z + + z_final = z + gating_out = self.precision_gate(z_final) + mu, logvar = gating_out.chunk(2, dim=-1) + precision = torch.exp(-logvar) + + gate_scores = torch.softmax(mu * precision, dim=-1) + + expert_outputs = [] + for i in range(self.n_experts): + h = F.silu(self.experts_up[i](z_final)) + h = self.dropout(h) + out = self.experts_down[i](h) + expert_outputs.append(out) + + expert_stack = torch.stack(expert_outputs, dim=-2) + routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=-2) + + output = base_logits.view(-1, self.out_dim) + self.alpha * routed_out + return output.view(*original_shape[:-1], self.out_dim) + + +class ChampionNetActiveInferenceExpert(nn.Module): + def __init__(self, inference_steps: int = 3, lr: float = 0.1, dropout: float = 0.1): + super().__init__() + from run import ChampionNet + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ActiveInferenceExpertHead( + 256, 10, + inference_steps=inference_steps, + lr=lr, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TestTimeDifferentialExpertHead(nn.Module): + """ + Experimental benchmark-focused head that combines: + - differential attention to suppress distractor context + - bounded hyper-connections for dynamic residual mixing + - a recurrent test-time refinement loop with energy/confidence halting + + The implementation is intentionally additive and keeps `weight` / `bias` + keys so warm-starting from a base ChampionNet checkpoint remains safe. + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + n_heads: int = 4, + n_experts: int = 6, + expert_hidden: int = 768, + max_thinking_steps: int = 4, + halt_threshold: float = 0.72, + energy_threshold: float = 0.20, + dropout: float = 0.1, + ): + super().__init__() + if in_dim % n_heads != 0: + raise ValueError(f"in_dim={in_dim} must be divisible by n_heads={n_heads}") + + self.in_dim = in_dim + self.out_dim = out_dim + self.n_heads = n_heads + self.n_experts = n_experts + self.head_dim = in_dim // n_heads + self.max_thinking_steps = int(max(1, max_thinking_steps)) + self.min_thinking_steps = 1 + self.halt_threshold = float(halt_threshold) + self.energy_threshold = float(energy_threshold) + + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + self.attn_norm = nn.LayerNorm(in_dim) + self.diff_q_pos = nn.Linear(in_dim, in_dim, bias=False) + self.diff_k_pos = nn.Linear(in_dim, in_dim, bias=False) + self.diff_q_neg = nn.Linear(in_dim, in_dim, bias=False) + self.diff_k_neg = nn.Linear(in_dim, in_dim, bias=False) + self.diff_v = nn.Linear(in_dim, in_dim, bias=False) + self.diff_out = nn.Linear(in_dim, in_dim, bias=False) + self.diff_lambda = nn.Parameter(torch.tensor(-2.0)) + + self.shared_norm = nn.LayerNorm(in_dim) + self.shared_up = nn.Linear(in_dim, expert_hidden, bias=False) + self.shared_down = nn.Linear(expert_hidden, in_dim, bias=False) + self.shared_scale = nn.Parameter(torch.tensor(0.0)) + + self.predictive_norm = nn.LayerNorm(in_dim) + self.predictive_prior = nn.Linear(in_dim, in_dim, bias=False) + self.error_out = nn.Linear(in_dim, in_dim, bias=False) + + self.hyper_diag = nn.Linear(in_dim, in_dim) + self.hyper_lowrank_in = nn.Linear(in_dim, max(32, in_dim // 2), bias=False) + self.hyper_lowrank_out = nn.Linear(max(32, in_dim // 2), in_dim, bias=False) + self.hyper_gate = nn.Linear(in_dim, in_dim) + + self.gate = nn.Linear(in_dim, n_experts, bias=False) + self.experts_up = nn.ModuleList([nn.Linear(in_dim, expert_hidden, bias=False) for _ in range(n_experts)]) + self.experts_down = nn.ModuleList([nn.Linear(expert_hidden, in_dim, bias=False) for _ in range(n_experts)]) + self.expert_norms = nn.ModuleList([nn.LayerNorm(in_dim) for _ in range(n_experts)]) + + self.delta_out = nn.Linear(in_dim, out_dim, bias=False) + self.reason_out = nn.Linear(in_dim, out_dim, bias=False) + self.energy_head = nn.Linear(in_dim, 1) + self.halt_gate = nn.Linear(in_dim, 1) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.beta = nn.Parameter(torch.tensor(0.0)) + self.gamma = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + self.last_thinking_steps: int = 0 + self.last_energy: float = 0.0 + self.last_confidence: float = 0.0 + self.last_update_norm: float = 0.0 + self.last_attention_gap: float = 0.0 + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5) + nn.init.zeros_(self.bias) + + for linear in ( + self.diff_q_pos, + self.diff_k_pos, + self.diff_q_neg, + self.diff_k_neg, + self.diff_v, + self.predictive_prior, + self.hyper_lowrank_in, + self.gate, + ): + nn.init.normal_(linear.weight, mean=0.0, std=0.02) + + for linear in ( + self.diff_out, + self.shared_up, + self.hyper_lowrank_out, + self.error_out, + ): + nn.init.normal_(linear.weight, mean=0.0, std=0.015) + + for linear in ( + self.shared_down, + ): + nn.init.zeros_(linear.weight) + nn.init.normal_(self.delta_out.weight, mean=0.0, std=0.01) + nn.init.normal_(self.reason_out.weight, mean=0.0, std=0.01) + + nn.init.zeros_(self.hyper_diag.weight) + nn.init.zeros_(self.hyper_diag.bias) + nn.init.zeros_(self.hyper_gate.weight) + nn.init.zeros_(self.hyper_gate.bias) + nn.init.normal_(self.energy_head.weight, mean=0.0, std=0.01) + nn.init.zeros_(self.energy_head.bias) + nn.init.normal_(self.halt_gate.weight, mean=0.0, std=0.01) + nn.init.zeros_(self.halt_gate.bias) + + for up, down in zip(self.experts_up, self.experts_down): + nn.init.normal_(up.weight, mean=0.0, std=0.02) + nn.init.normal_(down.weight, mean=0.0, std=0.01) + + def _ensure_sequence(self, x: torch.Tensor) -> Tuple[torch.Tensor, bool]: + if x.ndim == 2: + return x.unsqueeze(1), True + if x.ndim == 3: + return x, False + raise ValueError(f"Expected rank-2 or rank-3 tensor, got shape={tuple(x.shape)}") + + def _reshape_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch, steps, _ = tensor.shape + return tensor.view(batch, steps, self.n_heads, self.head_dim).transpose(1, 2) + + def compute_differential_attention_maps( + self, x: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sequence_x, _ = self._ensure_sequence(x) + normalized = self.attn_norm(sequence_x) + + q_pos = self._reshape_heads(self.diff_q_pos(normalized)) + k_pos = self._reshape_heads(self.diff_k_pos(normalized)) + q_neg = self._reshape_heads(self.diff_q_neg(normalized)) + k_neg = self._reshape_heads(self.diff_k_neg(normalized)) + + scale = self.head_dim ** -0.5 + pos_scores = torch.matmul(q_pos, k_pos.transpose(-1, -2)) * scale + neg_scores = torch.matmul(q_neg, k_neg.transpose(-1, -2)) * scale + + pos_attn = torch.softmax(pos_scores, dim=-1) + neg_attn = torch.softmax(neg_scores, dim=-1) + + diff_ratio = torch.sigmoid(self.diff_lambda) + diff_attn = pos_attn - diff_ratio * neg_attn + diff_attn = diff_attn / diff_attn.abs().sum(dim=-1, keepdim=True).clamp_min(1.0) + return pos_attn, neg_attn, diff_attn + + def _differential_context(self, state: torch.Tensor) -> Tuple[torch.Tensor, dict]: + normalized = self.attn_norm(state) + values = self._reshape_heads(self.diff_v(normalized)) + pos_attn, neg_attn, diff_attn = self.compute_differential_attention_maps(state) + + context = torch.matmul(diff_attn, values) + context = context.transpose(1, 2).contiguous().view_as(state) + context = self.diff_out(context) + stats = { + "pos_attn": pos_attn, + "neg_attn": neg_attn, + "diff_attn": diff_attn, + } + return context, stats + + def _expert_reason(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + flat = state.reshape(-1, self.in_dim) + gate_scores = torch.softmax(self.gate(flat), dim=-1) + + expert_outputs = [] + for expert_norm, expert_up, expert_down in zip(self.expert_norms, self.experts_up, self.experts_down): + expert_state = expert_norm(flat) + expert_hidden = self.dropout(F.silu(expert_up(expert_state))) + expert_outputs.append(expert_down(expert_hidden)) + + expert_stack = torch.stack(expert_outputs, dim=1) + mixed = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + return mixed.view_as(state), gate_scores.view(*state.shape[:-1], self.n_experts) + + def _hyper_connect(self, state: torch.Tensor, proposal: torch.Tensor) -> torch.Tensor: + normalized = self.predictive_norm(state) + diag = 1.0 + 0.15 * torch.tanh(self.hyper_diag(normalized)) + low_rank = self.hyper_lowrank_out(self.dropout(F.silu(self.hyper_lowrank_in(normalized)))) + low_rank = 0.10 * torch.tanh(low_rank) + blend = torch.sigmoid(self.hyper_gate(normalized)) + return state * diag + blend * proposal + low_rank + + def forward(self, x: torch.Tensor) -> torch.Tensor: + state, squeezed = self._ensure_sequence(x) + base_logits = F.linear(state, self.weight, self.bias) + + self.last_thinking_steps = 0 + self.last_energy = 0.0 + self.last_confidence = 0.0 + self.last_update_norm = 0.0 + self.last_attention_gap = 0.0 + + last_logits = base_logits + for step in range(self.max_thinking_steps): + context, attn_stats = self._differential_context(state) + shared_update = self.shared_down(self.dropout(F.silu(self.shared_up(self.shared_norm(state))))) + expert_update, _ = self._expert_reason(state) + + predictive_prior = self.predictive_prior(self.predictive_norm(state)) + prediction_error = context - predictive_prior + error_update = 0.25 * self.error_out(prediction_error) + + proposal = ( + context + + self.shared_scale * shared_update + + torch.tanh(self.beta) * expert_update + + error_update + ) + candidate = self._hyper_connect(state, proposal) + + delta_logits = self.delta_out(candidate) + reason_logits = self.reason_out(prediction_error) + last_logits = base_logits + self.alpha * delta_logits + torch.tanh(self.gamma) * reason_logits + + confidence = torch.softmax(last_logits, dim=-1).amax(dim=-1).mean() + energy = F.softplus(self.energy_head(prediction_error)).mean() + (candidate - state).pow(2).mean() + halt_score = torch.sigmoid(self.halt_gate(self.predictive_norm(candidate))).mean() + + self.last_thinking_steps = step + 1 + self.last_energy = float(energy.detach().cpu()) + self.last_confidence = float(confidence.detach().cpu()) + self.last_update_norm = float((candidate - state).norm(dim=-1).mean().detach().cpu()) + self.last_attention_gap = float( + (attn_stats["pos_attn"] - attn_stats["neg_attn"]).abs().mean().detach().cpu() + ) + + state = candidate + + if ( + not self.training + and step + 1 >= self.min_thinking_steps + and self.last_energy <= self.energy_threshold + and ( + self.last_confidence >= self.halt_threshold + or float(halt_score.detach().cpu()) >= self.halt_threshold + ) + ): + break + + return last_logits.squeeze(1) if squeezed else last_logits + + +class ChampionNetTestTimeDiffExpert(nn.Module): + """ + Backbone-compatible wrapper for TestTimeDifferentialExpertHead. + """ + + def __init__( + self, + max_thinking_steps: int = 4, + n_experts: int = 6, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append( + TestTimeDifferentialExpertHead( + 256, + 10, + n_experts=n_experts, + max_thinking_steps=max_thinking_steps, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class StateSpaceExpert(nn.Module): + """ + Continuous-Time State Space Model (inspired by S4/Mamba). + Approximates an LDS: h' = Ah + Bx, y = Ch + Dx. + A simple discrete Zero-Order Hold (ZOH) discretization is applied. + """ + def __init__(self, d_model: int, d_state: int = 64, dt_min: float = 0.001, dt_max: float = 0.1): + super().__init__() + self.d_model = d_model + self.d_state = d_state + + # State space parameters + self.A_log = nn.Parameter(torch.empty(d_model, d_state)) + self.B = nn.Parameter(torch.empty(d_model, d_state)) + self.C = nn.Parameter(torch.empty(d_model, d_state)) + self.D = nn.Parameter(torch.zeros(d_model)) + + self.dt_proj = nn.Linear(d_model, d_model) + self.dt_min = dt_min + self.dt_max = dt_max + self.reset_parameters() + + def reset_parameters(self): + # Initialize A to be stable (negative real parts) + nn.init.uniform_(self.A_log, -10.0, -1.0) + nn.init.normal_(self.B, std=0.02) + nn.init.normal_(self.C, std=0.02) + nn.init.zeros_(self.D) + nn.init.normal_(self.dt_proj.weight, std=0.01) + nn.init.constant_(self.dt_proj.bias, -3.0) + + def forward(self, x): + # x shape: (..., d_model) + # 1. Delta (time scale) + dt = F.softplus(self.dt_proj(x)) + dt = torch.clamp(dt, min=self.dt_min, max=self.dt_max) + + # Continuous A must be negative to be stable + A = -torch.exp(self.A_log) + + # 2. ZOH Discretization for A and B + # dA = exp(dt * A) + # dB = (exp(dt * A) - I) / A * B + + dA = torch.exp(dt.unsqueeze(-1) * A) + dB = (dA - 1.0) / (A + 1e-6) * self.B + + # Initialize sequence hidden state to zeros (assuming per-token routing execution, + # this acts like a depth-wise recurrent pass per token) + h = torch.zeros(x.shape[:-1] + (self.d_model, self.d_state), device=x.device, dtype=x.dtype) + + # Step the state 3 times recursively for depth without sequence + for _ in range(3): + h = dA * h + dB * x.unsqueeze(-1) + + y = (self.C * h).sum(dim=-1) + self.D * x + return F.silu(y) + + +class HolographicStateSpaceExpertHead(nn.Module): + """ + v28 Holographic State Space Expert Head. + 1. Holographic Memory: Binds discrete token features globally using FFT circular convolution. + 2. State Space Experts (S4-like): Handles mapping via continuous LDS. + 3. Dynamic Topology Routing: Chaining experts in an input-dependent sequence path. + """ + def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 4, dropout: float = 0.1): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Learned holographic associative memory matrix + self.holographic_memory = nn.Parameter(torch.empty(1, in_dim)) + self.holo_proj = nn.Linear(in_dim, in_dim, bias=False) + + self.ssm_experts = nn.ModuleList([StateSpaceExpert(in_dim) for _ in range(n_experts)]) + + # Adjacency topology router + # Outputs a (n_experts, n_experts) adjacency matrix per token + self.topology_router = nn.Linear(in_dim, n_experts * n_experts) + + self.out_proj = nn.Linear(in_dim, out_dim, bias=False) + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.zeros_(self.bias) + nn.init.normal_(self.holographic_memory, std=0.02) + nn.init.normal_(self.holo_proj.weight, std=0.01) + nn.init.zeros_(self.topology_router.weight) + nn.init.zeros_(self.topology_router.bias) + nn.init.zeros_(self.out_proj.weight) + + def _circular_convolution(self, a, b): + """HRR Binding via FFT.""" + # Convert to complex plane to perform cyclic convolution in frequency domain + a_f = torch.fft.rfft(a, dim=-1) + b_f = torch.fft.rfft(b, dim=-1) + c = torch.fft.irfft(a_f * b_f, n=a.shape[-1], dim=-1) + return c + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + + # 1. Holographic Memory Binding + # Bind the token state x with the learned persistent holographic_memory + x_bound = self._circular_convolution(x, self.holographic_memory.expand_as(x)) + x_holo = self.holo_proj(x_bound) + + # 2. Dynamic Topology Routing + # B x T x (N*N) -> B x T x N x N + adj_flat = self.topology_router(x_holo) + adj_logits = adj_flat.view(*x.shape[:-1], self.n_experts, self.n_experts) + # Convert to a soft permutation matrix via doubly stochastic approximation (Sinkhorn) + # For simplicity in testing, we use softmax over rows + adj = torch.softmax(adj_logits, dim=-1) + + # 3. State Space Execution Path + # Instead of parallel execution, we execute iteratively. + current_state = x_holo + accum_out = torch.zeros_like(x_holo) + + expert_outs = [] + for i in range(self.n_experts): + h = self.ssm_experts[i](current_state) + expert_outs.append(h) + expert_stack = torch.stack(expert_outs, dim=-2) # (..., N_experts, D) + + # Route outputs through adjacency. + # current_state = \sum_j adj[i, j] * expert_stack[j] + # In a generic formulation, we compute the aggregated state as a mixture. + # We do N passes, representing N hops in the dynamic graph. + for step in range(self.n_experts): + step_mix = torch.einsum('...ij,...jd->...id', adj, expert_stack) + # The next state input is the average mixture for the next nodes + current_state = step_mix.mean(dim=-2) + accum_out = accum_out + current_state + + # Re-evaluate experts via recurrence (dynamic routing unrolled) + expert_outs = [] + for i in range(self.n_experts): + h = self.ssm_experts[i](current_state) + expert_outs.append(h) + expert_stack = torch.stack(expert_outs, dim=-2) + + final_h = accum_out / self.n_experts + routed_logits = self.out_proj(self.dropout(final_h)) + + return base_logits + self.alpha * routed_logits + + +class ChampionNetHolographicStateSpaceExpert(nn.Module): + def __init__(self, dropout: float = 0.1): + super().__init__() + from run import ChampionNet + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(HolographicStateSpaceExpertHead( + 256, 10, + n_experts=4, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TemporalSSMAdapter(nn.Module): + """STAD-inspired temporal state-space adapter (Schirmer et al., arXiv 2407.12492). + Learns time-varying dynamics in hidden features and applies a corrective residual. + Uses a lightweight SSM: h' = Ah + Bx, correction = Ch + Dx. + """ + def __init__(self, d_model: int, d_state: int = 32): + super().__init__() + self.d_model = d_model + self.d_state = d_state + self.A = nn.Parameter(torch.empty(d_model, d_state)) + self.B = nn.Parameter(torch.empty(d_model, d_state)) + self.C = nn.Parameter(torch.empty(d_model, d_state)) + self.D = nn.Parameter(torch.zeros(d_model)) + self.gate = nn.Parameter(torch.tensor(0.0)) # starts at 0 for stable init + self.reset_parameters() + + def reset_parameters(self): + nn.init.uniform_(self.A, -5.0, -0.5) + nn.init.normal_(self.B, std=0.02) + nn.init.normal_(self.C, std=0.02) + nn.init.zeros_(self.D) + + def forward(self, x, h=None): + # x: (..., d_model) + # Use tanh(A) for state decay — keeps values in (-1, 1) with healthy gradients + A_decay = torch.tanh(self.A) + if h is None: + h = torch.zeros(x.shape[:-1] + (self.d_model, self.d_state), + device=x.device, dtype=x.dtype) + # Single step update + h_new = A_decay * h + self.B * x.unsqueeze(-1) + correction = (self.C * h_new).sum(dim=-1) + self.D * x + return self.gate * correction, h_new + + +class MixtureOfRouters(nn.Module): + """MoR-inspired multi-sub-router ensemble (Zhang et al., 2025). + K independent sub-routers score experts. A main router weights their decisions. + """ + def __init__(self, in_dim: int, n_experts: int, n_sub_routers: int = 3): + super().__init__() + self.n_sub_routers = n_sub_routers + self.n_experts = n_experts + self.sub_routers = nn.ModuleList([ + nn.Linear(in_dim, n_experts, bias=True) for _ in range(n_sub_routers) + ]) + self.main_router = nn.Linear(in_dim, n_sub_routers, bias=True) + self.reset_parameters() + + def reset_parameters(self): + for sr in self.sub_routers: + nn.init.normal_(sr.weight, std=0.01) + nn.init.zeros_(sr.bias) + nn.init.zeros_(self.main_router.weight) + nn.init.zeros_(self.main_router.bias) + + def forward(self, x): + # Each sub-router produces expert scores + sub_scores = torch.stack([sr(x) for sr in self.sub_routers], dim=-2) # (..., K, N) + # Main router produces sub-router weights + main_weights = torch.softmax(self.main_router(x), dim=-1) # (..., K) + # Weighted combination: (..., N) + fused_scores = torch.einsum('...k,...kn->...n', main_weights, sub_scores) + return fused_scores + + +class PaperFusionExpertHead(nn.Module): + """v29 Paper-Driven Expert Head combining three 2025 research techniques: + 1. MoR (Mixture of Routers) — Zhang et al., 2025: Multi-sub-router ensemble. + 2. RoMA (Routing Manifold Alignment) — Li et al., arXiv 2511.07419: Manifold + regularization aligning routing weights with input embeddings. + 3. STAD (Temporal SSM Adapter) — Schirmer et al., arXiv 2407.12492: SSM-based + test-time feature drift correction. + """ + def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 6, + n_sub_routers: int = 3, ssm_d_state: int = 32, dropout: float = 0.1): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.n_experts = n_experts + + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared expert (always active) + self.shared_up = nn.Linear(in_dim, 512, bias=False) + self.shared_down = nn.Linear(512, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(0.0)) + + # MoR: Mixture of Routers (Paper 1) + self.mor_router = MixtureOfRouters(in_dim, n_experts, n_sub_routers) + + # Routed experts with diverse activations + self.experts_up = nn.ModuleList([ + nn.Linear(in_dim, 1024, bias=False) for _ in range(n_experts) + ]) + self.experts_down = nn.ModuleList([ + nn.Linear(1024, out_dim, bias=False) for _ in range(n_experts) + ]) + self._expert_acts = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] + + # STAD: Temporal SSM Adapter (Paper 3) + self.ssm_adapter = TemporalSSMAdapter(in_dim, d_state=ssm_d_state) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # RoMA manifold alignment loss (Paper 2) — stored for training script access + self._manifold_loss = torch.tensor(0.0) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=5**0.5) + nn.init.zeros_(self.bias) + nn.init.normal_(self.shared_up.weight, std=0.01) + nn.init.zeros_(self.shared_down.weight) + for eu in self.experts_up: + nn.init.normal_(eu.weight, std=0.01) + for ed in self.experts_down: + nn.init.zeros_(ed.weight) + + def _compute_manifold_alignment_loss(self, x_flat, routing_weights): + """RoMA: Penalize routing dissimilarity between embedding-space neighbors. + Inputs should be (..., D) and (..., N_experts) respectively. + """ + if x_flat.shape[0] < 2: + return torch.tensor(0.0, device=x_flat.device) + # Pairwise cosine similarity of input embeddings + x_norm = F.normalize(x_flat, dim=-1) + embed_sim = torch.mm(x_norm, x_norm.t()) # (B, B) + # Pairwise cosine similarity of routing weights + r_norm = F.normalize(routing_weights, dim=-1) + route_sim = torch.mm(r_norm, r_norm.t()) # (B, B) + # MSE alignment loss + return F.mse_loss(route_sim, embed_sim.detach()) + + def forward(self, x): + base_logits = F.linear(x, self.weight, self.bias) + shape_prefix = x.shape[:-1] + x_flat = x.reshape(-1, self.in_dim) + + # Shared expert + shared_out = self.shared_norm(self.shared_down(F.silu(self.shared_up(x_flat)))) + + # STAD: Apply temporal SSM correction + ssm_correction, _ = self.ssm_adapter(x_flat) + x_corrected = x_flat + ssm_correction + + # MoR: Ensemble routing scores + gate_scores = torch.softmax(self.mor_router(x_corrected), dim=-1) # (B, N) + + # RoMA: Compute manifold alignment loss (training only) + if self.training: + self._manifold_loss = self._compute_manifold_alignment_loss(x_flat, gate_scores) + else: + self._manifold_loss = torch.tensor(0.0, device=x.device) + + # Expert execution + expert_outputs = [] + for i in range(self.n_experts): + act_fn = self._expert_acts[i % len(self._expert_acts)] + h = act_fn(self.experts_up[i](x_corrected)) + h = self.dropout(h) + expert_outputs.append(self.experts_down[i](h)) + + expert_stack = torch.stack(expert_outputs, dim=-2) # (B, N, out_dim) + routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=-2) # (B, out_dim) + + output = ( + base_logits.reshape(-1, self.out_dim) + + self.shared_scale * shared_out + + self.alpha * routed_out + ) + return output.view(*shape_prefix, self.out_dim) + + +class ChampionNetPaperFusionExpert(nn.Module): + """Backbone wrapper for v29 PaperFusionExpertHead.""" + def __init__(self, n_sub_routers: int = 3, dropout: float = 0.1): + super().__init__() + from run import ChampionNet + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(PaperFusionExpertHead( + 256, 10, + n_experts=6, + n_sub_routers=n_sub_routers, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class CognitiveLeapExpertHead(nn.Module): + """ + Generation v50: Cognitive Leap Expert + (Tiny Recursive Latent Reasoning + Adaptive Halting + Test-Time Compute Scaling) + + Research-grounded recursive reasoning head: + 1. TRM-style dual latents — a single tiny weight-tied core refines a + low-level scratchpad latent z_L and a high-level answer latent z_H + over outer/inner cycles (Jolicoeur-Martineau, arXiv:2510.04871). + 2. Recurrent-depth test-time compute — the core is weight-tied, so + `reasoning_cycles` can be raised at inference to spend more compute + on hard inputs without retraining (Geiping et al., arXiv:2502.05171). + 3. ACT-style differentiable halting — a halt head allocates probability + mass across cycles; the output is the halting-weighted mixture of + per-cycle decodes, so easy inputs effectively stop early + (HRM, arXiv:2506.21734). + 4. Latent-convergence regularizer — penalizes z_H drift between later + cycles, pushing the recursion toward a fixed point so extra test-time + cycles refine rather than wander (RD-VLA, arXiv:2602.07845). + 5. Hypersphere-normalized latents — each latent update is RMS-normalized + onto a learned-gain sphere for recursion stability (nGPT-inspired). + 6. Cycle conditioning — a learned per-cycle embedding tells the + weight-tied core where it is in the recursion, which recent TRM + analysis found load-bearing (arXiv:2512.11847). + 7. Deep improvement supervision — per-cycle decodes are cached during + training and `deep_supervision_loss(...)` applies progressively + weighted per-cycle losses so every cycle learns to improve on the + last (DIS, arXiv:2511.16886). + 8. Convergence early-exit — in eval mode `adaptive_compute=True` stops + cycling once the answer latent stops moving or the halting mass is + spent, skipping wasted cycles on easy inputs (arXiv:2506.02536). + """ + + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.latent_dim = latent_dim + self.core_dim = core_dim + self.n_cycles = n_cycles + self.inner_steps = inner_steps + self.max_cycles = max_cycles + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # Latent initializers: input -> scratchpad latent / answer latent + self.thought_init = nn.Linear(in_dim, latent_dim, bias=False) + self.answer_init = nn.Linear(in_dim, latent_dim, bias=False) + + # Tiny weight-tied recursive core: [x, z_L, z_H] -> delta z_L + self.recur_core = nn.Sequential( + nn.Linear(in_dim + 2 * latent_dim, core_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(core_dim, latent_dim), + ) + + # Outer answer update: [z_L, z_H] -> delta z_H + self.answer_update = nn.Sequential( + nn.Linear(2 * latent_dim, latent_dim), + nn.GELU(), + nn.Linear(latent_dim, latent_dim), + ) + + # Hypersphere gains for normalized latents (nGPT-style stability) + self.latent_gain = nn.Parameter(torch.ones(latent_dim)) + + # Cycle conditioning: tells the weight-tied core which cycle it is in + self.cycle_embed = nn.Embedding(max_cycles, latent_dim) + + # ACT-style halting head over the answer latent + self.halt_head = nn.Linear(latent_dim, 1) + + # Per-cycle decoder from answer latent to logits + self.decode_head = nn.Linear(latent_dim, out_dim, bias=False) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # Diagnostics / regularizers tracked for the training loop + self.register_buffer("last_ponder_cost", torch.tensor(0.0)) + self.register_buffer("last_consistency_loss", torch.tensor(0.0)) + self.register_buffer("last_cycles_used", torch.tensor(0.0)) + + # Per-cycle decode cache for deep improvement supervision (train only) + self._cycle_logits = [] + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.thought_init.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.answer_init.weight, a=math.sqrt(5)) + + for layer in list(self.recur_core) + list(self.answer_update): + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + # Last update layers start small so early recursion is near-identity + nn.init.normal_(self.recur_core[-1].weight, std=0.01) + nn.init.normal_(self.answer_update[-1].weight, std=0.01) + + nn.init.normal_(self.cycle_embed.weight, std=0.02) + + # Bias halting toward "continue" early in training so all cycles learn + nn.init.zeros_(self.halt_head.weight) + nn.init.constant_(self.halt_head.bias, -1.0) + + nn.init.normal_(self.decode_head.weight, std=0.01) + + def _sphere_norm(self, z): + # RMS-normalize onto the unit hypersphere, then apply learned gain + return F.normalize(z, p=2, dim=-1) * self.latent_gain + + def deep_supervision_loss(self, targets): + """Deep-improvement-supervision loss over cached per-cycle decodes. + + Practical DIS adaptation (arXiv:2511.16886): every cycle's decode is + supervised against the target with progressively increasing weight, + so each cycle learns to improve on the previous one rather than only + the final mixture receiving signal. Call after a training-mode + forward pass; returns a scalar loss. + """ + if not self._cycle_logits: + raise RuntimeError( + "deep_supervision_loss requires a training-mode forward pass first" + ) + targets_flat = targets.reshape(-1) + n = len(self._cycle_logits) + weights = torch.arange( + 1, n + 1, device=targets_flat.device, dtype=self._cycle_logits[0].dtype + ) + weights = weights / weights.sum() + loss = torch.zeros((), device=targets_flat.device, dtype=self._cycle_logits[0].dtype) + for w, logits in zip(weights, self._cycle_logits): + loss = loss + w * F.cross_entropy(logits, targets_flat) + return loss + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + ): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Global Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + n_cycles = int(reasoning_cycles) if reasoning_cycles is not None else self.n_cycles + n_cycles = max(1, n_cycles) + + # Initialize latents on the hypersphere + z_l = self._sphere_norm(self.thought_init(x_flat)) + z_h = self._sphere_norm(self.answer_init(x_flat)) + + cycle_logits = [] + halt_probs = [] + remaining = torch.ones(N, 1, device=x_flat.device, dtype=x_flat.dtype) + ponder = torch.zeros(N, 1, device=x_flat.device, dtype=x_flat.dtype) + consistency = torch.zeros((), device=x_flat.device, dtype=x_flat.dtype) + prev_z_h = None + cycles_used = n_cycles + + for t in range(n_cycles): + # Cycle conditioning: identity of the current cycle for the + # weight-tied core (clamped beyond the trained range) + cycle_idx = min(t, self.max_cycles - 1) + cycle_vec = self.cycle_embed.weight[cycle_idx].unsqueeze(0) + z_l_cond = z_l + cycle_vec + + # Inner scratchpad refinement (weight-tied tiny core) + for _ in range(self.inner_steps): + core_in = torch.cat([x_flat, z_l_cond, z_h], dim=-1) + z_l_cond = self._sphere_norm(z_l_cond + self.recur_core(core_in)) + z_l = z_l_cond + # Outer answer refinement + z_h = self._sphere_norm( + z_h + self.answer_update(torch.cat([z_l, z_h], dim=-1)) + ) + + cycle_logits.append(self.decode_head(z_h)) + + # Latent-convergence measurement (regularizer + early-exit signal) + if prev_z_h is not None: + delta = (z_h - prev_z_h).pow(2).mean() + consistency = consistency + delta + else: + delta = None + prev_z_h = z_h + + # Convergence early-exit: only at inference, never during training + converged = ( + adaptive_compute + and not self.training + and delta is not None + and (delta.item() < exit_tol or remaining.mean().item() < 0.05) + ) + + # Differentiable ACT halting: probability of stopping at cycle t + h_t = torch.sigmoid(self.halt_head(z_h)) # (N, 1) + if t == n_cycles - 1 or converged: + p_t = remaining # assign all leftover mass to the final cycle + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + cycles_used = t + 1 + break + p_t = remaining * h_t + remaining = remaining * (1.0 - h_t) + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + + self.last_cycles_used = torch.tensor( + float(cycles_used), device=x_flat.device + ) + if self.training: + self.last_ponder_cost = ponder.mean() + self.last_consistency_loss = consistency / max(1, cycles_used - 1) + self._cycle_logits = cycle_logits + else: + self._cycle_logits = [] + + # Halting-weighted mixture of per-cycle decodes + logits_stack = torch.stack(cycle_logits, dim=1) # (N, T, out_dim) + probs_stack = torch.stack(halt_probs, dim=1) # (N, T, 1) + recursive_out = (logits_stack * probs_stack).sum(dim=1) # (N, out_dim) + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * recursive_out + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetCognitiveLeapExpert(nn.Module): + """ + Backbone-compatible model with CognitiveLeapExpertHead (v50). + + forward() accepts an optional reasoning_cycles override so callers can + scale test-time compute without rebuilding the model. + """ + def __init__( + self, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(CognitiveLeapExpertHead( + 256, 10, + latent_dim=latent_dim, + core_dim=core_dim, + n_cycles=n_cycles, + inner_steps=inner_steps, + max_cycles=max_cycles, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + ): + for layer in self.layers: + if isinstance(layer, CognitiveLeapExpertHead): + x = layer( + x, + reasoning_cycles=reasoning_cycles, + adaptive_compute=adaptive_compute, + exit_tol=exit_tol, + ) + else: + x = layer(x) + return x + + def deep_supervision_loss(self, targets): + return self.layers[10].deep_supervision_loss(targets) + + +class LightweightMHA(nn.Module): + """ + Lightweight Multi-Head Attention for sequence of length 2 (z_l and z_h). + """ + def __init__(self, dim, num_heads=2, dropout=0.1): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=False) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(dropout) + + def forward(self, x): + B, T, C = x.shape + qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # (B, num_heads, T, head_dim) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + + x_out = (attn @ v).transpose(1, 2).reshape(B, T, C) + x_out = self.proj(x_out) + x_out = self.proj_drop(x_out) + return x_out + + +class CognitiveLeapUltraExpertHead(nn.Module): + """ + Generation v51: Cognitive Leap Ultra Expert + (Tiny Recursive Latent Reasoning with Cross-Attention + Recurrent MoE Core + Entropy Halting) + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + n_cores: int = 4, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.latent_dim = latent_dim + self.core_dim = core_dim + self.n_cycles = n_cycles + self.inner_steps = inner_steps + self.max_cycles = max_cycles + self.n_cores = n_cores + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # Latent initializers + self.thought_init = nn.Linear(in_dim, latent_dim, bias=False) + self.answer_init = nn.Linear(in_dim, latent_dim, bias=False) + + # Recurrent MoE Core + self.cores = nn.ModuleList([ + nn.Sequential( + nn.Linear(in_dim + 2 * latent_dim, core_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(core_dim, latent_dim), + ) for _ in range(n_cores) + ]) + self.core_router = nn.Linear(in_dim + 2 * latent_dim, n_cores) + + # Cross-Attention Refinement + self.cross_attn = LightweightMHA(latent_dim, num_heads=2, dropout=dropout) + + # Outer answer update + self.answer_update = nn.Sequential( + nn.Linear(2 * latent_dim, latent_dim), + nn.GELU(), + nn.Linear(latent_dim, latent_dim), + ) + + # Hypersphere gains + self.latent_gain = nn.Parameter(torch.ones(latent_dim)) + + # Cycle conditioning + self.cycle_embed = nn.Embedding(max_cycles, latent_dim) + + # ACT-style halting head + self.halt_head = nn.Linear(latent_dim, 1) + + # Per-cycle decoder + self.decode_head = nn.Linear(latent_dim, out_dim, bias=False) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # Diagnostics / regularizers + self.register_buffer("last_ponder_cost", torch.tensor(0.0)) + self.register_buffer("last_consistency_loss", torch.tensor(0.0)) + self.register_buffer("last_cycles_used", torch.tensor(0.0)) + self.register_buffer("last_gating_entropy", torch.tensor(0.0)) + self.register_buffer("last_prediction_streak", torch.tensor(0.0), persistent=False) + self.register_buffer("last_prediction_confidence_delta", torch.tensor(0.0), persistent=False) + self.register_buffer( + "last_prediction_topk_js_divergence", torch.tensor(0.0), persistent=False + ) + self.register_buffer( + "last_prediction_topk_js_divergence_max", torch.tensor(0.0), persistent=False + ) + self.last_exit_reason = "not_run" + + # Per-cycle decode cache for deep improvement supervision + self._cycle_logits = [] + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.thought_init.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.answer_init.weight, a=math.sqrt(5)) + + for core in self.cores: + for layer in core: + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + nn.init.normal_(core[-1].weight, std=0.01) + + nn.init.kaiming_uniform_(self.core_router.weight, a=math.sqrt(5)) + if self.core_router.bias is not None: + nn.init.zeros_(self.core_router.bias) + + nn.init.kaiming_uniform_(self.cross_attn.qkv.weight, a=math.sqrt(5)) + nn.init.normal_(self.cross_attn.proj.weight, std=0.01) + + for layer in self.answer_update: + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + nn.init.normal_(self.answer_update[-1].weight, std=0.01) + + nn.init.normal_(self.cycle_embed.weight, std=0.02) + + nn.init.zeros_(self.halt_head.weight) + nn.init.constant_(self.halt_head.bias, -1.0) + + nn.init.normal_(self.decode_head.weight, std=0.01) + + def _sphere_norm(self, z): + return F.normalize(z, p=2, dim=-1) * self.latent_gain + + def deep_supervision_loss(self, targets): + if not self._cycle_logits: + raise RuntimeError( + "deep_supervision_loss requires a training-mode forward pass first" + ) + targets_flat = targets.reshape(-1) + n = len(self._cycle_logits) + weights = torch.arange( + 1, n + 1, device=targets_flat.device, dtype=self._cycle_logits[0].dtype + ) + weights = weights / weights.sum() + loss = torch.zeros((), device=targets_flat.device, dtype=self._cycle_logits[0].dtype) + for w, logits in zip(weights, self._cycle_logits): + loss = loss + w * F.cross_entropy(logits, targets_flat) + return loss + + @staticmethod + def _topk_js_divergence( + previous_probs: torch.Tensor, + current_probs: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + """Return per-row top-k Jensen-Shannon drift with an other-mass bucket.""" + + k = min(max(1, int(top_k)), int(previous_probs.shape[-1])) + midpoint = 0.5 * (previous_probs + current_probs) + indices = midpoint.topk(k, dim=-1).indices + previous_top = previous_probs.gather(-1, indices) + current_top = current_probs.gather(-1, indices) + previous_other = (1.0 - previous_top.sum(dim=-1, keepdim=True)).clamp_min(0.0) + current_other = (1.0 - current_top.sum(dim=-1, keepdim=True)).clamp_min(0.0) + previous_compact = torch.cat([previous_top, previous_other], dim=-1) + current_compact = torch.cat([current_top, current_other], dim=-1) + compact_midpoint = 0.5 * (previous_compact + current_compact) + epsilon = torch.finfo(previous_compact.dtype).tiny + + def _kl(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + terms = left * ( + torch.log(left.clamp_min(epsilon)) + - torch.log(right.clamp_min(epsilon)) + ) + return terms.sum(dim=-1) + + return 0.5 * ( + _kl(previous_compact, compact_midpoint) + + _kl(current_compact, compact_midpoint) + ) + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + exit_entropy_threshold: float = 0.2, + prediction_stability_patience: int = 2, + prediction_stability_tol: float = 5e-3, + prediction_stability_top_k: int = 5, + ): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Global Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + n_cycles = int(reasoning_cycles) if reasoning_cycles is not None else self.n_cycles + n_cycles = max(1, n_cycles) + + # Initialize latents on the hypersphere + z_l = self._sphere_norm(self.thought_init(x_flat)) + z_h = self._sphere_norm(self.answer_init(x_flat)) + + cycle_logits = [] + halt_probs = [] + gate_probs_list = [] + remaining = torch.ones(N, 1, device=x_flat.device, dtype=x_flat.dtype) + ponder = torch.zeros(N, 1, device=x_flat.device, dtype=x_flat.dtype) + consistency = torch.zeros((), device=x_flat.device, dtype=x_flat.dtype) + prev_z_h = None + cycles_used = n_cycles + stability_patience = max(0, int(prediction_stability_patience)) + stability_tol = max(0.0, float(prediction_stability_tol)) + use_prediction_stability = adaptive_compute and not self.training and stability_patience > 0 + weighted_recursive = torch.zeros(N, self.out_dim, device=x_flat.device, dtype=x_flat.dtype) + previous_prediction = None + prediction_streak = torch.zeros(N, device=x_flat.device, dtype=torch.long) + stable_confidence_min = None + stable_confidence_max = None + prediction_confidence_delta = torch.zeros(N, device=x_flat.device, dtype=x_flat.dtype) + prediction_topk_js_divergence = torch.zeros( + N, device=x_flat.device, dtype=x_flat.dtype + ) + prediction_topk_js_divergence_max = torch.zeros_like( + prediction_topk_js_divergence + ) + previous_prediction_distribution = None + exit_reason = "max_cycles" + + for t in range(n_cycles): + cycle_idx = min(t, self.max_cycles - 1) + cycle_vec = self.cycle_embed.weight[cycle_idx].unsqueeze(0) + z_l_cond = z_l + cycle_vec + + # Inner scratchpad refinement (weight-tied recurrent MoE core) + for _ in range(self.inner_steps): + core_in = torch.cat([x_flat, z_l_cond, z_h], dim=-1) + gate_logits = self.core_router(core_in) + gate_probs = F.softmax(gate_logits, dim=-1) + gate_probs_list.append(gate_probs) + + # Route to recurrent cores + core_outs = torch.stack([core(core_in) for core in self.cores], dim=1) # (N, n_cores, latent_dim) + core_update = (core_outs * gate_probs.unsqueeze(-1)).sum(dim=1) # (N, latent_dim) + + z_l_cond = self._sphere_norm(z_l_cond + core_update) + z_l = z_l_cond + + # Attention-Based Cross-Latent Gating + z_seq = torch.stack([z_l, z_h], dim=1) # (N, 2, latent_dim) + z_seq = z_seq + self.cross_attn(z_seq) + z_l_refined = z_seq[:, 0, :] + z_h_refined = z_seq[:, 1, :] + + # Outer answer refinement + z_h = self._sphere_norm( + z_h_refined + self.answer_update(torch.cat([z_l_refined, z_h_refined], dim=-1)) + ) + + logits = self.decode_head(z_h) + cycle_logits.append(logits) + + # Prediction entropy + probs = F.softmax(logits, dim=-1) + entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1, keepdim=True) + + # Latent-convergence measurement + if prev_z_h is not None: + delta = (z_h - prev_z_h).pow(2).mean() + consistency = consistency + delta + else: + delta = None + prev_z_h = z_h + + # Treat the full prefix output as an internal verifier. If its top + # prediction and confidence stabilize, more recurrent work is not + # buying a meaningfully different answer for this request. + prediction_stable = False + if use_prediction_stability: + prefix_recursive = weighted_recursive + remaining * logits + prefix_output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * prefix_recursive + ) + prefix_probs = F.softmax(prefix_output, dim=-1) + if previous_prediction_distribution is not None: + prediction_topk_js_divergence = self._topk_js_divergence( + previous_prediction_distribution, + prefix_probs, + prediction_stability_top_k, + ) + prediction_topk_js_divergence_max = torch.maximum( + prediction_topk_js_divergence_max, + prediction_topk_js_divergence, + ) + previous_prediction_distribution = prefix_probs + prediction_confidence, prediction = prefix_probs.max(dim=-1) + if previous_prediction is None: + prediction_streak = torch.ones_like(prediction_streak) + stable_confidence_min = prediction_confidence + stable_confidence_max = prediction_confidence + else: + same_prediction = prediction.eq(previous_prediction) + prediction_streak = torch.where( + same_prediction, + prediction_streak + 1, + torch.ones_like(prediction_streak), + ) + stable_confidence_min = torch.where( + same_prediction, + torch.minimum(stable_confidence_min, prediction_confidence), + prediction_confidence, + ) + stable_confidence_max = torch.where( + same_prediction, + torch.maximum(stable_confidence_max, prediction_confidence), + prediction_confidence, + ) + prediction_confidence_delta = stable_confidence_max - stable_confidence_min + previous_prediction = prediction + stable_mask = ( + prediction_streak >= stability_patience + ) & ( + prediction_confidence_delta <= stability_tol + ) + prediction_stable = bool(stable_mask.all().item()) + + # Inference-only verifier and convergence exits. Prediction + # stability is checked first so diagnostics identify the new path + # when multiple exit criteria become true on the same cycle. + converged = False + if adaptive_compute and not self.training: + if prediction_stable: + converged = True + exit_reason = "prediction_stable" + elif delta is not None and delta.item() < exit_tol: + converged = True + exit_reason = "latent_converged" + elif entropy.mean().item() < exit_entropy_threshold: + converged = True + exit_reason = "low_entropy" + elif remaining.mean().item() < 0.05: + converged = True + exit_reason = "halt_mass" + + # Differentiable ACT halting + h_t = torch.sigmoid(self.halt_head(z_h)) # (N, 1) + if t == n_cycles - 1 or converged: + p_t = remaining + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + cycles_used = t + 1 + break + p_t = remaining * h_t + remaining = remaining * (1.0 - h_t) + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + if use_prediction_stability: + weighted_recursive = weighted_recursive + p_t * logits + + self.last_cycles_used = torch.tensor( + float(cycles_used), device=x_flat.device + ) + self.last_prediction_streak = prediction_streak.float().mean() + self.last_prediction_confidence_delta = prediction_confidence_delta.mean() + self.last_prediction_topk_js_divergence = ( + prediction_topk_js_divergence.mean().detach() + ) + self.last_prediction_topk_js_divergence_max = ( + prediction_topk_js_divergence_max.mean().detach() + ) + self.last_exit_reason = exit_reason + + # Calculate gating entropy for MoE cores + if gate_probs_list: + all_gate_probs = torch.stack(gate_probs_list, dim=0) # (Steps, N, n_cores) + mean_gate_probs = all_gate_probs.mean(dim=(0, 1)) # (n_cores,) + gating_entropy = -torch.sum(mean_gate_probs * torch.log(mean_gate_probs + 1e-9)) + else: + gating_entropy = torch.tensor(0.0, device=x_flat.device) + + if self.training: + self.last_ponder_cost = ponder.mean() + self.last_consistency_loss = consistency / max(1, cycles_used - 1) + self.last_gating_entropy = gating_entropy + self._cycle_logits = cycle_logits + else: + self.last_gating_entropy = gating_entropy + self._cycle_logits = [] + + # Halting-weighted mixture of per-cycle decodes + logits_stack = torch.stack(cycle_logits, dim=1) # (N, T, out_dim) + probs_stack = torch.stack(halt_probs, dim=1) # (N, T, 1) + recursive_out = (logits_stack * probs_stack).sum(dim=1) # (N, out_dim) + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * recursive_out + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetCognitiveLeapUltraExpert(nn.Module): + """ + Backbone-compatible model with CognitiveLeapUltraExpertHead (v51). + """ + def __init__( + self, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + n_cores: int = 4, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(CognitiveLeapUltraExpertHead( + 256, 10, + latent_dim=latent_dim, + core_dim=core_dim, + n_cycles=n_cycles, + inner_steps=inner_steps, + max_cycles=max_cycles, + dropout=dropout, + n_cores=n_cores, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + exit_entropy_threshold: float = 0.2, + prediction_stability_patience: int = 2, + prediction_stability_tol: float = 5e-3, + prediction_stability_top_k: int = 5, + ): + for layer in self.layers: + if isinstance(layer, CognitiveLeapUltraExpertHead): + x = layer( + x, + reasoning_cycles=reasoning_cycles, + adaptive_compute=adaptive_compute, + exit_tol=exit_tol, + exit_entropy_threshold=exit_entropy_threshold, + prediction_stability_patience=prediction_stability_patience, + prediction_stability_tol=prediction_stability_tol, + prediction_stability_top_k=prediction_stability_top_k, + ) + else: + x = layer(x) + return x + + def deep_supervision_loss(self, targets): + return self.layers[10].deep_supervision_loss(targets) + + +class ChampionNetHierarchicalExpert(nn.Module): + """ + Backbone-compatible model with HierarchicalMoEClassifierHead. + Uses the same layers 0-9 backbone from ChampionNet with the new + hierarchical MoE head at layer 10. + """ + def __init__( + self, + n_domain_groups: int = 2, + experts_per_group: int = 4, + top_k_domains: int = 1, + top_k_experts: int = 2, + lora_rank: int = 64, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(HierarchicalMoEClassifierHead( + 256, 10, + n_domain_groups=n_domain_groups, + experts_per_group=experts_per_group, + top_k_domains=top_k_domains, + top_k_experts=top_k_experts, + lora_rank=lora_rank, + dropout=dropout, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetDeepExpert(nn.Module): + """ + Backbone-compatible model with DeepExpertClassifierHead. + """ + def __init__(self, n_experts: int = 8, top_k: int = 2, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(DeepExpertClassifierHead(256, 10, n_experts=n_experts, top_k=top_k, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetExpertChoice(nn.Module): + """ + Backbone-compatible model with ExpertChoiceClassifierHead. + """ + def __init__(self, n_experts: int = 8, capacity_factor: float = 2.0, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ExpertChoiceClassifierHead(256, 10, n_experts=n_experts, capacity_factor=capacity_factor, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + return x + + +class ChampionNetSmarterExpert(nn.Module): + """ + Backbone-compatible model with SmarterExpertClassifierHead. + """ + def __init__(self, n_experts: int = 8, lora_rank: int = 32, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(SmarterExpertClassifierHead(256, 10, n_experts=n_experts, lora_rank=lora_rank, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + + +class ChampionNetThoughtExpert(nn.Module): + """ + Backbone-compatible model with ThoughtExpertClassifierHead. + """ + def __init__(self, n_experts: int = 8, reasoning_steps: int = 3, lora_rank: int = 32, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ThoughtExpertClassifierHead( + 256, 10, + n_experts=n_experts, + reasoning_steps=reasoning_steps, + lora_rank=lora_rank, + dropout=dropout + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetUltraExpert(nn.Module): + """ + Backbone-compatible model with GatedExpertClassifierHead. + """ + def __init__(self, n_experts: int = 6, top_k: int = 2, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(GatedExpertClassifierHead(256, 10, n_experts=n_experts, top_k=top_k, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetLarge(nn.Module): + """ + Backbone-compatible larger model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHead + """ + + def __init__(self, expansion_dim: int = 512, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(ExpandedClassifierHead(256, 10, expansion_dim=expansion_dim, dropout=dropout)) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetXL(nn.Module): + """ + Backbone-compatible xlarge model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHeadXL + """ + + def __init__(self, expansion_dim: int = 768, extra_expansion_dim: int = 1024, dropout: float = 0.1): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append( + ExpandedClassifierHeadXL( + 256, + 10, + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetXXL(nn.Module): + """ + Backbone-compatible xxlarge model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHeadXXL + """ + + def __init__( + self, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append( + ExpandedClassifierHeadXXL( + 256, + 10, + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + third_expansion_dim=third_expansion_dim, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetXXXL(nn.Module): + """ + Backbone-compatible xxxlarge model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHeadXXXL + """ + + def __init__( + self, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append( + ExpandedClassifierHeadXXXL( + 256, + 10, + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + third_expansion_dim=third_expansion_dim, + fourth_expansion_dim=fourth_expansion_dim, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetUltra(nn.Module): + """ + Backbone-compatible ultralarge model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHeadUltra + """ + + def __init__( + self, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + fifth_expansion_dim: int = 6144, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append( + ExpandedClassifierHeadUltra( + 256, + 10, + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + third_expansion_dim=third_expansion_dim, + fourth_expansion_dim=fourth_expansion_dim, + fifth_expansion_dim=fifth_expansion_dim, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class ChampionNetMega(nn.Module): + """ + Backbone-compatible megalarge model: + - keeps layers 0..9 and layer 11 from ChampionNet + - replaces layer 10 with ExpandedClassifierHeadMega + - adds GatedFFN after BitNetLinear for extra backbone capacity + """ + + def __init__( + self, + expansion_dim: int = 1024, + extra_expansion_dim: int = 2048, + third_expansion_dim: int = 3072, + fourth_expansion_dim: int = 4096, + fifth_expansion_dim: int = 6144, + sixth_expansion_dim: int = 8192, + dropout: float = 0.1, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + # Insert GatedFFN for extra backbone capacity (between BitNet and classifier). + layers.append(GatedFFN(d_model=256, d_inner=512)) + layers.append( + ExpandedClassifierHeadMega( + 256, + 10, + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + third_expansion_dim=third_expansion_dim, + fourth_expansion_dim=fourth_expansion_dim, + fifth_expansion_dim=fifth_expansion_dim, + sixth_expansion_dim=sixth_expansion_dim, + dropout=dropout, + ) + ) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +def build_model( + model_size: str = "base", + expansion_dim: int = 512, + dropout: float = 0.1, + extra_expansion_dim: Optional[int] = None, + third_expansion_dim: Optional[int] = None, + fourth_expansion_dim: Optional[int] = None, + fifth_expansion_dim: Optional[int] = None, + sixth_expansion_dim: Optional[int] = None, +) -> nn.Module: + """Construct a ChampionNet model variant by name. + + Factory function that dispatches to the appropriate backbone class + based on model_size. Expansion dims default to sensible values when + not provided. + + Args: + model_size: One of 'base', 'large', 'xlarge', 'xxlarge', 'xxxlarge', + 'ultralarge', 'megalarge', 'ultra_expert', 'hierarchical_expert', + 'deep_expert', 'expert_choice', 'smarter_expert', 'thought_expert', + 'recursive_expert'. + expansion_dim: Hidden size for the primary adapter branch. + dropout: Dropout rate for adapter layers. + extra_expansion_dim: Hidden size for the secondary adapter branch. + third_expansion_dim: Hidden size for the third adapter branch. + fourth_expansion_dim: Hidden size for the fourth adapter branch. + fifth_expansion_dim: Hidden size for the fifth adapter branch. + sixth_expansion_dim: Hidden size for the sixth adapter branch. + + Returns: + An nn.Module ready for training or inference. + + Raises: + ValueError: If model_size is not recognized. + """ + if model_size == "base": + return ChampionNet() + if model_size == "large": + return ChampionNetLarge(expansion_dim=expansion_dim, dropout=dropout) + if model_size == "xlarge": + extra_dim = int(extra_expansion_dim) if extra_expansion_dim is not None else int(max(1024, expansion_dim * 2)) + return ChampionNetXL(expansion_dim=expansion_dim, extra_expansion_dim=extra_dim, dropout=dropout) + if model_size == "xxlarge": + extra_dim = int(extra_expansion_dim) if extra_expansion_dim is not None else int(max(2048, expansion_dim * 2)) + third_dim = int(third_expansion_dim) if third_expansion_dim is not None else int(max(3072, extra_dim + expansion_dim)) + return ChampionNetXXL( + expansion_dim=expansion_dim, + extra_expansion_dim=extra_dim, + third_expansion_dim=third_dim, + dropout=dropout, + ) + if model_size == "xxxlarge": + extra_dim = int(extra_expansion_dim) if extra_expansion_dim is not None else int(max(2048, expansion_dim * 2)) + third_dim = int(third_expansion_dim) if third_expansion_dim is not None else int(max(3072, extra_dim + expansion_dim)) + fourth_dim = ( + int(fourth_expansion_dim) + if fourth_expansion_dim is not None + else int(max(4096, third_dim + expansion_dim)) + ) + return ChampionNetXXXL( + expansion_dim=expansion_dim, + extra_expansion_dim=extra_dim, + third_expansion_dim=third_dim, + fourth_expansion_dim=fourth_dim, + dropout=dropout, + ) + if model_size == "ultralarge": + extra_dim = int(extra_expansion_dim) if extra_expansion_dim is not None else int(max(2048, expansion_dim * 2)) + third_dim = int(third_expansion_dim) if third_expansion_dim is not None else int(max(3072, extra_dim + expansion_dim)) + fourth_dim = ( + int(fourth_expansion_dim) + if fourth_expansion_dim is not None + else int(max(4096, third_dim + expansion_dim)) + ) + fifth_dim = ( + int(fifth_expansion_dim) + if fifth_expansion_dim is not None + else int(max(6144, fourth_dim + expansion_dim)) + ) + return ChampionNetUltra( + expansion_dim=expansion_dim, + extra_expansion_dim=extra_dim, + third_expansion_dim=third_dim, + fourth_expansion_dim=fourth_dim, + fifth_expansion_dim=fifth_dim, + dropout=dropout, + ) + if model_size == "megalarge": + extra_dim = int(extra_expansion_dim) if extra_expansion_dim is not None else int(max(2048, expansion_dim * 2)) + third_dim = int(third_expansion_dim) if third_expansion_dim is not None else int(max(3072, extra_dim + expansion_dim)) + fourth_dim = ( + int(fourth_expansion_dim) + if fourth_expansion_dim is not None + else int(max(4096, third_dim + expansion_dim)) + ) + fifth_dim = ( + int(fifth_expansion_dim) + if fifth_expansion_dim is not None + else int(max(6144, fourth_dim + expansion_dim)) + ) + sixth_dim = ( + int(sixth_expansion_dim) + if sixth_expansion_dim is not None + else int(max(8192, fifth_dim + expansion_dim)) + ) + return ChampionNetMega( + expansion_dim=expansion_dim, + extra_expansion_dim=extra_expansion_dim, + third_expansion_dim=third_dim, + fourth_expansion_dim=fourth_dim, + fifth_expansion_dim=fifth_dim, + sixth_expansion_dim=sixth_dim, + dropout=dropout, + ) + if model_size == "ultra_expert": + return ChampionNetUltraExpert(dropout=dropout) + if model_size == "hierarchical_expert": + return ChampionNetHierarchicalExpert(dropout=dropout) + if model_size == "deep_expert": + return ChampionNetDeepExpert(dropout=dropout) + if model_size == "expert_choice": + return ChampionNetExpertChoice(dropout=dropout) + if model_size == "smarter_expert": + return ChampionNetSmarterExpert(dropout=dropout) + if model_size == "thought_expert": + return ChampionNetThoughtExpert(dropout=dropout) + if model_size == "recursive_expert": + return ChampionNetRecursiveExpert(dropout=dropout) + if model_size == "reflexive_expert": + return ChampionNetReflexiveExpert(dropout=dropout) + if model_size == "metacognitive_expert": + return ChampionNetMetaCognitiveExpert(dropout=dropout) + if model_size == "tree_of_thought_expert": + return ChampionNetTreeOfThoughtExpert(dropout=dropout) + if model_size == "consensus_expert": + return ChampionNetConsensusExpert(dropout=dropout) + if model_size == "deliberative_expert": + return ChampionNetDeliberativeExpert(dropout=dropout) + if model_size == "omniscient_expert": + return ChampionNetOmniscientExpert(dropout=dropout) + if model_size == "neurogenesis_expert": + return ChampionNetNeurogenesisExpert(dropout=dropout) + if model_size == "cognitive_expert": + return ChampionNetCognitiveExpert(dropout=dropout) + if model_size == "cognitive_leap_expert": + return ChampionNetCognitiveLeapExpert(dropout=dropout) + if model_size == "cognitive_leap_ultra_expert": + return ChampionNetCognitiveLeapUltraExpert(dropout=dropout) + if model_size == "transcendent_expert": + return ChampionNetTranscendentExpert(dropout=dropout) + if model_size == "omniversal_expert": + return ChampionNetOmniversalExpert(dropout=dropout) + if model_size == "fractal_expert": + return ChampionNetFractalExpert(dropout=dropout) + if model_size == "liquid_spiking_expert": + return ChampionNetLiquidSpikingExpert(dropout=dropout) + if model_size == "active_inference_expert": + return ChampionNetActiveInferenceExpert(dropout=dropout) + if model_size == "holographic_state_space_expert": + return ChampionNetHolographicStateSpaceExpert(dropout=dropout) + if model_size == "paper_fusion_expert": + return ChampionNetPaperFusionExpert(dropout=dropout) + if model_size == "frontier_expert": + from model_frontier_v33 import ChampionNetFrontierExpert + + return ChampionNetFrontierExpert(dropout=dropout) + if model_size == "frontier_collective_expert": + from model_frontier_v35 import ChampionNetFrontierCollectiveExpert + + return ChampionNetFrontierCollectiveExpert(dropout=dropout) + if model_size == "frontier_verifier_expert": + from model_frontier_v39 import ChampionNetFrontierVerifierExpert + + return ChampionNetFrontierVerifierExpert(dropout=dropout) + if model_size == "test_time_diff_expert": + return ChampionNetTestTimeDiffExpert(dropout=dropout) + if model_size == "titan_dreamer_expert": + from model_frontier_v43 import ChampionNetTitanDreamerExpert + + return ChampionNetTitanDreamerExpert(dropout=dropout) + + raise ValueError( + f"Unknown model_size={model_size!r}. Use one of {SUPPORTED_MODEL_SIZES}." + ) + + + +def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) -> Tuple[List[str], List[str]]: + """Load checkpoint weights into a model with backward compatibility. + + Handles warm-starting from smaller checkpoints by filtering out expected + missing keys (new head parameters) and stripping incompatible head weights + when upgrading across variant families. + + Args: + model: The target model to load weights into. + state_dict: Checkpoint state dictionary. + model_size: The variant name of the target model. + + Returns: + Tuple of (unexpected_missing_keys, unexpected_keys). Keys that are + expected to be missing for warm-start scenarios are filtered out. + + Raises: + RuntimeError: If weight dimensions are incompatible. + ValueError: If model_size is not recognized. + """ + if model_size == "base": + incompatible = model.load_state_dict(state_dict, strict=False) + return list(incompatible.missing_keys), list(incompatible.unexpected_keys) + + if model_size == "large": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load large checkpoint: likely head-dimension mismatch. " + "Use matching --expansion_dim or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when loading base checkpoint into large model. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + } + missing_filtered = [k for k in missing if k and k not in allowed_missing] + unexpected_filtered = [k for k in unexpected if k] + return missing_filtered, unexpected_filtered + + if model_size == "xlarge": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load xlarge checkpoint: likely expansion/aux dimension mismatch. " + "Use matching --expansion_dim/--extra_expansion_dim or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when warm-starting xlarge from base/large checkpoints. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + "layers.10.adapter_up_b.weight", + "layers.10.adapter_down_b.weight", + "layers.10.router.weight", + "layers.10.router.bias", + "layers.10.beta", + } + missing_filtered = [k for k in missing if k and k not in allowed_missing] + unexpected_filtered = [k for k in unexpected if k] + return missing_filtered, unexpected_filtered + + if model_size == "xxlarge": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load xxlarge checkpoint: likely expansion/aux/third dimension mismatch. " + "Use matching --expansion_dim/--extra_expansion_dim/--third_expansion_dim or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when warm-starting xxlarge from base/large/xlarge checkpoints. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + "layers.10.adapter_up_b.weight", + "layers.10.adapter_down_b.weight", + "layers.10.router.weight", + "layers.10.router.bias", + "layers.10.beta", + "layers.10.adapter_up_c.weight", + "layers.10.adapter_down_c.weight", + "layers.10.router2.weight", + "layers.10.router2.bias", + "layers.10.gamma", + } + missing_filtered = [k for k in missing if k and k not in allowed_missing] + unexpected_filtered = [k for k in unexpected if k] + return missing_filtered, unexpected_filtered + + if model_size == "xxxlarge": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load xxxlarge checkpoint: likely expansion/aux/third/fourth dimension mismatch. " + "Use matching dims or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when warm-starting xxxlarge from base/large/xlarge/xxlarge checkpoints. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + "layers.10.adapter_up_b.weight", + "layers.10.adapter_down_b.weight", + "layers.10.router.weight", + "layers.10.router.bias", + "layers.10.beta", + "layers.10.adapter_up_c.weight", + "layers.10.adapter_down_c.weight", + "layers.10.router2.weight", + "layers.10.router2.bias", + "layers.10.gamma", + "layers.10.adapter_up_d.weight", + "layers.10.adapter_down_d.weight", + "layers.10.router3.weight", + "layers.10.router3.bias", + "layers.10.delta", + } + missing_filtered = [k for k in missing if k and k not in allowed_missing] + unexpected_filtered = [k for k in unexpected if k] + return missing_filtered, unexpected_filtered + + if model_size == "ultralarge": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load ultralarge checkpoint: likely expansion/aux/third/fourth/fifth dimension mismatch. " + "Use matching dims or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when warm-starting ultralarge from smaller checkpoints. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + "layers.10.adapter_up_b.weight", + "layers.10.adapter_down_b.weight", + "layers.10.router.weight", + "layers.10.router.bias", + "layers.10.beta", + "layers.10.adapter_up_c.weight", + "layers.10.adapter_down_c.weight", + "layers.10.router2.weight", + "layers.10.router2.bias", + "layers.10.gamma", + "layers.10.adapter_up_d.weight", + "layers.10.adapter_down_d.weight", + "layers.10.router3.weight", + "layers.10.router3.bias", + "layers.10.delta", + "layers.10.adapter_up_e.weight", + "layers.10.adapter_down_e.weight", + "layers.10.router4.weight", + "layers.10.router4.bias", + "layers.10.epsilon", + "layers.10.pre_norm.weight", + "layers.10.pre_norm.bias", + "layers.10.domain_router.weight", + "layers.10.domain_router.bias", + "layers.10.domain_experts.weight", + "layers.10.calib_gate.weight", + "layers.10.calib_gate.bias", + "layers.10.zeta", + "layers.10.theta", + } + missing_filtered = [k for k in missing if k and k not in allowed_missing] + unexpected_filtered = [k for k in unexpected if k] + return missing_filtered, unexpected_filtered + + if model_size == "megalarge": + try: + incompatible = model.load_state_dict(state_dict, strict=False) + except RuntimeError as exc: + raise RuntimeError( + "Failed to load megalarge checkpoint: likely expansion dimension mismatch. " + "Use matching dims or auto-detect from checkpoint." + ) from exc + missing = list(incompatible.missing_keys) + unexpected = list(incompatible.unexpected_keys) + + # Expected when warm-starting megalarge from smaller checkpoints. + allowed_missing = { + "layers.10.adapter_up.weight", + "layers.10.adapter_down.weight", + "layers.10.alpha", + "layers.10.adapter_up_b.weight", + "layers.10.adapter_down_b.weight", + "layers.10.router.weight", + "layers.10.router.bias", + "layers.10.beta", + "layers.10.adapter_up_c.weight", + "layers.10.adapter_down_c.weight", + "layers.10.router2.weight", + "layers.10.router2.bias", + "layers.10.gamma", + "layers.10.adapter_up_d.weight", + "layers.10.adapter_down_d.weight", + "layers.10.router3.weight", + "layers.10.router3.bias", + "layers.10.delta", + "layers.10.adapter_up_e.weight", + "layers.10.adapter_down_e.weight", + "layers.10.router4.weight", + "layers.10.router4.bias", + "layers.10.epsilon", + "layers.10.pre_norm.weight", + "layers.10.pre_norm.bias", + "layers.10.domain_router.weight", + "layers.10.domain_router.bias", + "layers.10.domain_experts.weight", + "layers.10.calib_gate.weight", + "layers.10.calib_gate.bias", + "layers.10.zeta", + "layers.10.theta", + # Megalarge-only keys + "layers.10.adapter_up_f.weight", + "layers.10.adapter_down_f.weight", + "layers.10.router5.weight", + "layers.10.router5.bias", + "layers.10.iota", + "layers.10.cross_attn_fusion.q_proj.weight", + "layers.10.cross_attn_fusion.k_proj.weight", + "layers.10.cross_attn_fusion.v_proj.weight", + "layers.10.cross_attn_fusion.out_proj.weight", + "layers.10.cross_attn_fusion.fusion_weight.weight", + "layers.10.cross_attn_fusion.fusion_weight.bias", + "layers.10.kappa", + "layers.10.reasoning_gate.0.weight", + "layers.10.reasoning_gate.2.weight", + "layers.10.reasoning_gate.2.bias", + "layers.10.lambda_", + # GatedFFN backbone layer (layers.10 in mega = GatedFFN) + "layers.10.norm.weight", + "layers.10.up_proj.weight", + "layers.10.down_proj.weight", + "layers.10.layer_scale.gamma", + } + # Megalarge uses layers.10=GatedFFN, layers.11=Head, layers.12=norm + # So head keys are at layers.11.* not layers.10.* + mega_allowed = set() + for k in allowed_missing: + if k.startswith("layers.10."): + mega_key = "layers.11." + k[len("layers.10."):] + mega_allowed.add(mega_key) + mega_allowed.update(allowed_missing) + # Also allow GatedFFN keys at layers.10 + mega_allowed.update({ + "layers.10.norm.weight", + "layers.10.up_proj.weight", + "layers.10.down_proj.weight", + "layers.10.layer_scale.gamma", + }) + + # Robust loading: if checkpoint is not megalarge, strip head keys to avoid mismatch + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "megalarge": + filtered_sd = {k: v for k, v in state_dict.items() if not k.startswith("layers.10.") and not k.startswith("layers.11.")} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in mega_allowed] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "ultra_expert": + allowed_missing = { + "layers.10.experts_up.0.weight", "layers.10.experts_up.1.weight", + "layers.10.experts_up.2.weight", "layers.10.experts_up.3.weight", + "layers.10.experts_up.4.weight", "layers.10.experts_up.5.weight", + "layers.10.experts_down.0.weight", "layers.10.experts_down.1.weight", + "layers.10.experts_down.2.weight", "layers.10.experts_down.3.weight", + "layers.10.experts_down.4.weight", "layers.10.experts_down.5.weight", + "layers.10.gate.weight", "layers.10.noise_gate.weight", + "layers.10.alpha", "layers.10.calibration.weight", + "layers.10.calibration.bias", "layers.10.theta", + } + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "ultra_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "hierarchical_expert": + # Allow all new hierarchical head parameters to be missing when warm-starting + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", f"{head_pref}shared_scale", + f"{head_pref}domain_gate.weight", f"{head_pref}domain_noise_gate.weight", + f"{head_pref}alpha", f"{head_pref}theta", + f"{head_pref}calibration.weight", f"{head_pref}calibration.bias", + } + for i in range(2): # groups + allowed_missing.add(f"{head_pref}expert_gates.{i}.weight") + allowed_missing.add(f"{head_pref}expert_noise_gates.{i}.weight") + for i in range(8): # experts + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + allowed_missing.add(f"{head_pref}lora_down.{i}.weight") + allowed_missing.add(f"{head_pref}lora_up.{i}.weight") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "hierarchical_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "thought_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}expert_bias", + f"{head_pref}alpha", f"{head_pref}calibration.weight", + f"{head_pref}calibration.bias", f"{head_pref}theta", + f"{head_pref}cross_attn.q_proj.weight", f"{head_pref}cross_attn.k_proj.weight", + f"{head_pref}cross_attn.v_proj.weight", f"{head_pref}cross_attn.out_proj.weight", + f"{head_pref}cross_attn.fusion_weight.weight", f"{head_pref}cross_attn.fusion_weight.bias", + } + for i in range(8): # n_experts + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + allowed_missing.add(f"{head_pref}lora_down.{i}.weight") + allowed_missing.add(f"{head_pref}lora_up.{i}.weight") + for i in range(3): # reasoning_steps + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.bias") + allowed_missing.add(f"{head_pref}gates.{i}.weight") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "thought_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "recursive_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}expert_bias", + f"{head_pref}local_up.weight", f"{head_pref}local_down.weight", + f"{head_pref}local_norm.weight", f"{head_pref}local_norm.bias", + f"{head_pref}local_scale", + f"{head_pref}alpha", f"{head_pref}calibration.weight", + f"{head_pref}calibration.bias", f"{head_pref}theta", + f"{head_pref}cross_attn.q_proj.weight", f"{head_pref}cross_attn.k_proj.weight", + f"{head_pref}cross_attn.v_proj.weight", f"{head_pref}cross_attn.out_proj.weight", + f"{head_pref}cross_attn.fusion_weight.weight", f"{head_pref}cross_attn.fusion_weight.bias", + } + for i in range(8): # n_experts + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + allowed_missing.add(f"{head_pref}lora_down.{i}.weight") + allowed_missing.add(f"{head_pref}lora_up.{i}.weight") + for i in range(3): # reasoning_steps + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.bias") + allowed_missing.add(f"{head_pref}exit_gates.{i}.weight") + allowed_missing.add(f"{head_pref}exit_gates.{i}.bias") + for h in range(2): # n_heads + allowed_missing.add(f"{head_pref}gates.{i * 2 + h}.weight") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "recursive_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "reflexive_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}initial_gate.weight", + f"{head_pref}critique_gate.weight", + f"{head_pref}alpha", + f"{head_pref}beta", + } + for i in range(4): # default n_experts for reflexive + allowed_missing.add(f"{head_pref}initial_up.{i}.weight") + allowed_missing.add(f"{head_pref}initial_down.{i}.weight") + allowed_missing.add(f"{head_pref}initial_norms.{i}.weight") + allowed_missing.add(f"{head_pref}initial_norms.{i}.bias") + allowed_missing.add(f"{head_pref}critique_up.{i}.weight") + allowed_missing.add(f"{head_pref}critique_down.{i}.weight") + allowed_missing.add(f"{head_pref}critique_norms.{i}.weight") + allowed_missing.add(f"{head_pref}critique_norms.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "reflexive_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "metacognitive_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}proposal_bias", + f"{head_pref}alpha", f"{head_pref}beta", + } + for i in range(6): # n_proposal_experts + allowed_missing.add(f"{head_pref}proposal_up.{i}.weight") + allowed_missing.add(f"{head_pref}proposal_down.{i}.weight") + allowed_missing.add(f"{head_pref}proposal_norms.{i}.weight") + allowed_missing.add(f"{head_pref}proposal_norms.{i}.bias") + for i in range(4): # n_critique_experts + allowed_missing.add(f"{head_pref}critique_up.{i}.weight") + allowed_missing.add(f"{head_pref}critique_down.{i}.weight") + allowed_missing.add(f"{head_pref}critique_norms.{i}.weight") + allowed_missing.add(f"{head_pref}critique_norms.{i}.bias") + for i in range(3): # reasoning_steps + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.bias") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.bias") + allowed_missing.add(f"{head_pref}proposal_gates.{i}.weight") + allowed_missing.add(f"{head_pref}critique_gates.{i}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{i}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "metacognitive_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "tree_of_thought_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}final_up.weight", f"{head_pref}final_down.weight", + f"{head_pref}final_norm.weight", f"{head_pref}final_norm.bias", + } + for i in range(6): # n_action_experts + allowed_missing.add(f"{head_pref}action_up.{i}.weight") + allowed_missing.add(f"{head_pref}action_down.{i}.weight") + allowed_missing.add(f"{head_pref}action_norms.{i}.weight") + allowed_missing.add(f"{head_pref}action_norms.{i}.bias") + # Value net + allowed_missing.add(f"{head_pref}value_net.0.weight") + allowed_missing.add(f"{head_pref}value_net.0.bias") + allowed_missing.add(f"{head_pref}value_net.3.weight") + allowed_missing.add(f"{head_pref}value_net.3.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "tree_of_thought_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "consensus_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}alpha", + f"{head_pref}mlp_gate.weight", + f"{head_pref}attn_embed", + f"{head_pref}attn_q.weight", f"{head_pref}attn_k.weight", + f"{head_pref}attn_v.weight", f"{head_pref}attn_proj.weight", + f"{head_pref}attn_norm.weight", f"{head_pref}attn_norm.bias", + f"{head_pref}conv1.weight", f"{head_pref}conv1.bias", + f"{head_pref}conv2.weight", f"{head_pref}conv2.bias", + f"{head_pref}conv_proj.weight", + f"{head_pref}conv_norm.weight", f"{head_pref}conv_norm.bias", + f"{head_pref}consensus.0.weight", f"{head_pref}consensus.0.bias", + f"{head_pref}consensus.3.weight", f"{head_pref}consensus.3.bias", + f"{head_pref}consensus.5.weight", f"{head_pref}consensus.5.bias", + } + for i in range(6): # n_mlp_experts + allowed_missing.add(f"{head_pref}mlp_up.{i}.weight") + allowed_missing.add(f"{head_pref}mlp_down.{i}.weight") + allowed_missing.add(f"{head_pref}mlp_norms.{i}.weight") + allowed_missing.add(f"{head_pref}mlp_norms.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "consensus_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "deliberative_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}memory_keys", f"{head_pref}memory_values", + f"{head_pref}mem_read_query.weight", + f"{head_pref}mem_write_query.weight", + f"{head_pref}mem_write_gate.weight", f"{head_pref}mem_write_gate.bias", + # Consistency net + f"{head_pref}consistency_net.0.weight", f"{head_pref}consistency_net.0.bias", + f"{head_pref}consistency_net.3.weight", f"{head_pref}consistency_net.3.bias", + f"{head_pref}consistency_net.5.weight", f"{head_pref}consistency_net.5.bias", + } + for i in range(6): # n_experts + allowed_missing.add(f"{head_pref}expert_up.{i}.weight") + allowed_missing.add(f"{head_pref}expert_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + for step in range(3): # reasoning_steps + for d in range(3): # n_drafts + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.0.weight") + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.0.bias") + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.1.weight") + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.1.bias") + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.3.weight") + allowed_missing.add(f"{head_pref}draft_reasoning_cells.{step}.{d}.3.bias") + allowed_missing.add(f"{head_pref}draft_gates.{step}.{d}.weight") + allowed_missing.add(f"{head_pref}cross_draft_q.{step}.weight") + allowed_missing.add(f"{head_pref}cross_draft_k.{step}.weight") + allowed_missing.add(f"{head_pref}cross_draft_v.{step}.weight") + allowed_missing.add(f"{head_pref}cross_draft_norm.{step}.weight") + allowed_missing.add(f"{head_pref}cross_draft_norm.{step}.bias") + allowed_missing.add(f"{head_pref}halt_gates.{step}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{step}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "deliberative_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "omniscient_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}knowledge_core_keys", f"{head_pref}knowledge_core_values", + f"{head_pref}knowledge_query.weight", f"{head_pref}temperature", + f"{head_pref}node_fusion_q.weight", f"{head_pref}node_fusion_k.weight", + f"{head_pref}node_fusion_v.weight", + } + for i in range(6): # n_experts + allowed_missing.add(f"{head_pref}expert_up.{i}.weight") + allowed_missing.add(f"{head_pref}expert_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + for step in range(3): # reasoning_steps + allowed_missing.add(f"{head_pref}got_q.{step}.weight") + allowed_missing.add(f"{head_pref}got_k.{step}.weight") + allowed_missing.add(f"{head_pref}got_v.{step}.weight") + allowed_missing.add(f"{head_pref}got_norm.{step}.weight") + allowed_missing.add(f"{head_pref}got_norm.{step}.bias") + allowed_missing.add(f"{head_pref}route_mu.{step}.weight") + allowed_missing.add(f"{head_pref}route_logvar.{step}.weight") + allowed_missing.add(f"{head_pref}route_logvar.{step}.bias") + allowed_missing.add(f"{head_pref}critique_net.{step}.0.weight") + allowed_missing.add(f"{head_pref}critique_net.{step}.0.bias") + allowed_missing.add(f"{head_pref}critique_net.{step}.3.weight") + allowed_missing.add(f"{head_pref}critique_net.{step}.3.bias") + allowed_missing.add(f"{head_pref}node_halt_gates.{step}.weight") + allowed_missing.add(f"{head_pref}node_halt_gates.{step}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "omniscient_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "omniversal_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}holographic_memory", + f"{head_pref}real_lift.weight", f"{head_pref}real_lift.bias", + f"{head_pref}imag_lift.weight", f"{head_pref}imag_lift.bias", + f"{head_pref}uncertainty_gate.weight", f"{head_pref}uncertainty_gate.bias", + } + for sub in ['ode_func.net.0', 'ode_func.net.3', 'ode_func.q_proj', 'ode_func.k_proj', 'ode_func.v_proj']: + allowed_missing.add(f"{head_pref}{sub}.weight") + allowed_missing.add(f"{head_pref}{sub}.bias") + # measurement_proj + allowed_missing.add(f"{head_pref}measurement_proj.0.weight") + allowed_missing.add(f"{head_pref}measurement_proj.0.bias") + allowed_missing.add(f"{head_pref}measurement_proj.2.weight") + allowed_missing.add(f"{head_pref}measurement_proj.2.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "omniversal_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "transcendent_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}time_embed", f"{head_pref}log_betas", + f"{head_pref}attn_q.weight", f"{head_pref}attn_k.weight", f"{head_pref}attn_v.weight", + f"{head_pref}attn_norm.weight", f"{head_pref}attn_norm.bias", + f"{head_pref}meta_q_mod.weight", f"{head_pref}meta_k_mod.weight", + f"{head_pref}fitness_proj.weight", f"{head_pref}fitness_proj.bias", + } + for i in range(4): # decompose_heads + for j in [0,1]: + allowed_missing.add(f"{head_pref}decompose_heads.{i}.{j}.weight") + allowed_missing.add(f"{head_pref}decompose_heads.{i}.{j}.bias") + allowed_missing.add(f"{head_pref}decompose_heads.{i}.1.weight") + allowed_missing.add(f"{head_pref}decompose_heads.{i}.1.bias") + for t in range(4): # denoiser steps + allowed_missing.add(f"{head_pref}denoiser.{t}.0.weight") + allowed_missing.add(f"{head_pref}denoiser.{t}.0.bias") + allowed_missing.add(f"{head_pref}denoiser.{t}.3.weight") + allowed_missing.add(f"{head_pref}denoiser.{t}.3.bias") + for sub in ['world_transition', 'world_value', 'crossover_net', 'composer']: + for idx in range(5): + allowed_missing.add(f"{head_pref}{sub}.{idx}.weight") + allowed_missing.add(f"{head_pref}{sub}.{idx}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "transcendent_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "cognitive_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}memory_matrix", + f"{head_pref}mem_read_q.weight", f"{head_pref}mem_read_out.weight", + f"{head_pref}mem_write_gate.weight", f"{head_pref}mem_write_gate.bias", + f"{head_pref}mem_write_data.weight", f"{head_pref}mem_write_proj.weight", + f"{head_pref}causal_q.weight", f"{head_pref}causal_k.weight", f"{head_pref}causal_v.weight", + f"{head_pref}causal_norm.weight", f"{head_pref}causal_norm.bias", + } + for i in range(4): # hypernet + allowed_missing.add(f"{head_pref}hypernet.{i}.weight") + allowed_missing.add(f"{head_pref}hypernet.{i}.bias") + # orth paths + for i in range(4): # n_hypotheses + allowed_missing.add(f"{head_pref}path_projs.{i}.weight") + for i in range(3): # value_net + allowed_missing.add(f"{head_pref}value_net.{i}.weight") + allowed_missing.add(f"{head_pref}value_net.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "cognitive_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "cognitive_leap_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}thought_init.weight", f"{head_pref}answer_init.weight", + f"{head_pref}latent_gain", f"{head_pref}cycle_embed.weight", + f"{head_pref}halt_head.weight", f"{head_pref}halt_head.bias", + f"{head_pref}decode_head.weight", + } + for i in (0, 3): # recur_core linear layers + allowed_missing.add(f"{head_pref}recur_core.{i}.weight") + allowed_missing.add(f"{head_pref}recur_core.{i}.bias") + for i in (0, 2): # answer_update linear layers + allowed_missing.add(f"{head_pref}answer_update.{i}.weight") + allowed_missing.add(f"{head_pref}answer_update.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "cognitive_leap_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "cognitive_leap_ultra_expert": + head_pref = "layers.10." + allowed_missing = { + key + for key in model.state_dict().keys() + if key.startswith(head_pref) and key not in {f"{head_pref}weight", f"{head_pref}bias"} + } + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "cognitive_leap_ultra_expert": + filtered_sd = { + k: v + for k, v in state_dict.items() + if not (k.startswith(head_pref) and k not in [f"{head_pref}weight", f"{head_pref}bias"]) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + + if model_size == "neurogenesis_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + } + # Level up/down projections and knowledge banks + for lvl in range(3): # n_levels + allowed_missing.add(f"{head_pref}knowledge_keys.{lvl}") + allowed_missing.add(f"{head_pref}knowledge_values.{lvl}") + allowed_missing.add(f"{head_pref}knowledge_queries.{lvl}.weight") + for i in range(2): # n_levels - 1 + allowed_missing.add(f"{head_pref}level_up_projs.{i}.0.weight") + allowed_missing.add(f"{head_pref}level_up_projs.{i}.0.bias") + allowed_missing.add(f"{head_pref}level_up_projs.{i}.1.weight") + allowed_missing.add(f"{head_pref}level_up_projs.{i}.1.bias") + allowed_missing.add(f"{head_pref}level_down_projs.{i}.0.weight") + allowed_missing.add(f"{head_pref}level_down_projs.{i}.0.bias") + allowed_missing.add(f"{head_pref}level_down_projs.{i}.1.weight") + allowed_missing.add(f"{head_pref}level_down_projs.{i}.1.bias") + allowed_missing.add(f"{head_pref}distill_bridges.{i}.weight") + # Proposer + Adversary experts + for i in range(6): # n_experts + allowed_missing.add(f"{head_pref}proposer_up.{i}.weight") + allowed_missing.add(f"{head_pref}proposer_down.{i}.weight") + allowed_missing.add(f"{head_pref}proposer_norms.{i}.weight") + allowed_missing.add(f"{head_pref}proposer_norms.{i}.bias") + allowed_missing.add(f"{head_pref}adversary_up.{i}.weight") + allowed_missing.add(f"{head_pref}adversary_down.{i}.weight") + allowed_missing.add(f"{head_pref}adversary_norms.{i}.weight") + allowed_missing.add(f"{head_pref}adversary_norms.{i}.bias") + # Per-round components + for rnd in range(3): # debate_rounds + allowed_missing.add(f"{head_pref}proposer_gates.{rnd}.weight") + allowed_missing.add(f"{head_pref}adversary_gates.{rnd}.weight") + allowed_missing.add(f"{head_pref}synthesis_net.{rnd}.0.weight") + allowed_missing.add(f"{head_pref}synthesis_net.{rnd}.0.bias") + allowed_missing.add(f"{head_pref}synthesis_net.{rnd}.2.weight") + allowed_missing.add(f"{head_pref}synthesis_net.{rnd}.2.bias") + allowed_missing.add(f"{head_pref}resolution.{rnd}.0.weight") + allowed_missing.add(f"{head_pref}resolution.{rnd}.0.bias") + allowed_missing.add(f"{head_pref}resolution.{rnd}.3.weight") + allowed_missing.add(f"{head_pref}resolution.{rnd}.3.bias") + allowed_missing.add(f"{head_pref}confidence_gates.{rnd}.weight") + allowed_missing.add(f"{head_pref}confidence_gates.{rnd}.bias") + # Level fusion + allowed_missing.add(f"{head_pref}level_fusion.0.weight") + allowed_missing.add(f"{head_pref}level_fusion.0.bias") + allowed_missing.add(f"{head_pref}level_fusion.2.weight") + allowed_missing.add(f"{head_pref}level_fusion.2.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "neurogenesis_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "deep_expert": + allowed_missing = { + "layers.10.shared_up.weight", "layers.10.shared_down.weight", + "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", + "layers.10.shared_scale", "layers.10.expert_bias", + "layers.10.gate.weight", "layers.10.noise_gate.weight", + "layers.10.alpha", "layers.10.calibration.weight", + "layers.10.calibration.bias", "layers.10.theta", + } + for i in range(8): # n_experts max default + allowed_missing.add(f"layers.10.experts_up.{i}.weight") + allowed_missing.add(f"layers.10.experts_down.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "deep_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "expert_choice": + allowed_missing = { + "layers.10.gate.weight", "layers.10.noise_gate.weight", + "layers.10.alpha", "layers.10.calibration.weight", + "layers.10.calibration.bias", "layers.10.theta", + } + for i in range(8): # default max experts + allowed_missing.add(f"layers.10.experts_up.{i}.weight") + allowed_missing.add(f"layers.10.experts_down.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "expert_choice": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "smarter_expert": + allowed_missing = { + "layers.10.shared_up.weight", "layers.10.shared_down.weight", + "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", + "layers.10.shared_scale", "layers.10.expert_bias", + "layers.10.gate.weight", "layers.10.alpha", + "layers.10.calibration.weight", "layers.10.calibration.bias", "layers.10.theta", + } + for i in range(8): + allowed_missing.add(f"layers.10.experts_up.{i}.weight") + allowed_missing.add(f"layers.10.experts_down.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.weight") + allowed_missing.add(f"layers.10.expert_norms.{i}.bias") + allowed_missing.add(f"layers.10.lora_down.{i}.weight") + allowed_missing.add(f"layers.10.lora_up.{i}.weight") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "smarter_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + + if model_size == "omniversal_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}holographic_memory", + f"{head_pref}real_lift.weight", f"{head_pref}real_lift.bias", + f"{head_pref}imag_lift.weight", f"{head_pref}imag_lift.bias", + f"{head_pref}uncertainty_gate.weight", f"{head_pref}uncertainty_gate.bias", + } + for sub in ['ode_func.net.0', 'ode_func.net.3', 'ode_func.q_proj', 'ode_func.k_proj', 'ode_func.v_proj']: + allowed_missing.add(f"{head_pref}{sub}.weight") + allowed_missing.add(f"{head_pref}{sub}.bias") + # measurement_proj + allowed_missing.add(f"{head_pref}measurement_proj.0.weight") + allowed_missing.add(f"{head_pref}measurement_proj.0.bias") + allowed_missing.add(f"{head_pref}measurement_proj.2.weight") + allowed_missing.add(f"{head_pref}measurement_proj.2.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "omniversal_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "fractal_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}curvature", f"{head_pref}logic_atoms", + f"{head_pref}hyperbolic_proj.weight", f"{head_pref}hyperbolic_proj.bias", + f"{head_pref}assembly_net.0.weight", f"{head_pref}assembly_net.0.bias", + f"{head_pref}assembly_net.2.weight", f"{head_pref}assembly_net.2.bias", + f"{head_pref}final_projection.0.weight", f"{head_pref}final_projection.0.bias", + f"{head_pref}final_projection.3.weight", f"{head_pref}final_projection.3.bias", + } + for i in range(2): # fractal_depth + for sub in ['update_net.0', 'update_net.3', 'stochastic_gate']: + allowed_missing.add(f"{head_pref}fractal_evolvers.{i}.{sub}.weight") + allowed_missing.add(f"{head_pref}fractal_evolvers.{i}.{sub}.bias") + for i in range(1): # fractal_depth - 1 + allowed_missing.add(f"{head_pref}node_spawners.{i}.weight") + allowed_missing.add(f"{head_pref}node_spawners.{i}.bias") + allowed_missing.add(f"{head_pref}resonance_combiners.{i}.weight") + allowed_missing.add(f"{head_pref}resonance_combiners.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "fractal_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "liquid_spiking_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", + f"{head_pref}ttt_lora_A", f"{head_pref}ttt_lora_B", + f"{head_pref}lif_neuron.threshold", f"{head_pref}lif_neuron.alpha", + f"{head_pref}liquid_synapse.core_1", f"{head_pref}liquid_synapse.core_2", + f"{head_pref}liquid_synapse.tau", + } + for sub in ['ttt_decoder.weight', 'ttt_decoder.bias']: + allowed_missing.add(f"{head_pref}{sub}") + for sub in ['liquid_synapse.liquid_gate.weight', 'liquid_synapse.liquid_gate.bias', 'spike_decoder.weight', 'spike_decoder.bias']: + allowed_missing.add(f"{head_pref}{sub}") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "liquid_spiking_expert": + filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "frontier_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", f"{head_pref}beta", + f"{head_pref}memory_keys", f"{head_pref}memory_values", + f"{head_pref}memory_query.weight", f"{head_pref}memory_out.weight", + f"{head_pref}memory_write_gate.weight", f"{head_pref}memory_write_gate.bias", + f"{head_pref}memory_write_value.weight", f"{head_pref}memory_write_value.bias", + f"{head_pref}depth_query.weight", f"{head_pref}depth_key.weight", + f"{head_pref}depth_value.weight", f"{head_pref}depth_out.weight", + f"{head_pref}router_budget.weight", f"{head_pref}router_budget.bias", + f"{head_pref}expert_bias", + } + for i in range(4): + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}halt_gates.{i}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{i}.bias") + for i in range(10): + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + for sub in ( + "ssm_adapter.A", + "ssm_adapter.B", + "ssm_adapter.C", + "ssm_adapter.D", + "ssm_adapter.gate", + "mor_router.main_router.weight", + "mor_router.main_router.bias", + ): + allowed_missing.add(f"{head_pref}{sub}") + for i in range(4): + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.weight") + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "frontier_expert": + keep_head_keys = { + f"{head_pref}weight", + f"{head_pref}bias", + f"{head_pref}shared_up.weight", + f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", + f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", + f"{head_pref}alpha", + } + filtered_sd = { + k: v + for k, v in state_dict.items() + if (not k.startswith(head_pref)) or (k in keep_head_keys) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "frontier_collective_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", f"{head_pref}beta", + f"{head_pref}memory_keys", f"{head_pref}memory_values", + f"{head_pref}memory_query.weight", f"{head_pref}memory_out.weight", + f"{head_pref}memory_write_gate.weight", f"{head_pref}memory_write_gate.bias", + f"{head_pref}memory_write_value.weight", f"{head_pref}memory_write_value.bias", + f"{head_pref}depth_query.weight", f"{head_pref}depth_key.weight", + f"{head_pref}depth_value.weight", f"{head_pref}depth_out.weight", + f"{head_pref}router_budget.weight", f"{head_pref}router_budget.bias", + f"{head_pref}expert_bias", + f"{head_pref}shared_aux_up.weight", f"{head_pref}shared_aux_down.weight", + f"{head_pref}shared_aux_norm.weight", f"{head_pref}shared_aux_norm.bias", + f"{head_pref}shared_aux_gate.weight", f"{head_pref}shared_aux_gate.bias", + f"{head_pref}reflect_out.weight", + f"{head_pref}collective_q.weight", f"{head_pref}collective_k.weight", + f"{head_pref}collective_v.weight", f"{head_pref}collective_out.weight", + f"{head_pref}gamma", f"{head_pref}delta", + } + for i in range(5): + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}halt_gates.{i}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{i}.bias") + for i in range(10): + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + for i in range(2): + allowed_missing.add(f"{head_pref}reflect_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.norm.bias") + for sub in ( + "ssm_adapter.A", + "ssm_adapter.B", + "ssm_adapter.C", + "ssm_adapter.D", + "ssm_adapter.gate", + "mor_router.main_router.weight", + "mor_router.main_router.bias", + ): + allowed_missing.add(f"{head_pref}{sub}") + for i in range(4): + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.weight") + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.bias") + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "frontier_collective_expert": + keep_head_keys = { + f"{head_pref}weight", + f"{head_pref}bias", + f"{head_pref}shared_up.weight", + f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", + f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", + f"{head_pref}alpha", + f"{head_pref}beta", + f"{head_pref}memory_keys", + f"{head_pref}memory_values", + f"{head_pref}memory_query.weight", + f"{head_pref}memory_out.weight", + f"{head_pref}memory_write_gate.weight", + f"{head_pref}memory_write_gate.bias", + f"{head_pref}memory_write_value.weight", + f"{head_pref}memory_write_value.bias", + f"{head_pref}depth_query.weight", + f"{head_pref}depth_key.weight", + f"{head_pref}depth_value.weight", + f"{head_pref}depth_out.weight", + f"{head_pref}router_budget.weight", + f"{head_pref}router_budget.bias", + f"{head_pref}expert_bias", + f"{head_pref}ssm_adapter.A", + f"{head_pref}ssm_adapter.B", + f"{head_pref}ssm_adapter.C", + f"{head_pref}ssm_adapter.D", + f"{head_pref}ssm_adapter.gate", + f"{head_pref}mor_router.main_router.weight", + f"{head_pref}mor_router.main_router.bias", + } + for i in range(5): + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.up.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.down.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + keep_head_keys.add(f"{head_pref}halt_gates.{i}.weight") + keep_head_keys.add(f"{head_pref}halt_gates.{i}.bias") + for i in range(10): + keep_head_keys.add(f"{head_pref}experts_up.{i}.weight") + keep_head_keys.add(f"{head_pref}experts_down.{i}.weight") + keep_head_keys.add(f"{head_pref}expert_norms.{i}.weight") + keep_head_keys.add(f"{head_pref}expert_norms.{i}.bias") + for i in range(4): + keep_head_keys.add(f"{head_pref}mor_router.sub_routers.{i}.weight") + keep_head_keys.add(f"{head_pref}mor_router.sub_routers.{i}.bias") + filtered_sd = { + k: v + for k, v in state_dict.items() + if (not k.startswith(head_pref)) or (k in keep_head_keys) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size == "frontier_verifier_expert": + head_pref = "layers.10." + allowed_missing = { + f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", f"{head_pref}alpha", f"{head_pref}beta", + f"{head_pref}memory_keys", f"{head_pref}memory_values", + f"{head_pref}memory_query.weight", f"{head_pref}memory_out.weight", + f"{head_pref}memory_write_gate.weight", f"{head_pref}memory_write_gate.bias", + f"{head_pref}memory_write_value.weight", f"{head_pref}memory_write_value.bias", + f"{head_pref}depth_query.weight", f"{head_pref}depth_key.weight", + f"{head_pref}depth_value.weight", f"{head_pref}depth_out.weight", + f"{head_pref}router_budget.weight", f"{head_pref}router_budget.bias", + f"{head_pref}expert_bias", + f"{head_pref}shared_aux_up.weight", f"{head_pref}shared_aux_down.weight", + f"{head_pref}shared_aux_norm.weight", f"{head_pref}shared_aux_norm.bias", + f"{head_pref}shared_aux_gate.weight", f"{head_pref}shared_aux_gate.bias", + f"{head_pref}reflect_out.weight", + f"{head_pref}collective_q.weight", f"{head_pref}collective_k.weight", + f"{head_pref}collective_v.weight", f"{head_pref}collective_out.weight", + f"{head_pref}revisit_gate.weight", f"{head_pref}revisit_gate.bias", + f"{head_pref}verifier_q.weight", f"{head_pref}verifier_k.weight", + f"{head_pref}verifier_v.weight", f"{head_pref}verifier_out.weight", + f"{head_pref}verifier_gate.weight", f"{head_pref}verifier_gate.bias", + f"{head_pref}correction_out.weight", + f"{head_pref}gamma", f"{head_pref}delta", f"{head_pref}epsilon", f"{head_pref}zeta", + } + for i in range(6): + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + allowed_missing.add(f"{head_pref}halt_gates.{i}.weight") + allowed_missing.add(f"{head_pref}halt_gates.{i}.bias") + for i in range(12): + allowed_missing.add(f"{head_pref}experts_up.{i}.weight") + allowed_missing.add(f"{head_pref}experts_down.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.weight") + allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") + for i in range(3): + allowed_missing.add(f"{head_pref}reflect_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}reflect_cells.{i}.norm.bias") + for i in range(2): + allowed_missing.add(f"{head_pref}verifier_cells.{i}.up.weight") + allowed_missing.add(f"{head_pref}verifier_cells.{i}.down.weight") + allowed_missing.add(f"{head_pref}verifier_cells.{i}.norm.weight") + allowed_missing.add(f"{head_pref}verifier_cells.{i}.norm.bias") + for sub in ( + "ssm_adapter.A", + "ssm_adapter.B", + "ssm_adapter.C", + "ssm_adapter.D", + "ssm_adapter.gate", + "mor_router.main_router.weight", + "mor_router.main_router.bias", + ): + allowed_missing.add(f"{head_pref}{sub}") + for i in range(5): + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.weight") + allowed_missing.add(f"{head_pref}mor_router.sub_routers.{i}.bias") + + target_state = model.state_dict() + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "frontier_verifier_expert": + keep_head_keys = { + f"{head_pref}weight", + f"{head_pref}bias", + f"{head_pref}shared_up.weight", + f"{head_pref}shared_down.weight", + f"{head_pref}shared_norm.weight", + f"{head_pref}shared_norm.bias", + f"{head_pref}shared_scale", + f"{head_pref}alpha", + f"{head_pref}beta", + f"{head_pref}memory_keys", + f"{head_pref}memory_values", + f"{head_pref}memory_query.weight", + f"{head_pref}memory_out.weight", + f"{head_pref}memory_write_gate.weight", + f"{head_pref}memory_write_gate.bias", + f"{head_pref}memory_write_value.weight", + f"{head_pref}memory_write_value.bias", + f"{head_pref}depth_query.weight", + f"{head_pref}depth_key.weight", + f"{head_pref}depth_value.weight", + f"{head_pref}depth_out.weight", + f"{head_pref}router_budget.weight", + f"{head_pref}router_budget.bias", + f"{head_pref}expert_bias", + f"{head_pref}shared_aux_up.weight", + f"{head_pref}shared_aux_down.weight", + f"{head_pref}shared_aux_norm.weight", + f"{head_pref}shared_aux_norm.bias", + f"{head_pref}shared_aux_gate.weight", + f"{head_pref}shared_aux_gate.bias", + f"{head_pref}reflect_out.weight", + f"{head_pref}collective_q.weight", + f"{head_pref}collective_k.weight", + f"{head_pref}collective_v.weight", + f"{head_pref}collective_out.weight", + f"{head_pref}gamma", + f"{head_pref}delta", + f"{head_pref}ssm_adapter.A", + f"{head_pref}ssm_adapter.B", + f"{head_pref}ssm_adapter.C", + f"{head_pref}ssm_adapter.D", + f"{head_pref}ssm_adapter.gate", + f"{head_pref}mor_router.main_router.weight", + f"{head_pref}mor_router.main_router.bias", + } + for i in range(5): + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.up.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.down.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.norm.weight") + keep_head_keys.add(f"{head_pref}reasoning_cells.{i}.norm.bias") + keep_head_keys.add(f"{head_pref}halt_gates.{i}.weight") + keep_head_keys.add(f"{head_pref}halt_gates.{i}.bias") + for i in range(10): + keep_head_keys.add(f"{head_pref}experts_up.{i}.weight") + keep_head_keys.add(f"{head_pref}experts_down.{i}.weight") + keep_head_keys.add(f"{head_pref}expert_norms.{i}.weight") + keep_head_keys.add(f"{head_pref}expert_norms.{i}.bias") + for i in range(2): + keep_head_keys.add(f"{head_pref}reflect_cells.{i}.up.weight") + keep_head_keys.add(f"{head_pref}reflect_cells.{i}.down.weight") + keep_head_keys.add(f"{head_pref}reflect_cells.{i}.norm.weight") + keep_head_keys.add(f"{head_pref}reflect_cells.{i}.norm.bias") + for i in range(4): + keep_head_keys.add(f"{head_pref}mor_router.sub_routers.{i}.weight") + keep_head_keys.add(f"{head_pref}mor_router.sub_routers.{i}.bias") + filtered_sd = { + k: v + for k, v in state_dict.items() + if (((not k.startswith(head_pref)) or (k in keep_head_keys)) + and k in target_state + and tuple(target_state[k].shape) == tuple(v.shape)) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + filtered_sd = { + k: v + for k, v in state_dict.items() + if k in target_state and tuple(target_state[k].shape) == tuple(v.shape) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + if model_size in ("test_time_diff_expert", "titan_dreamer_expert"): + head_pref = "layers.10." + target_state = model.state_dict() + keep_head_keys = { + f"{head_pref}weight", + f"{head_pref}bias", + } + allowed_missing = { + key + for key in target_state.keys() + if key.startswith(head_pref) and key not in keep_head_keys + } + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != model_size: + candidate_state = { + key: value + for key, value in state_dict.items() + if (not key.startswith(head_pref)) or (key in keep_head_keys) + } + else: + candidate_state = dict(state_dict) + + filtered_sd = { + key: value + for key, value in candidate_state.items() + if key in target_state and tuple(target_state[key].shape) == tuple(value.shape) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + + raise ValueError(f"Unsupported model_size={model_size!r}") + + +def detect_model_size_from_state_dict(state_dict: dict) -> str: + """Infer the model variant from checkpoint keys. + + Examines the state_dict for variant-specific parameter names to + determine which model architecture produced the checkpoint. + + Args: + state_dict: Checkpoint state dictionary. + + Returns: + A model size string (e.g., 'base', 'ultra_expert', 'smarter_expert'). + """ + # Megalarge: has GatedFFN at layers.10 + head at layers.11 + if "layers.11.adapter_up_f.weight" in state_dict or "layers.11.router5.weight" in state_dict: + return "megalarge" + if "layers.10.adapter_up_e.weight" in state_dict or "layers.10.router4.weight" in state_dict: + return "ultralarge" + if "layers.10.adapter_up_d.weight" in state_dict or "layers.10.router3.weight" in state_dict: + return "xxxlarge" + if "layers.10.adapter_up_c.weight" in state_dict or "layers.10.router2.weight" in state_dict: + return "xxlarge" + if "layers.10.adapter_up_b.weight" in state_dict or "layers.10.router.weight" in state_dict: + return "xlarge" + if "layers.10.adapter_up.weight" in state_dict or "layers.10.adapter_down.weight" in state_dict: + return "large" + if "layers.10.diff_q_pos.weight" in state_dict and "layers.10.predictive_prior.weight" in state_dict: + return "test_time_diff_expert" + if "layers.10.titan_w1" in state_dict and "layers.10.recursion_router.weight" in state_dict: + return "titan_dreamer_expert" + if "layers.10.domain_gate.weight" in state_dict or "layers.10.expert_gates.0.weight" in state_dict: + return "hierarchical_expert" + if "layers.10.shared_up.weight" in state_dict and "layers.10.gate.weight" in state_dict: + if "layers.10.lora_up.0.weight" in state_dict: + return "smarter_expert" + return "deep_expert" + + if "layers.10.experts_up.0.weight" in state_dict and "layers.10.gate.weight" in state_dict and "layers.10.shared_scale" not in state_dict: + # ultra_expert also has gate.weight, but we distinguish EC by checking if it doesn't have EC specific but also no base keys maybe? + # A simpler way is to check the capacity_factor or something, but it's not in state dict. + # Actually both ultra_expert and expert_choice have gate.weight and experts_up.0.weight. + # We can detect expert_choice by checking for expert_norms which ultra_expert lacks. + if "layers.10.expert_norms.0.weight" in state_dict: + return "expert_choice" + return "ultra_expert" + + if "layers.10.halt_gates.0.weight" in state_dict and "layers.10.proposal_gates.0.weight" in state_dict: + return "metacognitive_expert" + if "layers.10.value_net.0.weight" in state_dict and "layers.10.action_up.0.weight" in state_dict: + return "tree_of_thought_expert" + if "layers.10.knowledge_core_keys" in state_dict and "layers.10.route_mu.0.weight" in state_dict: + return "omniscient_expert" + if "layers.10.hypernet.0.weight" in state_dict and "layers.10.memory_matrix" in state_dict: + return "cognitive_expert" + if "layers.10.core_router.weight" in state_dict and "layers.10.cross_attn.qkv.weight" in state_dict: + return "cognitive_leap_ultra_expert" + if "layers.10.recur_core.0.weight" in state_dict and "layers.10.halt_head.weight" in state_dict: + return "cognitive_leap_expert" + if "layers.10.denoiser.0.0.weight" in state_dict and "layers.10.world_transition.0.weight" in state_dict: + return "transcendent_expert" + if "layers.10.ode_func.net.0.weight" in state_dict and "layers.10.holographic_memory" in state_dict: + return "omniversal_expert" + if "layers.10.logic_atoms" in state_dict and "layers.10.curvature" in state_dict: + return "fractal_expert" + if "layers.10.ttt_lora_A" in state_dict and "layers.10.lif_neuron.threshold" in state_dict: + return "liquid_spiking_expert" + if "layers.10.world_model.weight" in state_dict and "layers.10.precision_gate.weight" in state_dict: + return "active_inference_expert" + if "layers.10.holographic_memory" in state_dict and "layers.10.topology_router.weight" in state_dict: + return "holographic_state_space_expert" + if "layers.10.verifier_q.weight" in state_dict and "layers.10.revisit_gate.weight" in state_dict: + return "frontier_verifier_expert" + if "layers.10.shared_aux_up.weight" in state_dict and "layers.10.collective_q.weight" in state_dict: + return "frontier_collective_expert" + if "layers.10.router_budget.weight" in state_dict and "layers.10.memory_write_gate.weight" in state_dict: + return "frontier_expert" + if "layers.10.mor_router.main_router.weight" in state_dict and "layers.10.ssm_adapter.A" in state_dict: + return "paper_fusion_expert" + if "layers.10.adversary_up.0.weight" in state_dict and "layers.10.proposer_up.0.weight" in state_dict: + return "neurogenesis_expert" + if "layers.10.memory_keys" in state_dict and "layers.10.draft_gates.0.0.weight" in state_dict: + return "deliberative_expert" + if "layers.10.mlp_gate.weight" in state_dict and "layers.10.consensus.0.weight" in state_dict: + return "consensus_expert" + if "layers.10.critique_gate.weight" in state_dict: + return "reflexive_expert" + if "layers.10.exit_gates.0.weight" in state_dict: + return "recursive_expert" + if "layers.10.reasoning_cells.0.up.weight" in state_dict: + return "thought_expert" + + if "layers.10.gate.weight" in state_dict or "layers.10.experts_up.0.weight" in state_dict: + return "ultra_expert" + return "base" + + +def detect_large_head_expansion_dim(state_dict: dict, default: int = 512) -> int: + """Detect the primary adapter expansion dim from a checkpoint.""" + key = "layers.10.adapter_up.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def detect_xlarge_aux_expansion_dim(state_dict: dict, default: int = 1024) -> int: + """Detect the secondary (XL) adapter expansion dim from a checkpoint.""" + key = "layers.10.adapter_up_b.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def detect_xxlarge_third_expansion_dim(state_dict: dict, default: int = 3072) -> int: + """Detect the third (XXL) adapter expansion dim from a checkpoint.""" + key = "layers.10.adapter_up_c.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def detect_xxxlarge_fourth_expansion_dim(state_dict: dict, default: int = 4096) -> int: + """Detect the fourth (XXXL) adapter expansion dim from a checkpoint.""" + key = "layers.10.adapter_up_d.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def detect_ultralarge_fifth_expansion_dim(state_dict: dict, default: int = 6144) -> int: + """Detect the fifth (Ultra) adapter expansion dim from a checkpoint.""" + key = "layers.10.adapter_up_e.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def detect_megalarge_sixth_expansion_dim(state_dict: dict, default: int = 8192) -> int: + """Detect the sixth (Mega) adapter expansion dim from a checkpoint.""" + key = "layers.11.adapter_up_f.weight" + weight = state_dict.get(key) + if isinstance(weight, torch.Tensor) and weight.ndim == 2: + return int(weight.shape[0]) + return int(default) + + +def list_trainable_parameter_names(model: nn.Module) -> Sequence[str]: + """Return the names of all parameters with requires_grad=True.""" + return [name for name, p in model.named_parameters() if p.requires_grad] diff --git a/source/benchmark_cognitive_leap_ultra_v51.py b/source/benchmark_cognitive_leap_ultra_v51.py new file mode 100644 index 000000000..060d39eee --- /dev/null +++ b/source/benchmark_cognitive_leap_ultra_v51.py @@ -0,0 +1,253 @@ +"""Train v51 CognitiveLeapUltraExpert on a bounded synthetic reasoning task. + +This is a controlled, local training run for the new ultra recursive head. It +does not require a base chat checkpoint; it trains from scratch on a chained +modular-arithmetic task and writes a reusable checkpoint plus metrics. + +Usage: + python source/benchmark_cognitive_leap_ultra_v51.py --epochs 4 --train_size 3000 +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +SOURCE_DIR = Path(__file__).resolve().parent +if str(SOURCE_DIR) not in sys.path: + sys.path.append(str(SOURCE_DIR)) + +from device_utils import resolve_device +from model_variants import ChampionNetCognitiveLeapUltraExpert + + +N_CLASSES = 10 +N_OPS = 4 +IN_DIM = 128 + + +def make_chained_task(n: int, seed: int) -> tuple[torch.Tensor, torch.Tensor]: + gen = torch.Generator().manual_seed(seed) + x = torch.zeros(n, IN_DIM) + y = torch.zeros(n, dtype=torch.long) + starts = torch.randint(0, 10, (n,), generator=gen) + op_types = torch.randint(0, 3, (n, N_OPS), generator=gen) # add/mul/sub + operands = torch.randint(1, 10, (n, N_OPS), generator=gen) + for i in range(n): + x[i, starts[i]] = 1.0 + acc = int(starts[i]) + for k in range(N_OPS): + base = 10 + k * 24 + x[i, base + int(op_types[i, k])] = 1.0 + x[i, base + 3 + int(operands[i, k])] = 1.0 + operand = int(operands[i, k]) + op = int(op_types[i, k]) + if op == 0: + acc = (acc + operand) % 10 + elif op == 1: + acc = (acc * operand) % 10 + else: + acc = (acc - operand) % 10 + y[i] = acc + x = x + 0.01 * torch.randn(n, IN_DIM, generator=gen) + return x.unsqueeze(1), y + + +def force_alpha(model: nn.Module, value: float = 1.0) -> None: + for name, param in model.named_parameters(): + if name == "layers.10.alpha": + param.data.fill_(value) + + +def aux_loss(model: ChampionNetCognitiveLeapUltraExpert, targets: torch.Tensor) -> torch.Tensor: + head = model.layers[10] + return ( + 0.08 * head.deep_supervision_loss(targets) + + 0.04 * head.last_consistency_loss + + 0.005 * head.last_ponder_cost + - 0.002 * head.last_gating_entropy + ) + + +def train_model( + model: ChampionNetCognitiveLeapUltraExpert, + x_train: torch.Tensor, + y_train: torch.Tensor, + epochs: int, + batch_size: int, + lr: float, + device: torch.device, +) -> Dict[str, Any]: + model.train() + optim = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01) + n = x_train.shape[0] + history = [] + start = time.perf_counter() + for epoch in range(1, epochs + 1): + epoch_start = time.perf_counter() + perm = torch.randperm(n) + total_loss = 0.0 + total_correct = 0 + total_seen = 0 + for i in range(0, n, batch_size): + idx = perm[i : i + batch_size] + xb = x_train[idx].to(device) + yb = y_train[idx].to(device) + logits = model(xb).squeeze(1) + loss = F.cross_entropy(logits, yb) + aux_loss(model, yb) + optim.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optim.step() + batch = int(yb.shape[0]) + total_loss += float(loss.item()) * batch + total_correct += int((logits.argmax(dim=-1) == yb).sum().item()) + total_seen += batch + row = { + "epoch": epoch, + "train_loss": round(total_loss / max(1, total_seen), 4), + "train_accuracy": round(total_correct / max(1, total_seen), 4), + "seconds": round(time.perf_counter() - epoch_start, 2), + } + history.append(row) + print( + f"epoch {epoch:02d}/{epochs} " + f"loss={row['train_loss']:.4f} acc={row['train_accuracy']:.4f} " + f"time={row['seconds']:.2f}s" + ) + return {"history": history, "train_seconds": round(time.perf_counter() - start, 2)} + + +@torch.no_grad() +def evaluate( + model: ChampionNetCognitiveLeapUltraExpert, + x_test: torch.Tensor, + y_test: torch.Tensor, + device: torch.device, + **forward_kwargs: Any, +) -> Dict[str, float]: + model.eval() + x = x_test.to(device) + y = y_test.to(device) + start = time.perf_counter() + logits = model(x, **forward_kwargs).squeeze(1) + latency_ms = (time.perf_counter() - start) * 1000.0 + loss = F.cross_entropy(logits, y) + acc = (logits.argmax(dim=-1) == y).float().mean() + head = model.layers[10] + return { + "loss": round(float(loss.item()), 4), + "accuracy": round(float(acc.item()), 4), + "latency_ms": round(latency_ms, 2), + "cycles_used": round(float(head.last_cycles_used.item()), 2), + "gating_entropy": round(float(head.last_gating_entropy.item()), 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--epochs", type=int, default=4) + ap.add_argument("--train_size", type=int, default=3000) + ap.add_argument("--test_size", type=int, default=1000) + ap.add_argument("--batch_size", type=int, default=128) + ap.add_argument("--lr", type=float, default=8e-4) + ap.add_argument("--seed", type=int, default=51) + ap.add_argument("--device", default="auto") + ap.add_argument("--device_preference", default="cuda,npu,xpu,dml,mps,cpu") + ap.add_argument( + "--output_dir", + default=str(SOURCE_DIR.parent / "output" / "benchmark_v51_cognitive_leap_ultra"), + ) + args = ap.parse_args() + + torch.manual_seed(args.seed) + device, device_info = resolve_device(args.device, preference=args.device_preference) + x_train, y_train = make_chained_task(args.train_size, seed=args.seed) + x_test, y_test = make_chained_task(args.test_size, seed=args.seed + 1) + + model = ChampionNetCognitiveLeapUltraExpert() + force_alpha(model) + model.to(device) + params = sum(p.numel() for p in model.parameters()) + + print(f"Task: chained modular arithmetic ({N_OPS} ops, mod {N_CLASSES})") + print(f"Train {args.train_size} / test {args.test_size}, epochs={args.epochs}, batch={args.batch_size}") + print(f"Device: {device_info.get('resolved', device)}") + print(f"Parameters: {params:,}") + + train = train_model( + model, + x_train, + y_train, + epochs=max(1, int(args.epochs)), + batch_size=max(1, int(args.batch_size)), + lr=float(args.lr), + device=device, + ) + + eval_default = evaluate(model, x_test, y_test, device) + eval_cycles = { + str(cycles): evaluate(model, x_test, y_test, device, reasoning_cycles=cycles) + for cycles in (1, 3, 8) + } + eval_adaptive = evaluate( + model, + x_test, + y_test, + device, + reasoning_cycles=8, + adaptive_compute=True, + exit_tol=1e-4, + exit_entropy_threshold=0.0, + ) + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = out_dir / "cognitive_leap_ultra_v51_trained.pth" + metrics_path = out_dir / "benchmark_results.json" + torch.save({k: v.detach().cpu() for k, v in model.state_dict().items()}, checkpoint_path) + + results = { + "created_at": datetime.now(timezone.utc).isoformat(), + "task": "chained_modular_arithmetic", + "model_size": "cognitive_leap_ultra_expert", + "params": params, + "train_size": int(args.train_size), + "test_size": int(args.test_size), + "epochs": int(args.epochs), + "batch_size": int(args.batch_size), + "lr": float(args.lr), + "seed": int(args.seed), + "device": str(device_info.get("resolved", device)), + "checkpoint_path": str(checkpoint_path), + **train, + "eval_default": eval_default, + "test_time_scaling": eval_cycles, + "adaptive_compute": eval_adaptive, + } + metrics_path.write_text(json.dumps(results, indent=2), encoding="utf-8") + + print("\nFinal evaluation:") + print(f"default: acc={eval_default['accuracy']:.4f} loss={eval_default['loss']:.4f} lat={eval_default['latency_ms']:.2f}ms") + for cycles, row in eval_cycles.items(): + print(f"cycles={cycles}: acc={row['accuracy']:.4f} loss={row['loss']:.4f} lat={row['latency_ms']:.2f}ms") + print( + "adaptive(8): " + f"acc={eval_adaptive['accuracy']:.4f} loss={eval_adaptive['loss']:.4f} " + f"cycles={eval_adaptive['cycles_used']:.0f} lat={eval_adaptive['latency_ms']:.2f}ms" + ) + print(f"\nCheckpoint written to {checkpoint_path}") + print(f"Metrics written to {metrics_path}") + + +if __name__ == "__main__": + main() diff --git a/source/benchmark_v51_prediction_stability.py b/source/benchmark_v51_prediction_stability.py new file mode 100644 index 000000000..d3887fcd4 --- /dev/null +++ b/source/benchmark_v51_prediction_stability.py @@ -0,0 +1,229 @@ +"""Benchmark v51 prediction-stability exits under request-sized inference. + +This intentionally disables the older latent/entropy exits so the comparison +isolates the new full-output stability verifier against a fixed-cycle baseline. +""" + +from __future__ import annotations + +import argparse +from collections import Counter +from datetime import datetime, timezone +import json +from pathlib import Path +import time +from typing import Any, Dict, List + +import torch + +from benchmark_cognitive_leap_ultra_v51 import make_chained_task +from device_utils import configure_torch_runtime, resolve_device +from model_variants import ChampionNetCognitiveLeapUltraExpert + + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_ARTIFACT_DIR = PROJECT_ROOT / "output" / "benchmark_v51_cognitive_leap_ultra_latest" +DEFAULT_WEIGHTS = DEFAULT_ARTIFACT_DIR / "cognitive_leap_ultra_v51_trained.pth" +DEFAULT_OUTPUT = DEFAULT_ARTIFACT_DIR / "prediction_stability_results.json" + + +def _load_model(weights: Path, device: torch.device) -> ChampionNetCognitiveLeapUltraExpert: + state = torch.load(weights, map_location="cpu", weights_only=True) + if isinstance(state, dict) and isinstance(state.get("state_dict"), dict): + state = state["state_dict"] + model = ChampionNetCognitiveLeapUltraExpert().to(device).eval() + model.load_state_dict(state) + return model + + +def _latency_summary(values: List[float]) -> Dict[str, float]: + tensor = torch.tensor(values, dtype=torch.float64) + return { + "mean_ms": round(float(tensor.mean().item()), 3), + "p50_ms": round(float(torch.quantile(tensor, 0.50).item()), 3), + "p95_ms": round(float(torch.quantile(tensor, 0.95).item()), 3), + } + + +def _scalar_summary(values: List[float]) -> Dict[str, float]: + tensor = torch.tensor(values, dtype=torch.float64) + return { + "mean": round(float(tensor.mean().item()), 8), + "p50": round(float(torch.quantile(tensor, 0.50).item()), 8), + "p95": round(float(torch.quantile(tensor, 0.95).item()), 8), + "max": round(float(tensor.max().item()), 8), + } + + +@torch.no_grad() +def benchmark_serving_requests( + model: ChampionNetCognitiveLeapUltraExpert, + x: torch.Tensor, + y: torch.Tensor, + *, + fixed_cycles: int, + max_cycles: int, + stability_patience: int, + stability_tol: float, + distribution_top_k: int, +) -> Dict[str, Any]: + head = model.layers[10] + fixed_predictions: List[int] = [] + adaptive_predictions: List[int] = [] + fixed_latencies: List[float] = [] + adaptive_latencies: List[float] = [] + adaptive_cycles: List[float] = [] + adaptive_latest_js_drift: List[float] = [] + adaptive_max_js_drift: List[float] = [] + exit_reasons: Counter[str] = Counter() + + # Keep initialization and one-time kernel setup outside measured requests. + model(x[:1], reasoning_cycles=fixed_cycles) + model( + x[:1], + reasoning_cycles=max_cycles, + adaptive_compute=True, + exit_tol=0.0, + exit_entropy_threshold=0.0, + prediction_stability_patience=stability_patience, + prediction_stability_tol=stability_tol, + prediction_stability_top_k=distribution_top_k, + ) + + for index in range(len(y)): + sample = x[index : index + 1] + + started = time.perf_counter() + fixed_output = model(sample, reasoning_cycles=fixed_cycles) + fixed_latencies.append((time.perf_counter() - started) * 1000.0) + fixed_predictions.append(int(fixed_output.argmax(dim=-1).item())) + + started = time.perf_counter() + adaptive_output = model( + sample, + reasoning_cycles=max_cycles, + adaptive_compute=True, + exit_tol=0.0, + exit_entropy_threshold=0.0, + prediction_stability_patience=stability_patience, + prediction_stability_tol=stability_tol, + prediction_stability_top_k=distribution_top_k, + ) + adaptive_latencies.append((time.perf_counter() - started) * 1000.0) + adaptive_predictions.append(int(adaptive_output.argmax(dim=-1).item())) + adaptive_cycles.append(float(head.last_cycles_used.item())) + adaptive_latest_js_drift.append( + float(head.last_prediction_topk_js_divergence.item()) + ) + adaptive_max_js_drift.append( + float(head.last_prediction_topk_js_divergence_max.item()) + ) + exit_reasons[str(head.last_exit_reason)] += 1 + + truth = y.detach().cpu() + fixed_tensor = torch.tensor(fixed_predictions) + adaptive_tensor = torch.tensor(adaptive_predictions) + fixed_accuracy = float((fixed_tensor == truth).float().mean().item()) + adaptive_accuracy = float((adaptive_tensor == truth).float().mean().item()) + agreement = float((fixed_tensor == adaptive_tensor).float().mean().item()) + fixed_latency = _latency_summary(fixed_latencies) + adaptive_latency = _latency_summary(adaptive_latencies) + mean_cycles = sum(adaptive_cycles) / max(1, len(adaptive_cycles)) + + return { + "fixed": { + "cycles": int(fixed_cycles), + "accuracy": round(fixed_accuracy, 6), + "latency": fixed_latency, + }, + "prediction_stability": { + "max_cycles": int(max_cycles), + "patience": int(stability_patience), + "confidence_tolerance": float(stability_tol), + "accuracy": round(adaptive_accuracy, 6), + "mean_cycles_used": round(mean_cycles, 4), + "cycle_counts": { + str(cycle): count + for cycle, count in sorted(Counter(adaptive_cycles).items()) + }, + "exit_reasons": dict(sorted(exit_reasons.items())), + "distribution_drift": { + "metric": "top_k_jensen_shannon_divergence_nats", + "top_k": int(distribution_top_k), + "role": "shadow_diagnostic_only", + "latest_cycle_pair": _scalar_summary(adaptive_latest_js_drift), + "max_cycle_pair": _scalar_summary(adaptive_max_js_drift), + }, + "latency": adaptive_latency, + }, + "comparison": { + "accuracy_delta": round(adaptive_accuracy - fixed_accuracy, 6), + "prediction_agreement": round(agreement, 6), + "cycle_reduction_percent": round( + 100.0 * (float(fixed_cycles) - mean_cycles) / max(1.0, float(fixed_cycles)), + 3, + ), + "mean_latency_reduction_percent": round( + 100.0 + * (fixed_latency["mean_ms"] - adaptive_latency["mean_ms"]) + / max(1e-9, fixed_latency["mean_ms"]), + 3, + ), + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--weights", default=str(DEFAULT_WEIGHTS)) + parser.add_argument("--output", default=str(DEFAULT_OUTPUT)) + parser.add_argument("--samples", type=int, default=64) + parser.add_argument("--seed", type=int, default=52) + parser.add_argument("--fixed_cycles", type=int, default=3) + parser.add_argument("--max_cycles", type=int, default=8) + parser.add_argument("--stability_patience", type=int, default=2) + parser.add_argument("--stability_tol", type=float, default=0.005) + parser.add_argument("--distribution_top_k", type=int, default=5) + parser.add_argument("--device", default="auto") + parser.add_argument("--device_preference", default="cuda,npu,xpu,dml,mps,cpu") + parser.add_argument("--torch_num_threads", type=int, default=0) + args = parser.parse_args() + + configure_torch_runtime(torch_num_threads=int(args.torch_num_threads)) + device, device_info = resolve_device(args.device, preference=args.device_preference) + weights = Path(args.weights) + if not weights.exists(): + raise FileNotFoundError(f"Missing v51 checkpoint: {weights}") + + sample_count = max(1, int(args.samples)) + x, y = make_chained_task(sample_count, seed=int(args.seed)) + x = x.to(device) + y = y.to(device) + model = _load_model(weights, device) + result = benchmark_serving_requests( + model, + x, + y, + fixed_cycles=max(1, int(args.fixed_cycles)), + max_cycles=max(1, int(args.max_cycles)), + stability_patience=max(0, int(args.stability_patience)), + stability_tol=max(0.0, float(args.stability_tol)), + distribution_top_k=max(1, int(args.distribution_top_k)), + ) + payload = { + "created_at": datetime.now(timezone.utc).isoformat(), + "weights": str(weights), + "device": str(device_info.get("resolved", device)), + "samples": sample_count, + "seed": int(args.seed), + **result, + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(json.dumps(payload, indent=2)) + print(f"Results written to {output}") + + +if __name__ == "__main__": + main() diff --git a/source/materialize_v51_chat_demo.py b/source/materialize_v51_chat_demo.py new file mode 100644 index 000000000..ffcce26bd --- /dev/null +++ b/source/materialize_v51_chat_demo.py @@ -0,0 +1,206 @@ +"""Materialize and verify chat-demo metadata for the v51 ultra checkpoint. + +The v51 checkpoint produced by ``benchmark_cognitive_leap_ultra_v51.py`` is a +synthetic reasoning benchmark model, not a general chat fine-tune. This script +creates a small retrieval metadata file so the existing chat UI can load the +checkpoint and expose runtime-compute controls consistently. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import torch + +SOURCE_DIR = Path(__file__).resolve().parent +if str(SOURCE_DIR) not in sys.path: + sys.path.insert(0, str(SOURCE_DIR)) +PROJECT_ROOT = SOURCE_DIR.parent + +from chat_pipeline import featurize_text +from chat_web_app import Engine + + +DEFAULT_ARTIFACT_DIR = SOURCE_DIR.parent / "output" / "benchmark_v51_cognitive_leap_ultra" +DEFAULT_WEIGHTS = DEFAULT_ARTIFACT_DIR / "cognitive_leap_ultra_v51_trained.pth" +DEFAULT_METRICS = DEFAULT_ARTIFACT_DIR / "benchmark_results.json" +DEFAULT_META = DEFAULT_ARTIFACT_DIR / "chat_demo_meta.json" + + +PROMPT_RESPONSE_ROWS: Tuple[Tuple[str, str], ...] = ( + ( + "v51 ultra demo checkpoint status and limitations", + "This is the v51 Cognitive Leap Ultra demo checkpoint. It was trained on a small synthetic chained-arithmetic task, so treat this chat surface as a runtime demo rather than a polished assistant.", + ), + ( + "training results metrics checkpoint reload", + "Training finished successfully on CPU and produced a reloadable checkpoint plus metrics under output/benchmark_v51_cognitive_leap_ultra.", + ), + ( + "runtime compute reasoning cycles adaptive exit entropy", + "Runtime compute controls are available: set reasoning cycles, adaptive compute, exit tolerance, and entropy threshold, then compare cycles used and latency.", + ), + ( + "synthetic chained modular arithmetic task", + "The training task applies chained modular arithmetic operations. It is useful for testing recursive latent reasoning, not broad natural-language capability.", + ), + ( + "how to test the model in chat interface", + "Use this interface to verify loading, inference, candidate retrieval, and compute diagnostics. For serious chat quality, run a real chat fine-tune with conversation data.", + ), + ( + "explain confidence uncertainty for this demo", + "Uncertainty is high for broad chat prompts because this checkpoint was not trained on conversation data. Prefer asking about the run, metrics, or compute controls.", + ), + ( + "debug model response and inspect candidates", + "For debugging, enable top candidates and compare timing, cycles used, and selected responses. If output is poor, the limitation is training data, not the web interface.", + ), + ( + "next training step for real chat fine tune", + "The next useful training step is a bounded chat fine-tune from a compatible base checkpoint with real conversation manifests and recursive-head auxiliary losses enabled.", + ), + ( + "short concise answer mode", + "Short answer: the v51 model is loaded for demo inference, but it is not a full chat model yet.", + ), + ( + "technical analyst answer mode", + "Technical readout: v51 ultra uses a recurrent MoE latent core, cross-latent attention, ACT-style halting, and runtime-compute diagnostics.", + ), +) + + +def _load_metrics(path: Path) -> Dict[str, Any]: + if not path.exists(): + return {} + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + + +def _portable_project_path(path: Path) -> str: + resolved = Path(path).resolve() + try: + return resolved.relative_to(PROJECT_ROOT.resolve()).as_posix() + except ValueError: + return str(resolved) + + +def build_demo_metadata(weights: Path, metrics_path: Path) -> Dict[str, Any]: + metrics = _load_metrics(metrics_path) + stability_metrics_path = metrics_path.with_name("prediction_stability_results.json") + stability_metrics = _load_metrics(stability_metrics_path) + buckets: Dict[str, List[Dict[str, Any]]] = {} + for label, (prompt, response) in enumerate(PROMPT_RESPONSE_ROWS): + buckets[str(label)] = [ + { + "text": response, + "vec": featurize_text(response).tolist(), + "ctx_vec": featurize_text(prompt).tolist(), + "count": 1, + } + ] + + return { + "created_by": "materialize_v51_chat_demo.py", + "purpose": "demo metadata for v51 cognitive_leap_ultra_expert chat_web_app", + "model_size": "cognitive_leap_ultra_expert", + "feature_mode": "context_mix_v4", + "num_classes": 10, + "training_task": "synthetic chained modular arithmetic benchmark", + "fine_tuned_weights": str(weights.name), + "checkpoint_path": _portable_project_path(weights), + "benchmark_metrics": _portable_project_path(metrics_path) if metrics else None, + "prediction_stability_metrics": ( + _portable_project_path(stability_metrics_path) if stability_metrics else None + ), + "benchmark_summary": { + "train_seconds": metrics.get("train_seconds"), + "eval_default": metrics.get("eval_default"), + "test_time_scaling": metrics.get("test_time_scaling"), + "adaptive_compute": metrics.get("adaptive_compute"), + "prediction_stability": stability_metrics.get("prediction_stability"), + "prediction_stability_comparison": stability_metrics.get("comparison"), + }, + "runtime_defaults": { + "reasoning_cycles": 3, + "adaptive_compute": True, + "adaptive_exit_tol": 0.001, + "adaptive_exit_entropy": 0.2, + "prediction_stability_patience": 2, + "prediction_stability_tol": 0.005, + }, + "buckets": buckets, + "label_priors": {str(i): 0.1 for i in range(10)}, + } + + +def materialize(weights: Path, metrics: Path, meta: Path) -> Path: + payload = build_demo_metadata(weights=weights, metrics_path=metrics) + meta.parent.mkdir(parents=True, exist_ok=True) + meta.write_text(json.dumps(payload, indent=2), encoding="utf-8") + return meta + + +def check_load(weights: Path, meta: Path) -> Dict[str, Any]: + engine = Engine( + torch.device("cpu"), + {"resolved": "cpu"}, + { + "model_size": "auto", + "pool_mode": "all", + "reasoning_cycles": 3, + "adaptive_compute": True, + "adaptive_exit_entropy": 0.2, + "prediction_stability_patience": 2, + "prediction_stability_tol": 0.005, + }, + ) + status = engine.load(str(weights), str(meta)) + chat = engine.chat( + session_id="v51-check", + user_text="how should I test this v51 demo?", + show_top_responses=3, + reasoning_cycles=3, + adaptive_compute=True, + adaptive_exit_entropy=0.2, + prediction_stability_patience=2, + prediction_stability_tol=0.005, + ) + return { + "loaded": status, + "chat_response": chat.get("response"), + "compute": chat.get("compute"), + "top_candidates": chat.get("top_candidates"), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(DEFAULT_WEIGHTS)) + ap.add_argument("--metrics", default=str(DEFAULT_METRICS)) + ap.add_argument("--meta", default=str(DEFAULT_META)) + ap.add_argument("--check", action="store_true", help="Load the checkpoint and run a one-turn smoke chat.") + args = ap.parse_args() + + weights = Path(args.weights) + metrics = Path(args.metrics) + meta = Path(args.meta) + if not weights.exists(): + raise FileNotFoundError(f"Missing v51 checkpoint: {weights}") + + materialize(weights=weights, metrics=metrics, meta=meta) + print(f"Wrote demo metadata: {meta}") + if args.check: + result = check_load(weights=weights, meta=meta) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/source/model_variants.py b/source/model_variants.py index f3a7705dc..8fde3df37 100644 --- a/source/model_variants.py +++ b/source/model_variants.py @@ -53,6 +53,7 @@ "neurogenesis_expert", "cognitive_expert", "cognitive_leap_expert", + "cognitive_leap_ultra_expert", "transcendent_expert", "omniversal_expert", "fractal_expert", @@ -702,7 +703,7 @@ def __init__( nn.Linear(1024 + (i * 512), out_dim, bias=False) for i in range(n_experts) ]) self.activations = [F.silu, F.gelu, F.mish, F.relu, F.selu, torch.tanh] - + self.gate = nn.Linear(in_dim, n_experts, bias=False) self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) self.alpha = nn.Parameter(torch.tensor(0.0)) @@ -724,7 +725,7 @@ def reset_parameters(self): def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) - + # Gating logic with Noisy Top-K clean_logits = self.gate(x) if self.training: @@ -744,15 +745,15 @@ def forward(self, x): out = self.experts_down[i](self.dropout(self.activations[i](self.experts_up[i](x)))) expert_outputs.append(out) expert_stack = torch.stack(expert_outputs, dim=-2) # (B, T, N, D) - + # Gather top experts # We need to broadcast the indices for gather gather_idx = top_idx.unsqueeze(-1).expand(*top_idx.shape, base_logits.shape[-1]) selected_experts = torch.gather(expert_stack, -2, gather_idx) # (B, T, K, D) - + expert_logits = (selected_experts * top_weights.unsqueeze(-1)).sum(dim=-2) calib = self.calibration(x) - + return base_logits + self.alpha * expert_logits + self.theta * calib @@ -1035,12 +1036,12 @@ def __init__( # Routing mechanism self.gate = nn.Linear(in_dim, n_experts, bias=False) self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) - + # Aux-free load balancing terms # These biases are added to the routing logits during training. # They don't require gradients; we update them manually. self.register_buffer("expert_bias", torch.zeros(n_experts)) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.calibration = nn.Linear(in_dim, out_dim, bias=True) self.theta = nn.Parameter(torch.tensor(0.0)) @@ -1145,11 +1146,11 @@ def __init__( super().__init__() self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) - + self.n_experts = n_experts # How many tokens an expert will pick, relative to average token load per expert self.capacity_factor = capacity_factor - + # Routed Experts self.experts_up = nn.ModuleList([ nn.Linear(in_dim, 1024 + (i * 256), bias=False) for i in range(n_experts) @@ -1162,10 +1163,10 @@ def __init__( self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) # Routing mechanism: Token-to-expert affinity - # We need a score for each (token, expert) pair. + # We need a score for each (token, expert) pair. self.gate = nn.Linear(in_dim, n_experts, bias=False) self.noise_gate = nn.Linear(in_dim, n_experts, bias=False) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.calibration = nn.Linear(in_dim, out_dim, bias=True) self.theta = nn.Parameter(torch.tensor(0.0)) @@ -1215,33 +1216,33 @@ def forward(self, x): expert_to_token_probs = token_to_expert_probs.t() # (n_experts, BT) expert_outputs = torch.zeros((BT, base_logits.shape[-1]), device=x.device, dtype=x.dtype) - - # We need a scaling factor later. If an expert picks a token, its output is scaled by + + # We need a scaling factor later. If an expert picks a token, its output is scaled by # the token_to_expert_prob. - + # 2. Each expert processes its top tokens for i in range(self.n_experts): probs_i = expert_to_token_probs[i] # (BT,) - + # Select top-k tokens for this expert top_probs, top_token_idx = torch.topk(probs_i, k=expert_capacity, dim=0) # (capacity,) - + # Gather tokens selected_tokens = x_flat[top_token_idx] # (capacity, D) - + # Process through expert i act = self.activations[i] out = self.experts_down[i](self.dropout(act(self.experts_up[i](selected_tokens)))) out = self.expert_norms[i](out) # (capacity, out_dim) - + # Scale by routing probability out = out * top_probs.unsqueeze(-1) # (capacity, out_dim) - + # Add to the global output buffer using `scatter_add_` or index_add - # out is (capacity, out_dim). top_token_idx is (capacity,). + # out is (capacity, out_dim). top_token_idx is (capacity,). # expert_outputs is (BT, out_dim). expert_outputs.index_add_(0, top_token_idx, out) - + expert_outputs = expert_outputs.view(*shape_prefix, -1) calib = self.calibration(x) @@ -1302,7 +1303,7 @@ def __init__( # 4. Routing Mechanism (Sigma Gating) self.gate = nn.Linear(in_dim, n_experts, bias=False) self.register_buffer("expert_bias", torch.zeros(n_experts)) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.calibration = nn.Linear(in_dim, out_dim, bias=True) self.theta = nn.Parameter(torch.tensor(0.0)) @@ -1315,16 +1316,16 @@ def reset_parameters(self): if fan_in != 0: bound = 1 / math.sqrt(fan_in) nn.init.uniform_(self.bias, -bound, bound) - + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.zeros_(self.shared_down.weight) - + for i in range(self.n_experts): nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) nn.init.zeros_(self.experts_down[i].weight) nn.init.normal_(self.lora_down[i].weight, std=0.02) nn.init.zeros_(self.lora_up[i].weight) - + nn.init.normal_(self.gate.weight, std=0.02) nn.init.zeros_(self.calibration.weight) nn.init.zeros_(self.calibration.bias) @@ -1344,9 +1345,9 @@ def forward(self, x): gate_logits = clean_logits + self.expert_bias else: gate_logits = clean_logits - + # Sigma activation - gate_scores = torch.sigmoid(gate_logits) + gate_scores = torch.sigmoid(gate_logits) all_expert_logits = [] for i in range(self.n_experts): @@ -1354,9 +1355,9 @@ def forward(self, x): lora = self.lora_up[i](self.lora_down[i](x_flat)) out = self.expert_norms[i](core + lora) all_expert_logits.append(out) - - expert_stack = torch.stack(all_expert_logits, dim=1) - routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) + + expert_stack = torch.stack(all_expert_logits, dim=1) + routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) if self.training: with torch.no_grad(): @@ -1392,28 +1393,28 @@ def __init__(self, out_dim: int, n_branches: int, n_heads: int = 2): self.k_proj = nn.Linear(out_dim, out_dim, bias=False) self.v_proj = nn.Linear(out_dim, out_dim, bias=False) self.out_proj = nn.Linear(out_dim, out_dim, bias=False) - + self.fusion_weight = nn.Linear(out_dim, 1) def forward(self, x): # x: (B, T, N, D) B, T, N, D = x.shape x_flat = x.view(B * T, N, D) - + q = self.q_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x_flat).view(B * T, N, self.n_heads, self.head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) attn = F.softmax(attn, dim=-1) - + out = (attn @ v).transpose(1, 2).reshape(B * T, N, D) out = self.out_proj(out) - + # Weighted reduction to single D-dim vector per token scores = torch.sigmoid(self.fusion_weight(out)) # (B*T, N, 1) fused = (out * scores).sum(dim=1) / (scores.sum(dim=1) + 1e-6) - + return fused.view(B, T, D) @@ -1489,7 +1490,7 @@ def __init__( self.reasoning_cells = nn.ModuleList([ ReasoningCell(in_dim, inner_dim=512, dropout=dropout) for _ in range(reasoning_steps) ]) - + # 5. Iterative Routing (Sigma Gating) self.gates = nn.ModuleList([ nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps) @@ -1498,7 +1499,7 @@ def __init__( # 6. Cross-Expert Attention Fusion self.cross_attn = CrossAttentionFusion(out_dim=out_dim, n_branches=n_experts, n_heads=2) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.calibration = nn.Linear(in_dim, out_dim, bias=True) self.theta = nn.Parameter(torch.tensor(0.0)) @@ -1511,19 +1512,19 @@ def reset_parameters(self): if fan_in != 0: bound = 1 / math.sqrt(fan_in) nn.init.uniform_(self.bias, -bound, bound) - + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.zeros_(self.shared_down.weight) - + for i in range(self.n_experts): nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) nn.init.zeros_(self.experts_down[i].weight) nn.init.normal_(self.lora_down[i].weight, std=0.02) nn.init.zeros_(self.lora_up[i].weight) - + for g in self.gates: nn.init.normal_(g.weight, std=0.02) - + nn.init.zeros_(self.calibration.weight) nn.init.zeros_(self.calibration.bias) @@ -1531,7 +1532,7 @@ def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) shape_prefix = x.shape[:-1] x_flat = x.reshape(-1, x.shape[-1]) - + # Shared Anchor processing shared_out = self.shared_norm(self.shared_down(self.dropout(F.silu(self.shared_up(x_flat))))) @@ -1542,30 +1543,30 @@ def forward(self, x): lora = self.lora_up[i](self.lora_down[i](x_flat)) out = self.expert_norms[i](core + lora) all_expert_logits.append(out) - expert_stack = torch.stack(all_expert_logits, dim=1) + expert_stack = torch.stack(all_expert_logits, dim=1) # Iterative Reasoning Loop current_features = x_flat total_routed_out = 0 - + for step in range(self.reasoning_steps): # 1. Refine features current_features = self.reasoning_cells[step](current_features) - + # 2. Gate experts based on refined features gate_logits = self.gates[step](current_features) if self.training: gate_logits = gate_logits + self.expert_bias gate_scores = torch.sigmoid(gate_logits) # (B*T, n_experts) - + # 3. Fuse experts for this step step_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) - + # 4. Attention-based cross-expert correction (residual) # Re-weight experts with attention to find hidden correlations # Shape: expert_stack is (BT, N, D), we need (BT, 1, N, D) for CrossAttentionFusion attn_fused = self.cross_attn(expert_stack.unsqueeze(1)).squeeze(1) - + total_routed_out = total_routed_out + step_out + 0.1 * attn_fused if self.training: @@ -1650,7 +1651,7 @@ def __init__( self.exit_gates = nn.ModuleList([ nn.Linear(in_dim, 1) for _ in range(reasoning_steps) ]) - + # 5. Multi-Head Sigma Gating (Flattened for registration) self.gates = nn.ModuleList([ nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps * n_heads) @@ -1659,7 +1660,7 @@ def __init__( # 6. Cross-Expert Attention Fusion self.cross_attn = CrossAttentionFusion(out_dim=out_dim, n_branches=n_experts, n_heads=2) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.calibration = nn.Linear(in_dim, out_dim, bias=True) self.theta = nn.Parameter(torch.tensor(0.0)) @@ -1672,21 +1673,21 @@ def reset_parameters(self): if fan_in != 0: bound = 1 / math.sqrt(fan_in) nn.init.uniform_(self.bias, -bound, bound) - + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.zeros_(self.shared_down.weight) nn.init.kaiming_uniform_(self.local_up.weight, a=math.sqrt(5)) nn.init.zeros_(self.local_down.weight) - + for i in range(self.n_experts): nn.init.kaiming_uniform_(self.experts_up[i].weight, a=math.sqrt(5)) nn.init.normal_(self.experts_down[i].weight, std=0.01) # Small initial signal nn.init.normal_(self.lora_down[i].weight, std=0.02) nn.init.zeros_(self.lora_up[i].weight) - + for g in self.gates: nn.init.normal_(g.weight, std=0.02) - + for e in self.exit_gates: nn.init.constant_(e.bias, -3.0) # Start with low exit prob @@ -1697,7 +1698,7 @@ def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) shape_prefix = x.shape[:-1] x_flat = x.reshape(-1, x.shape[-1]) - + # 1. Global Shared Expert shared_out = self.shared_norm(self.shared_down(self.dropout(F.silu(self.shared_up(x_flat))))) # 2. Local Shared Expert @@ -1710,19 +1711,19 @@ def forward(self, x): lora = self.lora_up[i](self.lora_down[i](x_flat)) out = self.expert_norms[i](core + lora) all_expert_logits.append(out) - expert_stack = torch.stack(all_expert_logits, dim=1) + expert_stack = torch.stack(all_expert_logits, dim=1) # Iterative Reasoning Loop with ACE current_features = x_flat total_routed_out = torch.zeros_like(expert_stack[:, 0, :]) cumulative_exit_prob = torch.zeros(x_flat.shape[0], 1, device=x.device) depth_count = 0 - + for step in range(self.reasoning_steps): depth_count += 1 # A. Refine features current_features = self.reasoning_cells[step](current_features) - + # B. Multi-Head Sigma Gating head_logits = [] for h in range(self.n_heads): @@ -1730,24 +1731,24 @@ def forward(self, x): head_logits.append(self.gates[idx](current_features)) # Average head logits (Sigma Gating) gate_logits = torch.stack(head_logits, dim=0).mean(dim=0) - + if self.training: gate_logits = gate_logits + self.expert_bias gate_scores = torch.sigmoid(gate_logits) # (B*T, n_experts) - + # C. ACE: Early Exit check exit_logit = self.exit_gates[step](current_features) exit_prob = torch.sigmoid(exit_logit) # (B*T, 1) - + # D. Fuse experts for this step step_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) - + # E. Attention-based cross-expert correction attn_fused = self.cross_attn(expert_stack.unsqueeze(1)).squeeze(1) - + # Weighted contribution based on current depth total_routed_out = total_routed_out + (1.0 - cumulative_exit_prob) * (step_out + 0.1 * attn_fused) - + # Update cumulative exit probability cumulative_exit_prob = cumulative_exit_prob + (1.0 - cumulative_exit_prob) * exit_prob @@ -1756,7 +1757,7 @@ def forward(self, x): target_load = 1.0 / self.n_experts actual_load = gate_scores.mean(dim=0) self.expert_bias.add_(0.01 * (target_load - actual_load)) - + # Early exit condition if not self.training and cumulative_exit_prob.mean() > 0.9: break @@ -1785,10 +1786,10 @@ def __init__(self, n_experts: int = 8, reasoning_steps: int = 3, lora_rank: int base = ChampionNet() layers = [base.layers[i] for i in range(10)] layers.append(RecursiveThoughtExpertHead( - 256, 10, - n_experts=n_experts, - reasoning_steps=reasoning_steps, - lora_rank=lora_rank, + 256, 10, + n_experts=n_experts, + reasoning_steps=reasoning_steps, + lora_rank=lora_rank, dropout=dropout )) layers.append(base.layers[11]) @@ -1851,7 +1852,7 @@ def reset_parameters(self): def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) - + # --- Initial Pass --- gate_init = torch.softmax(self.initial_gate(x), dim=-1) # (B, T, E) init_outs = [] @@ -1860,12 +1861,12 @@ def forward(self, x): init_outs.append(self.initial_norms[i](self.initial_down[i](h))) init_stack = torch.stack(init_outs, dim=-2) # (B, T, E, D) init_fused = (init_stack * gate_init.unsqueeze(-1)).sum(dim=-2) - + # Intermediate prediction inter_logits = base_logits + self.alpha * init_fused - + # --- Critique Pass --- - crit_in = torch.cat([x, inter_logits], dim=-1) + crit_in = torch.cat([x, inter_logits], dim=-1) gate_crit = torch.softmax(self.critique_gate(crit_in), dim=-1) crit_outs = [] for i in range(self.n_experts): @@ -1873,7 +1874,7 @@ def forward(self, x): crit_outs.append(self.critique_norms[i](self.critique_down[i](h))) crit_stack = torch.stack(crit_outs, dim=-2) crit_fused = (crit_stack * gate_crit.unsqueeze(-1)).sum(dim=-2) - + # Final prediction return inter_logits + self.beta * crit_fused @@ -2153,17 +2154,17 @@ def __init__( self.n_action_experts = n_action_experts self.beam_size = beam_size self.reasoning_steps = reasoning_steps - + # Base projection self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) - + # 1. Global Shared Expert (always-on) self.shared_up = nn.Linear(in_dim, 2048, bias=False) self.shared_down = nn.Linear(2048, out_dim, bias=False) self.shared_norm = nn.LayerNorm(out_dim) self.shared_scale = nn.Parameter(torch.tensor(1.0)) - + # 2. Value Network (Scorer) # Evaluates the "goodness" of a latent state. self.value_net = nn.Sequential( @@ -2172,7 +2173,7 @@ def __init__( nn.Dropout(dropout), nn.Linear(512, 1) ) - + # 3. Action MoE (Proposer) # We use a set of experts that propose feature modifications self.action_up = nn.ModuleList([ @@ -2184,12 +2185,12 @@ def __init__( self.action_norms = nn.ModuleList([ nn.LayerNorm(in_dim) for _ in range(n_action_experts) ]) - + # Final projection from feature to out_dim self.final_up = nn.Linear(in_dim, 1024, bias=False) self.final_down = nn.Linear(1024, out_dim, bias=False) self.final_norm = nn.LayerNorm(out_dim) - + self.alpha = nn.Parameter(torch.tensor(0.0)) # tree scale self.dropout = nn.Dropout(dropout) self.reset_parameters() @@ -2200,20 +2201,20 @@ def reset_parameters(self): if fan_in != 0: bound = 1 / math.sqrt(fan_in) nn.init.uniform_(self.bias, -bound, bound) - + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.normal_(self.shared_down.weight, std=0.01) - + for p in self.value_net.parameters(): if p.dim() > 1: nn.init.kaiming_uniform_(p, a=math.sqrt(5)) else: nn.init.zeros_(p) - + for i in range(self.n_action_experts): nn.init.kaiming_uniform_(self.action_up[i].weight, a=math.sqrt(5)) nn.init.normal_(self.action_down[i].weight, std=0.01) - + nn.init.kaiming_uniform_(self.final_up.weight, a=math.sqrt(5)) nn.init.normal_(self.final_down.weight, std=0.01) @@ -2221,65 +2222,65 @@ def forward(self, x): N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] shape_prefix = x.shape[:-1] x_flat = x.reshape(N, self.in_dim) - + base_logits = F.linear(x_flat, self.weight, self.bias) - + # 1. Shared Expert shared_out = self.shared_norm( self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) ) - + # 2. Tree-of-Thought Search # Initial beam state: K=1 beam_states = x_flat.unsqueeze(1) # (N, 1, D) - + for step in range(self.reasoning_steps): K = beam_states.shape[1] - + # A. Expansion: Each action expert proposes a modification to every state in the beam nk_states = beam_states.reshape(N * K, self.in_dim) - + candidates = [] for i in range(self.n_action_experts): h = self.dropout(F.gelu(self.action_up[i](nk_states))) delta = self.action_norms[i](self.action_down[i](h)) new_state = nk_states + delta candidates.append(new_state) - + # Stack candidates cand_stack = torch.stack(candidates, dim=1) cand_stack = cand_stack.reshape(N, K * self.n_action_experts, self.in_dim) - + # B. Evaluation (Value Network) flat_cands = cand_stack.reshape(N * K * self.n_action_experts, self.in_dim) scores = self.value_net(flat_cands).squeeze(-1) # (N * K * M) scores = scores.reshape(N, K * self.n_action_experts) # (N, K*M) - + # C. Selection (Beam Search) K_next = min(self.beam_size, K * self.n_action_experts) - + topk_scores, topk_indices = torch.topk(scores, k=K_next, dim=1) # (N, K_next) - + expanded_indices = topk_indices.unsqueeze(-1).expand(-1, -1, self.in_dim) beam_states = torch.gather(cand_stack, 1, expanded_indices) # (N, K_next, D) - + # 3. Aggregation across the final beam final_K = beam_states.shape[1] flat_final = beam_states.reshape(N * final_K, self.in_dim) - + h_final = self.dropout(F.silu(self.final_up(flat_final))) final_preds = self.final_norm(self.final_down(h_final)) # (N * final_K, out_dim) final_preds = final_preds.reshape(N, final_K, self.out_dim) # (N, K, out_dim) - + final_scores = self.value_net(flat_final).squeeze(-1).reshape(N, final_K) # (N, K) agg_weights = torch.softmax(final_scores, dim=1) # (N, K) - + tree_fused = (final_preds * agg_weights.unsqueeze(-1)).sum(dim=1) - + base_logits = base_logits.view(*shape_prefix, -1) shared_out = shared_out.view(*shape_prefix, -1) tree_fused = tree_fused.view(*shape_prefix, -1) - + return ( base_logits + self.shared_scale * shared_out @@ -2743,7 +2744,7 @@ def forward(self, x): write_scores = torch.sigmoid(self.mem_write_gate(write_q)) # (N, M) # Differentiable soft memory update: weighted retrieval from write signal mem_retrieval = torch.matmul(write_scores, self.memory_values) # (N, D) - # Inject memory write contribution into draft features for next step + # Inject memory write contribution into draft features for next step for d in range(self.n_drafts): draft_features[d] = draft_features[d] + 0.1 * mem_retrieval @@ -2897,7 +2898,7 @@ def __init__( ]) self.expert_acts = [expert_acts_fns[i % len(expert_acts_fns)] for i in range(n_experts)] self.expert_norms = nn.ModuleList([nn.LayerNorm(out_dim) for _ in range(n_experts)]) - + # Bayesian Routers (predicts mean and log-variance for expert weighting) self.route_mu = nn.ModuleList([ nn.Linear(in_dim, n_experts, bias=False) for _ in range(reasoning_steps) @@ -2920,7 +2921,7 @@ def __init__( self.node_halt_gates = nn.ModuleList([ nn.Linear(in_dim, 1, bias=True) for _ in range(reasoning_steps) ]) - + # Final Node Fusion self.node_fusion_q = nn.Linear(out_dim, out_dim, bias=False) self.node_fusion_k = nn.Linear(out_dim, out_dim, bias=False) @@ -2928,7 +2929,7 @@ def __init__( self.alpha = nn.Parameter(torch.tensor(0.0)) self.dropout = nn.Dropout(dropout) - + # Divergence temperature for Bayesian routing during inference self.register_buffer("temperature", torch.tensor(0.2)) @@ -2954,17 +2955,17 @@ def reset_parameters(self): nn.init.kaiming_uniform_(self.got_q[step].weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.got_k[step].weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.got_k[step].weight, a=math.sqrt(5)) - + nn.init.normal_(self.route_mu[step].weight, std=0.02) nn.init.normal_(self.route_logvar[step].weight, std=0.02) nn.init.constant_(self.route_logvar[step].bias, -2.0) # start with low variance - + for p in self.critique_net[step].parameters(): if p.dim() > 1: nn.init.kaiming_uniform_(p, a=math.sqrt(5)) else: nn.init.zeros_(p) - + nn.init.constant_(self.node_halt_gates[step].bias, -2.0) - + nn.init.kaiming_uniform_(self.node_fusion_q.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.node_fusion_k.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.node_fusion_v.weight, a=math.sqrt(5)) @@ -2994,7 +2995,7 @@ def forward(self, x): for i in range(self.n_nodes): noise = torch.randn_like(x_flat) * 0.05 nodes.append(x_flat + knowledge_context + noise) - + node_stack = torch.stack(nodes, dim=1) # (N, V, D) node_cumulative_halt = torch.zeros(N, self.n_nodes, 1, device=x.device) total_node_preds = torch.zeros(N, self.n_nodes, self.out_dim, device=x.device) @@ -3009,42 +3010,42 @@ def forward(self, x): kl_loss = 0.0 for step in range(self.reasoning_steps): - + # A. Graph Attention: Nodes exchange information q_got = self.got_q[step](node_stack) # (N, V, D) k_got = self.got_k[step](node_stack) # (N, V, D) v_got = self.got_v[step](node_stack) # (N, V, D) - + got_attn = torch.bmm(q_got, k_got.transpose(1, 2)) * self.knowledge_scale # (N, V, V) got_attn = self.dropout(torch.softmax(got_attn, dim=-1)) got_updates = torch.bmm(got_attn, v_got) # (N, V, D) - + node_stack = self.got_norm[step](node_stack + got_updates) - + # B. Meta-Cognitive Critique # Global view of the graph spots logic errors, broadcasts correction global_state = node_stack.reshape(N, -1) # (N, V*D) global_correction = self.critique_net[step](global_state) # (N, D) - + # Apply correction to all active nodes node_stack = node_stack + global_correction.unsqueeze(1) - + # C. Stochastic Bayesian MoE Routing & Halting step_node_preds = [] step_halt_probs = [] - + for v in range(self.n_nodes): node_feat = node_stack[:, v, :] # (N, D) - + # Halt eval halt_prob = torch.sigmoid(self.node_halt_gates[step](node_feat)) # (N, 1) step_halt_probs.append(halt_prob) - - # Bayesian Routing + + # Bayesian Routing mu = self.route_mu[step](node_feat) # (N, E) logvar = self.route_logvar[step](node_feat) # (N, E) var = torch.exp(logvar) - + if self.training: # Reparameterization trick: sample routing weights eps = torch.randn_like(mu) @@ -3057,7 +3058,7 @@ def forward(self, x): sample = mu + eps * torch.sqrt(var) * self.temperature gate_scores = torch.sigmoid(sample) # (N, E) - + # Execute experts and fusion routed = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=1) # (N, D_out) step_node_preds.append(routed) @@ -3065,14 +3066,14 @@ def forward(self, x): # D. Accumulate Predictions step_preds_stack = torch.stack(step_node_preds, dim=1) # (N, V, D_out) step_halts_stack = torch.stack(step_halt_probs, dim=1) # (N, V, 1) - + remaining = 1.0 - node_cumulative_halt total_node_preds = total_node_preds + remaining * step_preds_stack node_cumulative_halt = node_cumulative_halt + remaining * step_halts_stack - + if not self.training and node_cumulative_halt.mean() > 0.95: break - + # Optional: return KL loss to be added to total loss via an attribute self.last_kl_loss = kl_loss * 0.001 @@ -3080,7 +3081,7 @@ def forward(self, x): fusion_q = self.node_fusion_q(total_node_preds).mean(dim=1, keepdim=True) # (N, 1, D_out) fusion_k = self.node_fusion_k(total_node_preds) # (N, V, D_out) fusion_v = self.node_fusion_v(total_node_preds) # (N, V, D_out) - + pool_attn = torch.bmm(fusion_q, fusion_k.transpose(1, 2)) * (self.out_dim**-0.5) # (N, 1, V) pool_attn = torch.softmax(pool_attn, dim=-1) fused_got_out = torch.bmm(pool_attn, fusion_v).squeeze(1) # (N, D_out) @@ -3513,7 +3514,7 @@ def __init__( self.path_projs = nn.ModuleList([ nn.Linear(in_dim, in_dim, bias=False) for _ in range(n_hypotheses) ]) - + # 4. Latent Tree-Search Value Network # Evaluates the "goodness" or logical consistency of a hypothesis self.value_net = nn.Sequential( @@ -3530,7 +3531,7 @@ def __init__( self.alpha = nn.Parameter(torch.tensor(0.0)) self.dropout = nn.Dropout(dropout) - + # Track orthogonal penalty loss for debugging/regularization self.register_buffer("last_ortho_loss", torch.tensor(0.0)) @@ -3559,7 +3560,7 @@ def reset_parameters(self): nn.init.normal_(layer.weight, std=0.01) if layer.bias is not None: nn.init.zeros_(layer.bias) - + # Output of hypernetwork should be initialized small to prevent explosion nn.init.normal_(self.hypernet[-1].weight, std=0.001) @@ -3605,18 +3606,18 @@ def forward(self, x): attn = torch.softmax(attn, dim=-1) # (N, M) mem_retrieval = torch.matmul(attn, self.memory_matrix) # (N, mem_dim) mem_context = self.mem_read_out(mem_retrieval) # (N, in_dim) - + # 3. Hypernetwork Weight Generation # Context is x_flat + episodic memory recall hyper_input = torch.cat([x_flat, mem_context], dim=-1) # (N, 2*in_dim) generated_weights = self.hypernet(hyper_input) # (N, W) - + w1_size = self.in_dim * self.hyper_dim w2_size = self.hyper_dim * self.out_dim - + w1_flat = generated_weights[:, :w1_size] w2_flat = generated_weights[:, w1_size:] - + # Reshape generated weights: (N, hyper_dim, in_dim) and (N, out_dim, hyper_dim) W1 = w1_flat.view(N, self.hyper_dim, self.in_dim) W2 = w2_flat.view(N, self.out_dim, self.hyper_dim) @@ -3628,14 +3629,14 @@ def forward(self, x): for k in range(self.n_hypotheses): h_in = self.path_projs[k](x_flat) # (N, D) hypotheses.append(h_in) - + # Execute dynamically generated expert layer 1: (N, hyper_dim) h_hidden = self.dropout(F.gelu(torch.bmm(W1, h_in.unsqueeze(-1)).squeeze(-1))) - + # Execute dynamically generated expert layer 2: (N, out_dim) h_out = torch.bmm(W2, h_hidden.unsqueeze(-1)).squeeze(-1) custom_outs.append(h_out) - + hypotheses_stack = torch.stack(hypotheses, dim=1) # (N, K, D) custom_outs_stack = torch.stack(custom_outs, dim=1) # (N, K, D_out) @@ -3648,40 +3649,40 @@ def forward(self, x): values = self.value_net(custom_outs_stack).squeeze(-1) # (N, K) # Select best hypotheses (softmax converts to probabilities) path_probs = torch.softmax(values, dim=-1) # (N, K) - + # 6. Causal Logic Refinement # Enforce that the top hypotheses are logically coherent sequences causal_in = custom_outs_stack # (N, K, D_out) q = self.causal_q(causal_in) k = self.causal_k(causal_in) v = self.causal_v(causal_in) - + # Masked attention (causal path flow) causal_attn = torch.bmm(q, k.transpose(1, 2)) / math.sqrt(self.out_dim) mask = torch.tril(torch.ones(self.n_hypotheses, self.n_hypotheses, device=x.device)) causal_attn = causal_attn.masked_fill(mask == 0, float('-inf')) causal_attn = self.dropout(torch.softmax(causal_attn, dim=-1)) - + refined_hypotheses = self.causal_norm(causal_in + torch.bmm(causal_attn, v)) # (N, K, D_out) - + # 7. Episodic Memory Write # Write the selected (highest value) path logic back to memory for future steps best_path_idx = path_probs.argmax(dim=-1) # (N) best_h_in = hypotheses_stack[torch.arange(N), best_path_idx] # (N, D) - + write_gates = torch.sigmoid(self.mem_write_gate(best_h_in)) # (N, M) write_data = self.mem_write_data(best_h_in) # (N, mem_dim) - + # Differentiable update to parameter matrix # This operates as a residual accumulation gradient update during training # For true recurrent state changes across tokens, an external hidden state is needed, # but the Parameter acts as a static "long term episodic registry" mem_update = torch.matmul(write_gates.t(), write_data) / (write_gates.sum(dim=0, keepdim=True).t() + 1e-6) - + # Add to predictions # Weigh refined hypotheses by their tree-search value fused_cognitive = (refined_hypotheses * path_probs.unsqueeze(-1)).sum(dim=1) # (N, D_out) - + # Add a residual from the memory write operation so write heads receive gradients fused_write = self.mem_write_proj(write_data * write_gates.mean(dim=-1, keepdim=True)) fused_cognitive = fused_cognitive + 0.1 * fused_write @@ -4052,33 +4053,33 @@ def __init__(self, dim: int, dropout: float = 0.1): nn.Dropout(dropout), nn.Linear(dim * 2, dim * 2) ) - + # Entanglement Routing approximations self.q_proj = nn.Linear(dim * 2, dim) self.k_proj = nn.Linear(dim * 2, dim) self.v_proj = nn.Linear(dim * 2, dim * 2) - + def forward(self, t: torch.Tensor, z: torch.Tensor): # z is (N, dim*2) representing complex state [Real, Imag] N, D2 = z.shape D = D2 // 2 - + # 1. Neural Mechanics Step t_tensor = t.expand(N, 1) z_t = torch.cat([z, t_tensor], dim=-1) dz_dt_base = self.net(z_t) # (N, D*2) - + # 2. Entanglement Routing # Instead of dot product similarity, we use a cheap pseudo-Kronecker entanglement q = self.q_proj(z) # (N, D) k = self.k_proj(z) # (N, D) v = self.v_proj(z) # (N, D*2) - + # Entangle Q and K via element-wise multiplication (diagonal of Kronecker) # then non-linear mix entangled = torch.relu(q * k).sum(dim=-1, keepdim=True) / math.sqrt(D) # (N, 1) routing_update = torch.sigmoid(entangled) * v # (N, D*2) - + return dz_dt_base + 0.1 * routing_update @@ -4086,7 +4087,7 @@ class OmniversalQuantumExpertHead(nn.Module): """ Generation v24: Omniversal Quantum Simulation Expert (Complex Superposition + Neural ODE + Holographic HRR + Entanglement) - + Abandons discrete real-valued layers for Continuous-Time Complex integration. Allows mutual contradiction handling via wave interference before Measurement Collapse. """ @@ -4102,7 +4103,7 @@ def __init__( self.out_dim = out_dim self.ode_steps = ode_steps # Discretization steps for Euler integrator self.complex_dim = in_dim # Operate in complex space of size in_dim - + # Base projection (backward-compatible residual) self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) @@ -4112,20 +4113,20 @@ def __init__( self.shared_down = nn.Linear(2048, out_dim, bias=False) self.shared_norm = nn.LayerNorm(out_dim) self.shared_scale = nn.Parameter(torch.tensor(1.0)) - + # ═══ 1. Quantum Lift (Measurement Prep) ═══ # Lifts Real-valued input to Complex-valued Wavefunction (Real, Imag) self.real_lift = nn.Linear(in_dim, self.complex_dim) self.imag_lift = nn.Linear(in_dim, self.complex_dim) - + # ═══ 2. Holographic Knowledge Binding (HRR Context) ═══ # Global holographic memory trace (learnable contextual superposition) self.holographic_memory = nn.Parameter(torch.randn(self.complex_dim * 2)) - + # ═══ 3. Continuous-Time Neural ODE ═══ # Defines the continuous derivative of the complex thought process self.ode_func = ODEFunc(dim=self.complex_dim, dropout=dropout) - + # ═══ 4. Measurement Collapse ═══ # Collapses the final Complex state back into distinct Real-valued logic (logits) self.measurement_proj = nn.Sequential( @@ -4133,7 +4134,7 @@ def __init__( nn.GELU(), nn.Linear(256, out_dim) ) - + # Measurement uncertainty modulator (Phase variance modulates confidence) self.uncertainty_gate = nn.Linear(self.complex_dim, 1) @@ -4150,24 +4151,24 @@ def reset_parameters(self): nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.normal_(self.shared_down.weight, std=0.01) - + nn.init.kaiming_uniform_(self.real_lift.weight, a=math.sqrt(5)) nn.init.zeros_(self.real_lift.bias) nn.init.kaiming_uniform_(self.imag_lift.weight, a=math.sqrt(5)) nn.init.zeros_(self.imag_lift.bias) - + for m in self.ode_func.modules(): if isinstance(m, nn.Linear): nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) if m.bias is not None: nn.init.zeros_(m.bias) - + for m in self.measurement_proj: if isinstance(m, nn.Linear): nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) if m.bias is not None: nn.init.zeros_(m.bias) - + nn.init.kaiming_uniform_(self.uncertainty_gate.weight, a=math.sqrt(5)) nn.init.zeros_(self.uncertainty_gate.bias) @@ -4198,15 +4199,15 @@ def forward(self, x): # Create Complex Wavefunction: Ψ = Real + i*Imag z_real = self.real_lift(x_flat) # (N, D) z_imag = self.imag_lift(x_flat) # (N, D) - + # Pack into (N, D*2) for compatibility with real-valued PyTorch ops z0 = torch.cat([z_real, z_imag], dim=-1) # (N, D*2) - + # ═══ PHASE 2: Holographic Knowledge Binding ═══ # Bind the current state with the global holographic memory trace memory_expanded = self.holographic_memory.unsqueeze(0).expand(N, -1) # (N, D*2) z_bound = self.circular_convolution(z0, memory_expanded) # (N, D*2) - + # Add bound knowledge into superposition z_t = z0 + 0.1 * z_bound @@ -4215,29 +4216,29 @@ def forward(self, x): # using a simple Euler integrator (differentiable) dt = 1.0 / self.ode_steps t = torch.zeros(1, device=x.device) - + for _ in range(self.ode_steps): # Compute derivative dz_dt = self.ode_func(t, z_t) # Step forward in continuous time z_t = z_t + dz_dt * dt t = t + dt - + # z_t is now the final complex state at t=1.0 # ═══ PHASE 4: Measurement Collapse ═══ # Extract Phase and Magnitude to measure uncertainty final_real, final_imag = z_t.chunk(2, dim=-1) # (N, D), (N, D) - + # Phase = atan2(Imag, Real). Represents conceptual angle/ambiguity. phase = torch.atan2(final_imag, final_real + 1e-8) # (N, D) - + # Uncertainty is based on phase complexity. High variance = high uncertainty. collapse_confidence = torch.sigmoid(self.uncertainty_gate(phase)) # (N, 1) - + # Collapse the Complex wavefunction back into Real-valued Action (Logits) collapsed_logits = self.measurement_proj(z_t) # (N, out_dim) - + # Apply confidence: if uncertainty is high, the model's absolute amplitude drops final_quantum_logic = collapsed_logits * collapse_confidence @@ -4285,7 +4286,7 @@ def __init__(self, dim: int, grid_size: int = 4, ca_steps: int = 3, dropout: flo self.dim = dim self.grid_size = grid_size self.ca_steps = ca_steps - + # State update network (operates on 3x3 local neighborhood) self.update_net = nn.Sequential( nn.Conv2d(dim, dim * 2, kernel_size=3, padding=1), @@ -4300,16 +4301,16 @@ def forward(self, x: torch.Tensor): N = x.shape[0] # Reshape flat vector into a 2D grid grid = x.view(N, self.dim, self.grid_size, self.grid_size) - + for _ in range(self.ca_steps): # NCA Update: compute delta dx = self.update_net(grid) - + # Stochastic update gate (allows cells to preserve state or change) update_prob = torch.sigmoid(self.stochastic_gate(grid)) # In training, we use soft updates. In true NCA this would be stochastic mask. grid = grid + update_prob * dx - + return grid.view(N, -1) @@ -4317,7 +4318,7 @@ class FractalGenesisExpertHead(nn.Module): """ Generation v25: Fractal Genesis Expert (Hyperbolic Space + Fractal Graph + Cellular Automata + Self-Assembly) - + Transforms reasoning from a flat vector sequence into a dynamically spawning hyperbolic computational fractal tree. """ @@ -4336,7 +4337,7 @@ def __init__( self.fractal_depth = fractal_depth self.atoms_count = atoms_count self.grid_size = grid_size - + # Base Euclidean projection (residual fallback) self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) @@ -4346,12 +4347,12 @@ def __init__( self.shared_down = nn.Linear(2048, out_dim, bias=False) self.shared_norm = nn.LayerNorm(out_dim) self.shared_scale = nn.Parameter(torch.tensor(1.0)) - + # ═══ 1. Hyperbolic Projection ═══ # Map Euclidean features to the Poincaré Disk manifold parameters self.hyperbolic_proj = nn.Linear(in_dim, in_dim) self.curvature = nn.Parameter(torch.tensor(1.0)) # Learned negative curvature c - + # ═══ 2. Semantic Self-Assembly ═══ # Pool of 'Logic Atoms' that dynamically construct the fractal node circuits self.logic_atoms = nn.Parameter(torch.randn(atoms_count, in_dim)) @@ -4360,29 +4361,29 @@ def __init__( nn.GELU(), nn.Linear(in_dim, in_dim) ) - + # ═══ 3. Fractal Graph Recursion & CA ═══ # Each level of the fractal tree has its own CA grid evolver # and a mechanism to spawn K=2 child nodes. self.fractal_evolvers = nn.ModuleList([ CellularAutomataGrid( - dim=in_dim // (grid_size**2), - grid_size=grid_size, + dim=in_dim // (grid_size**2), + grid_size=grid_size, dropout=dropout ) for _ in range(fractal_depth) ]) - + # Node spawn projections: splits a parent node into 2 children self.node_spawners = nn.ModuleList([ nn.Linear(in_dim, in_dim * 2) for _ in range(fractal_depth - 1) ]) - + # ═══ 4. Resonance Collapse Aggregation ═══ # Collapses 2 children back into a parent via phase resonance self.resonance_combiners = nn.ModuleList([ nn.Linear(in_dim * 2, in_dim) for _ in range(fractal_depth - 1) ]) - + self.final_projection = nn.Sequential( nn.Linear(in_dim, 256), nn.GELU(), @@ -4403,27 +4404,27 @@ def reset_parameters(self): nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) nn.init.normal_(self.shared_down.weight, std=0.01) - + nn.init.kaiming_uniform_(self.hyperbolic_proj.weight, a=math.sqrt(5)) nn.init.zeros_(self.hyperbolic_proj.bias) - + std = 1.0 / math.sqrt(self.in_dim) nn.init.uniform_(self.logic_atoms, -std, std) - + for m in self.assembly_net: if isinstance(m, nn.Linear): nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) if m.bias is not None: nn.init.zeros_(m.bias) - + for m in self.node_spawners: nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) nn.init.zeros_(m.bias) - + for m in self.resonance_combiners: nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) nn.init.zeros_(m.bias) - + for m in self.final_projection: if isinstance(m, nn.Linear): nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5)) @@ -4448,23 +4449,23 @@ def poincare_dist(self, x: torch.Tensor, y: torch.Tensor, c: torch.Tensor) -> to # Using arcosh trick dist = (2 / sqrt_c) * torch.log(arg + torch.sqrt(arg**2 - 1 + 1e-6)) return dist - + def assemble_node(self, node_state: torch.Tensor) -> torch.Tensor: """Dynamically binds logic atoms to the node state using hyperbolic distance.""" # Map node to hyperbolic space c = torch.clamp(F.softplus(self.curvature), min=1e-4) # Ensure c > 0 h_node = self.exp_map_zero(self.hyperbolic_proj(node_state), c) # (N, D) h_atoms = self.exp_map_zero(self.logic_atoms, c) # (A, D) - + # Compute pairwise hyperbolic distances: (N, 1, D) and (1, A, D) -> (N, A) dists = self.poincare_dist(h_node.unsqueeze(1), h_atoms.unsqueeze(0), c).squeeze(-1) - + # Attention over atoms (closer = higher weight) atom_weights = torch.softmax(-dists, dim=-1) # (N, A) - + # Assemble bounded circuit state bound_atoms = torch.matmul(atom_weights, self.logic_atoms) # (N, D) - + # Feed back to Euclidean space for processing assembled = self.assembly_net(torch.cat([node_state, bound_atoms], dim=-1)) # (N, D) return node_state + 0.1 * assembled @@ -4477,26 +4478,26 @@ def recursive_fractal_pass(self, node_state: torch.Tensor, depth: int) -> torch. """ # 1. Self-Assembly: bind logic atoms to current node state assembled_node = self.assemble_node(node_state) - + # 2. Base case (leaf nodes) if depth == 0: # Evolve via Cellular Automata and return return self.fractal_evolvers[0](assembled_node) - + # 3. Recursive step (internal nodes) # Spawn two children from current node children_raw = self.node_spawners[depth-1](assembled_node) # (N, D*2) child_left, child_right = children_raw.chunk(2, dim=-1) # (N, D), (N, D) - + # Recurse on children evolved_left = self.recursive_fractal_pass(child_left, depth - 1) evolved_right = self.recursive_fractal_pass(child_right, depth - 1) - + # 4. Resonance Collapse # Combine children back into parent. Instead of simple addition, we use # a resonance combiner to simulate constructive/destructive logic interference combined = self.resonance_combiners[depth-1](torch.cat([evolved_left, evolved_right], dim=-1)) - + # Finally evolve the combined parent node via its CA grid final_node = self.fractal_evolvers[depth](combined + assembled_node) return final_node @@ -4517,7 +4518,7 @@ def forward(self, x): # Launch the recursive fractal graph processing from the root root_state = x_flat final_euclidean_state = self.recursive_fractal_pass(root_state, depth=self.fractal_depth - 1) - + # Project back to logits fractal_logic = self.final_projection(final_euclidean_state) @@ -4577,11 +4578,11 @@ def backward(ctx, grad_output): # Backward pass: Fast sigmoid surrogate gradient sgax = (input - threshold) * alpha S_dtps = torch.exp(-torch.abs(sgax)) / (1 + torch.exp(-torch.abs(sgax)))**2 - + grad_input = grad_output * alpha * S_dtps grad_threshold = (grad_output * (-alpha) * S_dtps).sum() grad_alpha = (grad_output * (input - threshold) * S_dtps).sum() - + return grad_input, grad_threshold, grad_alpha def surrogate_spike(input, threshold, alpha=torch.tensor(2.0)): @@ -4598,52 +4599,52 @@ def __init__(self, in_features: int, decay: float = 0.9, threshold: float = 1.0) self.decay = decay self.threshold = nn.Parameter(torch.tensor(threshold)) self.alpha = nn.Parameter(torch.tensor(2.0)) # Surrogate gradient sharpness - + def forward(self, x: torch.Tensor, v_mem: torch.Tensor): # Update membrane potential: decay previous + integrate new input v_mem = v_mem * self.decay + x - + # Fire spike if v_mem > threshold spike = surrogate_spike(v_mem, self.threshold, self.alpha) - + # Soft reset: subtract threshold from membrane if fired v_mem = v_mem - spike * self.threshold - + return spike, v_mem class LiquidSynapseMPS(nn.Module): """ - Combines Tensor Networks (1D Matrix Product State approximation) + Combines Tensor Networks (1D Matrix Product State approximation) with Liquid Neural Network synaptic dynamics (State-dependent differential scaling). """ def __init__(self, dim: int, rank: int = 16): super().__init__() self.dim = dim self.rank = rank - + # Matrix Product State (MPS) core tensors (highly compressed representation of dense weights) # Instead of dim x dim, we use (dim x rank) and (rank x dim) to emulate a tensor train self.core_1 = nn.Parameter(torch.randn(dim, rank) / math.sqrt(dim)) self.core_2 = nn.Parameter(torch.randn(rank, dim) / math.sqrt(rank)) - + # Liquid Synaptic Time Constants (LNN mechanics) self.tau = nn.Parameter(torch.ones(dim)) # Non-linear input processor for the liquid equation self.liquid_gate = nn.Linear(dim, dim) - + def forward(self, x: torch.Tensor, liquid_state: torch.Tensor, dt: float = 0.1): # 1. Tensor Network (MPS) transformation # x is (..., dim) hidden = torch.matmul(x, self.core_1) # (..., rank) mps_transform = torch.matmul(hidden, self.core_2) # (..., dim) - + # 2. Liquid Synapse update: dx/dt = -x/tau + S(x, I) sensory_input = torch.sigmoid(self.liquid_gate(mps_transform)) - + # Differential equation Euler step d_state = -liquid_state / torch.clamp(F.softplus(self.tau), min=1e-3) + sensory_input new_state = liquid_state + d_state * dt - + return new_state @@ -4651,7 +4652,7 @@ class LiquidSpikingTensorExpertHead(nn.Module): """ Generation v26: Liquid Spiking Tensor Expert (Spiking Neural Networks + Liquid Synapses + Test-Time Training + Tensor Networks) - + A bio-plausible, highly compressed architectured adapting dynamically during inference. """ def __init__( @@ -4666,7 +4667,7 @@ def __init__( self.in_dim = in_dim self.out_dim = out_dim self.spiking_steps = spiking_steps - + # Base Euclidean projection self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) @@ -4676,20 +4677,20 @@ def __init__( self.shared_down = nn.Linear(2048, out_dim, bias=False) self.shared_norm = nn.LayerNorm(out_dim) self.shared_scale = nn.Parameter(torch.tensor(1.0)) - + # ═══ 1. Test-Time Training (TTT) Fast-Weight Module ═══ # Internal LoRA adapter that updates *during the forward pass* self.ttt_lora_A = nn.Parameter(torch.randn(in_dim, 8) / math.sqrt(in_dim)) self.ttt_lora_B = nn.Parameter(torch.zeros(8, out_dim)) # Self-supervised auxiliary task (Reconstruct input from bottleneck) self.ttt_decoder = nn.Linear(8, in_dim) - + # ═══ 2. Liquid Synapse / Tensor Network Core ═══ self.liquid_synapse = LiquidSynapseMPS(dim=in_dim, rank=mps_rank) - + # ═══ 3. Spiking Neural Mechanics (LIF Neurons) ═══ self.lif_neuron = LIFNeuron(in_features=in_dim) - + # ═══ 4. Neuro-Symbolic Thresholding ═══ self.spike_decoder = nn.Linear(in_dim, out_dim) @@ -4709,10 +4710,10 @@ def reset_parameters(self): nn.init.normal_(self.ttt_lora_B, std=0.01) self.lif_neuron.threshold.data.fill_(0.1) - + nn.init.kaiming_uniform_(self.ttt_decoder.weight, a=math.sqrt(5)) nn.init.zeros_(self.ttt_decoder.bias) - + nn.init.kaiming_uniform_(self.spike_decoder.weight, a=math.sqrt(5)) nn.init.zeros_(self.spike_decoder.bias) @@ -4723,25 +4724,25 @@ def test_time_training_step(self, x: torch.Tensor) -> torch.Tensor: """ # Detach x from the global computation graph so inner loop grads don't leak x_inner = x.detach().requires_grad_(True) - + # Self-supervised forward: Encode using shared representation, decode using aux task latent = torch.matmul(x_inner, self.ttt_lora_A) reconstruction = self.ttt_decoder(latent) - + # Loss: Mean Squared Error representation loss = F.mse_loss(reconstruction, x_inner) - + # Calculate gradients for the Lora A encoder and aux decoder grads = torch.autograd.grad( - loss, - [self.ttt_lora_A, self.ttt_decoder.weight, self.ttt_decoder.bias], - create_graph=self.training, + loss, + [self.ttt_lora_A, self.ttt_decoder.weight, self.ttt_decoder.bias], + create_graph=self.training, retain_graph=True ) - + # TTT SGD Step (lr = 0.01). Update only the shared encoder fast-weights. fast_lora_A = self.ttt_lora_A - 0.01 * grads[0] - + # Compute TTT contribution for this batch using adapted encoder and frozen task-decoder ttt_output = torch.matmul(torch.matmul(x, fast_lora_A), self.ttt_lora_B) # (N, out_dim) return ttt_output @@ -4766,17 +4767,17 @@ def forward(self, x): v_mem = torch.zeros_like(x_flat) liquid_state = torch.zeros_like(x_flat) spike_train_accum = torch.zeros_like(x_flat) - + for _ in range(self.spiking_steps): # Liquid Synapse (MPS) update liquid_state = self.liquid_synapse(x_flat, liquid_state) - + # Feed continuously changing liquid state into LIF Neuron spike, v_mem = self.lif_neuron(liquid_state, v_mem) - + # Accumulate discrete spikes spike_train_accum = spike_train_accum + spike - + # ═══ 4. Neuro-Symbolic Thresholding ═══ # Average firing rate decoded back into continuous real logits firing_rate = spike_train_accum / self.spiking_steps @@ -4823,10 +4824,10 @@ def forward(self, x): class ActiveInferenceExpertHead(nn.Module): """ v27 Active Inference Expert Head based on the Free Energy Principle. - It introduces a generative World Model and minimizes Variational Free Energy + It introduces a generative World Model and minimizes Variational Free Energy at test time before routing to experts. """ - def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 6, + def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 6, inference_steps: int = 3, lr: float = 0.1, dropout: float = 0.1): super().__init__() self.in_dim = in_dim @@ -4834,17 +4835,17 @@ def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 6, self.n_experts = n_experts self.inference_steps = inference_steps self.lr = lr - + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) - + self.world_model = nn.Linear(in_dim, in_dim, bias=True) self.world_prior = nn.Linear(in_dim, in_dim, bias=False) self.precision_gate = nn.Linear(in_dim, n_experts * 2, bias=True) - + self.experts_up = nn.ModuleList([nn.Linear(in_dim, 1024, bias=False) for _ in range(n_experts)]) self.experts_down = nn.ModuleList([nn.Linear(1024, out_dim, bias=False) for _ in range(n_experts)]) - + self.alpha = nn.Parameter(torch.tensor(0.0)) self.dropout = nn.Dropout(dropout) self.reset_parameters() @@ -4866,7 +4867,7 @@ def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) original_shape = x.shape x_flat = x.view(-1, self.in_dim) - + # Test-time inner loop optimization with torch.enable_grad(): z = x_flat.clone().detach().requires_grad_(True) @@ -4875,27 +4876,27 @@ def forward(self, x): surprise = F.mse_loss(predicted_state, x_flat, reduction='sum') complexity = F.mse_loss(z, self.world_prior(x_flat), reduction='sum') free_energy = surprise + 0.1 * complexity - + grad_z = torch.autograd.grad(free_energy, z, create_graph=self.training)[0] z = z - self.lr * grad_z - + z_final = z gating_out = self.precision_gate(z_final) mu, logvar = gating_out.chunk(2, dim=-1) precision = torch.exp(-logvar) - + gate_scores = torch.softmax(mu * precision, dim=-1) - + expert_outputs = [] for i in range(self.n_experts): h = F.silu(self.experts_up[i](z_final)) h = self.dropout(h) out = self.experts_down[i](h) expert_outputs.append(out) - + expert_stack = torch.stack(expert_outputs, dim=-2) routed_out = (expert_stack * gate_scores.unsqueeze(-1)).sum(dim=-2) - + output = base_logits.view(-1, self.out_dim) + self.alpha * routed_out return output.view(*original_shape[:-1], self.out_dim) @@ -5224,13 +5225,13 @@ def __init__(self, d_model: int, d_state: int = 64, dt_min: float = 0.001, dt_ma super().__init__() self.d_model = d_model self.d_state = d_state - + # State space parameters self.A_log = nn.Parameter(torch.empty(d_model, d_state)) self.B = nn.Parameter(torch.empty(d_model, d_state)) self.C = nn.Parameter(torch.empty(d_model, d_state)) self.D = nn.Parameter(torch.zeros(d_model)) - + self.dt_proj = nn.Linear(d_model, d_model) self.dt_min = dt_min self.dt_max = dt_max @@ -5250,25 +5251,25 @@ def forward(self, x): # 1. Delta (time scale) dt = F.softplus(self.dt_proj(x)) dt = torch.clamp(dt, min=self.dt_min, max=self.dt_max) - + # Continuous A must be negative to be stable A = -torch.exp(self.A_log) - + # 2. ZOH Discretization for A and B # dA = exp(dt * A) # dB = (exp(dt * A) - I) / A * B - + dA = torch.exp(dt.unsqueeze(-1) * A) dB = (dA - 1.0) / (A + 1e-6) * self.B - - # Initialize sequence hidden state to zeros (assuming per-token routing execution, + + # Initialize sequence hidden state to zeros (assuming per-token routing execution, # this acts like a depth-wise recurrent pass per token) h = torch.zeros(x.shape[:-1] + (self.d_model, self.d_state), device=x.device, dtype=x.dtype) - + # Step the state 3 times recursively for depth without sequence for _ in range(3): h = dA * h + dB * x.unsqueeze(-1) - + y = (self.C * h).sum(dim=-1) + self.D * x return F.silu(y) @@ -5285,24 +5286,24 @@ def __init__(self, in_dim: int = 256, out_dim: int = 10, n_experts: int = 4, dro self.in_dim = in_dim self.out_dim = out_dim self.n_experts = n_experts - + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) self.bias = nn.Parameter(torch.zeros(out_dim)) - + # Learned holographic associative memory matrix self.holographic_memory = nn.Parameter(torch.empty(1, in_dim)) self.holo_proj = nn.Linear(in_dim, in_dim, bias=False) - + self.ssm_experts = nn.ModuleList([StateSpaceExpert(in_dim) for _ in range(n_experts)]) - + # Adjacency topology router # Outputs a (n_experts, n_experts) adjacency matrix per token self.topology_router = nn.Linear(in_dim, n_experts * n_experts) - + self.out_proj = nn.Linear(in_dim, out_dim, bias=False) self.alpha = nn.Parameter(torch.tensor(0.0)) self.dropout = nn.Dropout(dropout) - + self.reset_parameters() def reset_parameters(self): @@ -5324,12 +5325,12 @@ def _circular_convolution(self, a, b): def forward(self, x): base_logits = F.linear(x, self.weight, self.bias) - + # 1. Holographic Memory Binding # Bind the token state x with the learned persistent holographic_memory x_bound = self._circular_convolution(x, self.holographic_memory.expand_as(x)) x_holo = self.holo_proj(x_bound) - + # 2. Dynamic Topology Routing # B x T x (N*N) -> B x T x N x N adj_flat = self.topology_router(x_holo) @@ -5337,18 +5338,18 @@ def forward(self, x): # Convert to a soft permutation matrix via doubly stochastic approximation (Sinkhorn) # For simplicity in testing, we use softmax over rows adj = torch.softmax(adj_logits, dim=-1) - + # 3. State Space Execution Path # Instead of parallel execution, we execute iteratively. current_state = x_holo accum_out = torch.zeros_like(x_holo) - + expert_outs = [] for i in range(self.n_experts): h = self.ssm_experts[i](current_state) expert_outs.append(h) expert_stack = torch.stack(expert_outs, dim=-2) # (..., N_experts, D) - + # Route outputs through adjacency. # current_state = \sum_j adj[i, j] * expert_stack[j] # In a generic formulation, we compute the aggregated state as a mixture. @@ -5356,19 +5357,19 @@ def forward(self, x): for step in range(self.n_experts): step_mix = torch.einsum('...ij,...jd->...id', adj, expert_stack) # The next state input is the average mixture for the next nodes - current_state = step_mix.mean(dim=-2) + current_state = step_mix.mean(dim=-2) accum_out = accum_out + current_state - + # Re-evaluate experts via recurrence (dynamic routing unrolled) expert_outs = [] for i in range(self.n_experts): h = self.ssm_experts[i](current_state) expert_outs.append(h) expert_stack = torch.stack(expert_outs, dim=-2) - + final_h = accum_out / self.n_experts routed_logits = self.out_proj(self.dropout(final_h)) - + return base_logits + self.alpha * routed_logits @@ -5916,13 +5917,523 @@ def deep_supervision_loss(self, targets): return self.layers[10].deep_supervision_loss(targets) -class ChampionNetHierarchicalExpert(nn.Module): +class LightweightMHA(nn.Module): """ - Backbone-compatible model with HierarchicalMoEClassifierHead. - Uses the same layers 0-9 backbone from ChampionNet with the new - hierarchical MoE head at layer 10. + Lightweight Multi-Head Attention for sequence of length 2 (z_l and z_h). """ - def __init__( + def __init__(self, dim, num_heads=2, dropout=0.1): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=False) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(dropout) + + def forward(self, x): + B, T, C = x.shape + qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # (B, num_heads, T, head_dim) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + + x_out = (attn @ v).transpose(1, 2).reshape(B, T, C) + x_out = self.proj(x_out) + x_out = self.proj_drop(x_out) + return x_out + + +class CognitiveLeapUltraExpertHead(nn.Module): + """ + Generation v51: Cognitive Leap Ultra Expert + (Tiny Recursive Latent Reasoning with Cross-Attention + Recurrent MoE Core + Entropy Halting) + """ + def __init__( + self, + in_dim: int = 256, + out_dim: int = 10, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + n_cores: int = 4, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.latent_dim = latent_dim + self.core_dim = core_dim + self.n_cycles = n_cycles + self.inner_steps = inner_steps + self.max_cycles = max_cycles + self.n_cores = n_cores + + # Base projection (backward-compatible) + self.weight = nn.Parameter(torch.empty(out_dim, in_dim)) + self.bias = nn.Parameter(torch.zeros(out_dim)) + + # Shared Expert + self.shared_up = nn.Linear(in_dim, 2048, bias=False) + self.shared_down = nn.Linear(2048, out_dim, bias=False) + self.shared_norm = nn.LayerNorm(out_dim) + self.shared_scale = nn.Parameter(torch.tensor(1.0)) + + # Latent initializers + self.thought_init = nn.Linear(in_dim, latent_dim, bias=False) + self.answer_init = nn.Linear(in_dim, latent_dim, bias=False) + + # Recurrent MoE Core + self.cores = nn.ModuleList([ + nn.Sequential( + nn.Linear(in_dim + 2 * latent_dim, core_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(core_dim, latent_dim), + ) for _ in range(n_cores) + ]) + self.core_router = nn.Linear(in_dim + 2 * latent_dim, n_cores) + + # Cross-Attention Refinement + self.cross_attn = LightweightMHA(latent_dim, num_heads=2, dropout=dropout) + + # Outer answer update + self.answer_update = nn.Sequential( + nn.Linear(2 * latent_dim, latent_dim), + nn.GELU(), + nn.Linear(latent_dim, latent_dim), + ) + + # Hypersphere gains + self.latent_gain = nn.Parameter(torch.ones(latent_dim)) + + # Cycle conditioning + self.cycle_embed = nn.Embedding(max_cycles, latent_dim) + + # ACT-style halting head + self.halt_head = nn.Linear(latent_dim, 1) + + # Per-cycle decoder + self.decode_head = nn.Linear(latent_dim, out_dim, bias=False) + + self.alpha = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(dropout) + + # Diagnostics / regularizers + self.register_buffer("last_ponder_cost", torch.tensor(0.0)) + self.register_buffer("last_consistency_loss", torch.tensor(0.0)) + self.register_buffer("last_cycles_used", torch.tensor(0.0)) + self.register_buffer("last_gating_entropy", torch.tensor(0.0)) + self.register_buffer("last_prediction_streak", torch.tensor(0.0), persistent=False) + self.register_buffer("last_prediction_confidence_delta", torch.tensor(0.0), persistent=False) + self.register_buffer( + "last_prediction_topk_js_divergence", torch.tensor(0.0), persistent=False + ) + self.register_buffer( + "last_prediction_topk_js_divergence_max", torch.tensor(0.0), persistent=False + ) + self.last_exit_reason = "not_run" + + # Per-cycle decode cache for deep improvement supervision + self._cycle_logits = [] + + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + nn.init.kaiming_uniform_(self.shared_up.weight, a=math.sqrt(5)) + nn.init.normal_(self.shared_down.weight, std=0.01) + + nn.init.kaiming_uniform_(self.thought_init.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.answer_init.weight, a=math.sqrt(5)) + + for core in self.cores: + for layer in core: + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + nn.init.normal_(core[-1].weight, std=0.01) + + nn.init.kaiming_uniform_(self.core_router.weight, a=math.sqrt(5)) + if self.core_router.bias is not None: + nn.init.zeros_(self.core_router.bias) + + nn.init.kaiming_uniform_(self.cross_attn.qkv.weight, a=math.sqrt(5)) + nn.init.normal_(self.cross_attn.proj.weight, std=0.01) + + for layer in self.answer_update: + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight, a=math.sqrt(5)) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + nn.init.normal_(self.answer_update[-1].weight, std=0.01) + + nn.init.normal_(self.cycle_embed.weight, std=0.02) + + nn.init.zeros_(self.halt_head.weight) + nn.init.constant_(self.halt_head.bias, -1.0) + + nn.init.normal_(self.decode_head.weight, std=0.01) + + def _sphere_norm(self, z): + return F.normalize(z, p=2, dim=-1) * self.latent_gain + + def deep_supervision_loss(self, targets): + if not self._cycle_logits: + raise RuntimeError( + "deep_supervision_loss requires a training-mode forward pass first" + ) + targets_flat = targets.reshape(-1) + n = len(self._cycle_logits) + weights = torch.arange( + 1, n + 1, device=targets_flat.device, dtype=self._cycle_logits[0].dtype + ) + weights = weights / weights.sum() + loss = torch.zeros((), device=targets_flat.device, dtype=self._cycle_logits[0].dtype) + for w, logits in zip(weights, self._cycle_logits): + loss = loss + w * F.cross_entropy(logits, targets_flat) + return loss + + @staticmethod + def _topk_js_divergence( + previous_probs: torch.Tensor, + current_probs: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + """Return per-row top-k Jensen-Shannon drift with an other-mass bucket.""" + + k = min(max(1, int(top_k)), int(previous_probs.shape[-1])) + midpoint = 0.5 * (previous_probs + current_probs) + indices = midpoint.topk(k, dim=-1).indices + previous_top = previous_probs.gather(-1, indices) + current_top = current_probs.gather(-1, indices) + previous_other = (1.0 - previous_top.sum(dim=-1, keepdim=True)).clamp_min(0.0) + current_other = (1.0 - current_top.sum(dim=-1, keepdim=True)).clamp_min(0.0) + previous_compact = torch.cat([previous_top, previous_other], dim=-1) + current_compact = torch.cat([current_top, current_other], dim=-1) + compact_midpoint = 0.5 * (previous_compact + current_compact) + epsilon = torch.finfo(previous_compact.dtype).tiny + + def _kl(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + terms = left * ( + torch.log(left.clamp_min(epsilon)) + - torch.log(right.clamp_min(epsilon)) + ) + return terms.sum(dim=-1) + + return 0.5 * ( + _kl(previous_compact, compact_midpoint) + + _kl(current_compact, compact_midpoint) + ) + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + exit_entropy_threshold: float = 0.2, + prediction_stability_patience: int = 2, + prediction_stability_tol: float = 5e-3, + prediction_stability_top_k: int = 5, + ): + N = x.shape[0] * x.shape[1] if x.dim() == 3 else x.shape[0] + shape_prefix = x.shape[:-1] + x_flat = x.reshape(N, self.in_dim) + + base_logits = F.linear(x_flat, self.weight, self.bias) + + # Global Shared Expert + shared_out = self.shared_norm( + self.shared_down(self.dropout(F.silu(self.shared_up(x_flat)))) + ) + + n_cycles = int(reasoning_cycles) if reasoning_cycles is not None else self.n_cycles + n_cycles = max(1, n_cycles) + + # Initialize latents on the hypersphere + z_l = self._sphere_norm(self.thought_init(x_flat)) + z_h = self._sphere_norm(self.answer_init(x_flat)) + + cycle_logits = [] + halt_probs = [] + gate_probs_list = [] + remaining = torch.ones(N, 1, device=x_flat.device, dtype=x_flat.dtype) + ponder = torch.zeros(N, 1, device=x_flat.device, dtype=x_flat.dtype) + consistency = torch.zeros((), device=x_flat.device, dtype=x_flat.dtype) + prev_z_h = None + cycles_used = n_cycles + stability_patience = max(0, int(prediction_stability_patience)) + stability_tol = max(0.0, float(prediction_stability_tol)) + use_prediction_stability = adaptive_compute and not self.training and stability_patience > 0 + weighted_recursive = torch.zeros(N, self.out_dim, device=x_flat.device, dtype=x_flat.dtype) + previous_prediction = None + prediction_streak = torch.zeros(N, device=x_flat.device, dtype=torch.long) + stable_confidence_min = None + stable_confidence_max = None + prediction_confidence_delta = torch.zeros(N, device=x_flat.device, dtype=x_flat.dtype) + prediction_topk_js_divergence = torch.zeros( + N, device=x_flat.device, dtype=x_flat.dtype + ) + prediction_topk_js_divergence_max = torch.zeros_like( + prediction_topk_js_divergence + ) + previous_prediction_distribution = None + exit_reason = "max_cycles" + + for t in range(n_cycles): + cycle_idx = min(t, self.max_cycles - 1) + cycle_vec = self.cycle_embed.weight[cycle_idx].unsqueeze(0) + z_l_cond = z_l + cycle_vec + + # Inner scratchpad refinement (weight-tied recurrent MoE core) + for _ in range(self.inner_steps): + core_in = torch.cat([x_flat, z_l_cond, z_h], dim=-1) + gate_logits = self.core_router(core_in) + gate_probs = F.softmax(gate_logits, dim=-1) + gate_probs_list.append(gate_probs) + + # Route to recurrent cores + core_outs = torch.stack([core(core_in) for core in self.cores], dim=1) # (N, n_cores, latent_dim) + core_update = (core_outs * gate_probs.unsqueeze(-1)).sum(dim=1) # (N, latent_dim) + + z_l_cond = self._sphere_norm(z_l_cond + core_update) + z_l = z_l_cond + + # Attention-Based Cross-Latent Gating + z_seq = torch.stack([z_l, z_h], dim=1) # (N, 2, latent_dim) + z_seq = z_seq + self.cross_attn(z_seq) + z_l_refined = z_seq[:, 0, :] + z_h_refined = z_seq[:, 1, :] + + # Outer answer refinement + z_h = self._sphere_norm( + z_h_refined + self.answer_update(torch.cat([z_l_refined, z_h_refined], dim=-1)) + ) + + logits = self.decode_head(z_h) + cycle_logits.append(logits) + + # Prediction entropy + probs = F.softmax(logits, dim=-1) + entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1, keepdim=True) + + # Latent-convergence measurement + if prev_z_h is not None: + delta = (z_h - prev_z_h).pow(2).mean() + consistency = consistency + delta + else: + delta = None + prev_z_h = z_h + + # Treat the full prefix output as an internal verifier. If its top + # prediction and confidence stabilize, more recurrent work is not + # buying a meaningfully different answer for this request. + prediction_stable = False + if use_prediction_stability: + prefix_recursive = weighted_recursive + remaining * logits + prefix_output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * prefix_recursive + ) + prefix_probs = F.softmax(prefix_output, dim=-1) + if previous_prediction_distribution is not None: + prediction_topk_js_divergence = self._topk_js_divergence( + previous_prediction_distribution, + prefix_probs, + prediction_stability_top_k, + ) + prediction_topk_js_divergence_max = torch.maximum( + prediction_topk_js_divergence_max, + prediction_topk_js_divergence, + ) + previous_prediction_distribution = prefix_probs + prediction_confidence, prediction = prefix_probs.max(dim=-1) + if previous_prediction is None: + prediction_streak = torch.ones_like(prediction_streak) + stable_confidence_min = prediction_confidence + stable_confidence_max = prediction_confidence + else: + same_prediction = prediction.eq(previous_prediction) + prediction_streak = torch.where( + same_prediction, + prediction_streak + 1, + torch.ones_like(prediction_streak), + ) + stable_confidence_min = torch.where( + same_prediction, + torch.minimum(stable_confidence_min, prediction_confidence), + prediction_confidence, + ) + stable_confidence_max = torch.where( + same_prediction, + torch.maximum(stable_confidence_max, prediction_confidence), + prediction_confidence, + ) + prediction_confidence_delta = stable_confidence_max - stable_confidence_min + previous_prediction = prediction + stable_mask = ( + prediction_streak >= stability_patience + ) & ( + prediction_confidence_delta <= stability_tol + ) + prediction_stable = bool(stable_mask.all().item()) + + # Inference-only verifier and convergence exits. Prediction + # stability is checked first so diagnostics identify the new path + # when multiple exit criteria become true on the same cycle. + converged = False + if adaptive_compute and not self.training: + if prediction_stable: + converged = True + exit_reason = "prediction_stable" + elif delta is not None and delta.item() < exit_tol: + converged = True + exit_reason = "latent_converged" + elif entropy.mean().item() < exit_entropy_threshold: + converged = True + exit_reason = "low_entropy" + elif remaining.mean().item() < 0.05: + converged = True + exit_reason = "halt_mass" + + # Differentiable ACT halting + h_t = torch.sigmoid(self.halt_head(z_h)) # (N, 1) + if t == n_cycles - 1 or converged: + p_t = remaining + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + cycles_used = t + 1 + break + p_t = remaining * h_t + remaining = remaining * (1.0 - h_t) + halt_probs.append(p_t) + ponder = ponder + p_t * float(t + 1) + if use_prediction_stability: + weighted_recursive = weighted_recursive + p_t * logits + + self.last_cycles_used = torch.tensor( + float(cycles_used), device=x_flat.device + ) + self.last_prediction_streak = prediction_streak.float().mean() + self.last_prediction_confidence_delta = prediction_confidence_delta.mean() + self.last_prediction_topk_js_divergence = ( + prediction_topk_js_divergence.mean().detach() + ) + self.last_prediction_topk_js_divergence_max = ( + prediction_topk_js_divergence_max.mean().detach() + ) + self.last_exit_reason = exit_reason + + # Calculate gating entropy for MoE cores + if gate_probs_list: + all_gate_probs = torch.stack(gate_probs_list, dim=0) # (Steps, N, n_cores) + mean_gate_probs = all_gate_probs.mean(dim=(0, 1)) # (n_cores,) + gating_entropy = -torch.sum(mean_gate_probs * torch.log(mean_gate_probs + 1e-9)) + else: + gating_entropy = torch.tensor(0.0, device=x_flat.device) + + if self.training: + self.last_ponder_cost = ponder.mean() + self.last_consistency_loss = consistency / max(1, cycles_used - 1) + self.last_gating_entropy = gating_entropy + self._cycle_logits = cycle_logits + else: + self.last_gating_entropy = gating_entropy + self._cycle_logits = [] + + # Halting-weighted mixture of per-cycle decodes + logits_stack = torch.stack(cycle_logits, dim=1) # (N, T, out_dim) + probs_stack = torch.stack(halt_probs, dim=1) # (N, T, 1) + recursive_out = (logits_stack * probs_stack).sum(dim=1) # (N, out_dim) + + output = ( + base_logits + + self.shared_scale * shared_out + + (self.alpha + 1e-4) * recursive_out + ) + + return output.view(*shape_prefix, -1) + + +class ChampionNetCognitiveLeapUltraExpert(nn.Module): + """ + Backbone-compatible model with CognitiveLeapUltraExpertHead (v51). + """ + def __init__( + self, + latent_dim: int = 128, + core_dim: int = 256, + n_cycles: int = 3, + inner_steps: int = 2, + max_cycles: int = 16, + dropout: float = 0.1, + n_cores: int = 4, + ): + super().__init__() + base = ChampionNet() + layers = [base.layers[i] for i in range(10)] + layers.append(CognitiveLeapUltraExpertHead( + 256, 10, + latent_dim=latent_dim, + core_dim=core_dim, + n_cycles=n_cycles, + inner_steps=inner_steps, + max_cycles=max_cycles, + dropout=dropout, + n_cores=n_cores, + )) + layers.append(base.layers[11]) + self.layers = nn.ModuleList(layers) + + def forward( + self, + x, + reasoning_cycles: Optional[int] = None, + adaptive_compute: bool = False, + exit_tol: float = 1e-3, + exit_entropy_threshold: float = 0.2, + prediction_stability_patience: int = 2, + prediction_stability_tol: float = 5e-3, + prediction_stability_top_k: int = 5, + ): + for layer in self.layers: + if isinstance(layer, CognitiveLeapUltraExpertHead): + x = layer( + x, + reasoning_cycles=reasoning_cycles, + adaptive_compute=adaptive_compute, + exit_tol=exit_tol, + exit_entropy_threshold=exit_entropy_threshold, + prediction_stability_patience=prediction_stability_patience, + prediction_stability_tol=prediction_stability_tol, + prediction_stability_top_k=prediction_stability_top_k, + ) + else: + x = layer(x) + return x + + def deep_supervision_loss(self, targets): + return self.layers[10].deep_supervision_loss(targets) + + +class ChampionNetHierarchicalExpert(nn.Module): + """ + Backbone-compatible model with HierarchicalMoEClassifierHead. + Uses the same layers 0-9 backbone from ChampionNet with the new + hierarchical MoE head at layer 10. + """ + def __init__( self, n_domain_groups: int = 2, experts_per_group: int = 4, @@ -6013,10 +6524,10 @@ def __init__(self, n_experts: int = 8, reasoning_steps: int = 3, lora_rank: int base = ChampionNet() layers = [base.layers[i] for i in range(10)] layers.append(ThoughtExpertClassifierHead( - 256, 10, - n_experts=n_experts, - reasoning_steps=reasoning_steps, - lora_rank=lora_rank, + 256, 10, + n_experts=n_experts, + reasoning_steps=reasoning_steps, + lora_rank=lora_rank, dropout=dropout )) layers.append(base.layers[11]) @@ -6401,6 +6912,8 @@ def build_model( return ChampionNetCognitiveExpert(dropout=dropout) if model_size == "cognitive_leap_expert": return ChampionNetCognitiveLeapExpert(dropout=dropout) + if model_size == "cognitive_leap_ultra_expert": + return ChampionNetCognitiveLeapUltraExpert(dropout=dropout) if model_size == "transcendent_expert": return ChampionNetTranscendentExpert(dropout=dropout) if model_size == "omniversal_expert": @@ -6711,7 +7224,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) "layers.10.down_proj.weight", "layers.10.layer_scale.gamma", }) - + # Robust loading: if checkpoint is not megalarge, strip head keys to avoid mismatch ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "megalarge": @@ -6719,7 +7232,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in mega_allowed] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -6736,14 +7249,14 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) "layers.10.alpha", "layers.10.calibration.weight", "layers.10.calibration.bias", "layers.10.theta", } - + ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "ultra_expert": filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -6768,14 +7281,14 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) allowed_missing.add(f"{head_pref}expert_norms.{i}.bias") allowed_missing.add(f"{head_pref}lora_down.{i}.weight") allowed_missing.add(f"{head_pref}lora_up.{i}.weight") - + ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "hierarchical_expert": filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -6784,7 +7297,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) head_pref = "layers.10." allowed_missing = { f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", - f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", f"{head_pref}shared_scale", f"{head_pref}expert_bias", f"{head_pref}alpha", f"{head_pref}calibration.weight", f"{head_pref}calibration.bias", f"{head_pref}theta", @@ -6814,7 +7327,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -6823,10 +7336,10 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) head_pref = "layers.10." allowed_missing = { f"{head_pref}shared_up.weight", f"{head_pref}shared_down.weight", - f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", + f"{head_pref}shared_norm.weight", f"{head_pref}shared_norm.bias", f"{head_pref}shared_scale", f"{head_pref}expert_bias", f"{head_pref}local_up.weight", f"{head_pref}local_down.weight", - f"{head_pref}local_norm.weight", f"{head_pref}local_norm.bias", + f"{head_pref}local_norm.weight", f"{head_pref}local_norm.bias", f"{head_pref}local_scale", f"{head_pref}alpha", f"{head_pref}calibration.weight", f"{head_pref}calibration.bias", f"{head_pref}theta", @@ -6859,7 +7372,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -6888,7 +7401,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -7226,6 +7739,29 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered + if model_size == "cognitive_leap_ultra_expert": + head_pref = "layers.10." + allowed_missing = { + key + for key in model.state_dict().keys() + if key.startswith(head_pref) and key not in {f"{head_pref}weight", f"{head_pref}bias"} + } + + ckpt_size = detect_model_size_from_state_dict(state_dict) + if ckpt_size != "cognitive_leap_ultra_expert": + filtered_sd = { + k: v + for k, v in state_dict.items() + if not (k.startswith(head_pref) and k not in [f"{head_pref}weight", f"{head_pref}bias"]) + } + incompatible = model.load_state_dict(filtered_sd, strict=False) + else: + incompatible = model.load_state_dict(state_dict, strict=False) + + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] + unexpected_filtered = [k for k in incompatible.unexpected_keys if k] + return missing_filtered, unexpected_filtered + if model_size == "neurogenesis_expert": head_pref = "layers.10." @@ -7293,7 +7829,7 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) if model_size == "deep_expert": allowed_missing = { "layers.10.shared_up.weight", "layers.10.shared_down.weight", - "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", + "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", "layers.10.shared_scale", "layers.10.expert_bias", "layers.10.gate.weight", "layers.10.noise_gate.weight", "layers.10.alpha", "layers.10.calibration.weight", @@ -7304,14 +7840,14 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) allowed_missing.add(f"layers.10.experts_down.{i}.weight") allowed_missing.add(f"layers.10.expert_norms.{i}.weight") allowed_missing.add(f"layers.10.expert_norms.{i}.bias") - + ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "deep_expert": filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -7327,14 +7863,14 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) allowed_missing.add(f"layers.10.experts_down.{i}.weight") allowed_missing.add(f"layers.10.expert_norms.{i}.weight") allowed_missing.add(f"layers.10.expert_norms.{i}.bias") - + ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "expert_choice": filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -7342,9 +7878,9 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) if model_size == "smarter_expert": allowed_missing = { "layers.10.shared_up.weight", "layers.10.shared_down.weight", - "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", + "layers.10.shared_norm.weight", "layers.10.shared_norm.bias", "layers.10.shared_scale", "layers.10.expert_bias", - "layers.10.gate.weight", "layers.10.alpha", + "layers.10.gate.weight", "layers.10.alpha", "layers.10.calibration.weight", "layers.10.calibration.bias", "layers.10.theta", } for i in range(8): @@ -7354,14 +7890,14 @@ def load_weights_for_model(model: nn.Module, state_dict: dict, model_size: str) allowed_missing.add(f"layers.10.expert_norms.{i}.bias") allowed_missing.add(f"layers.10.lora_down.{i}.weight") allowed_missing.add(f"layers.10.lora_up.{i}.weight") - + ckpt_size = detect_model_size_from_state_dict(state_dict) if ckpt_size != "smarter_expert": filtered_sd = {k: v for k, v in state_dict.items() if not (k.startswith("layers.10.") and k not in ["layers.10.weight", "layers.10.bias"])} incompatible = model.load_state_dict(filtered_sd, strict=False) else: incompatible = model.load_state_dict(state_dict, strict=False) - + missing_filtered = [k for k in incompatible.missing_keys if k and k not in allowed_missing] unexpected_filtered = [k for k in incompatible.unexpected_keys if k] return missing_filtered, unexpected_filtered @@ -7874,7 +8410,7 @@ def detect_model_size_from_state_dict(state_dict: dict) -> str: if "layers.10.expert_norms.0.weight" in state_dict: return "expert_choice" return "ultra_expert" - + if "layers.10.halt_gates.0.weight" in state_dict and "layers.10.proposal_gates.0.weight" in state_dict: return "metacognitive_expert" if "layers.10.value_net.0.weight" in state_dict and "layers.10.action_up.0.weight" in state_dict: @@ -7883,6 +8419,8 @@ def detect_model_size_from_state_dict(state_dict: dict) -> str: return "omniscient_expert" if "layers.10.hypernet.0.weight" in state_dict and "layers.10.memory_matrix" in state_dict: return "cognitive_expert" + if "layers.10.core_router.weight" in state_dict and "layers.10.cross_attn.qkv.weight" in state_dict: + return "cognitive_leap_ultra_expert" if "layers.10.recur_core.0.weight" in state_dict and "layers.10.halt_head.weight" in state_dict: return "cognitive_leap_expert" if "layers.10.denoiser.0.0.weight" in state_dict and "layers.10.world_transition.0.weight" in state_dict: diff --git a/source/sync_runtime_model_variants.py b/source/sync_runtime_model_variants.py new file mode 100644 index 000000000..9fa427585 --- /dev/null +++ b/source/sync_runtime_model_variants.py @@ -0,0 +1,76 @@ +"""Synchronize the self-contained runtime model-variant snapshot. + +The packaged ``runtime_python`` surface must not import back into ``source``: +that works in a checkout but fails as soon as the runtime directory is copied or +shipped on its own. This helper keeps the snapshot mechanical and makes drift a +cheap CI/test check. +""" + +from __future__ import annotations + +import argparse +import hashlib +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SOURCE_PATH = REPO_ROOT / "source" / "model_variants.py" +RUNTIME_PATH = REPO_ROOT / "runtime_python" / "model_variants.py" + + +def _digest(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def snapshot_status() -> tuple[bool, str]: + source_payload = SOURCE_PATH.read_bytes() + runtime_payload = RUNTIME_PATH.read_bytes() if RUNTIME_PATH.exists() else b"" + matches = source_payload == runtime_payload + detail = ( + f"source={_digest(source_payload)[:12]} " + f"runtime={_digest(runtime_payload)[:12] if runtime_payload else 'missing'} " + f"bytes={len(source_payload)}" + ) + return matches, detail + + +def sync_snapshot() -> str: + source_payload = SOURCE_PATH.read_bytes() + RUNTIME_PATH.parent.mkdir(parents=True, exist_ok=True) + temporary_path = RUNTIME_PATH.with_suffix(".py.tmp") + temporary_path.write_bytes(source_payload) + temporary_path.replace(RUNTIME_PATH) + return f"synced {RUNTIME_PATH.relative_to(REPO_ROOT)} ({len(source_payload)} bytes)" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--check", + action="store_true", + help="Fail when runtime_python/model_variants.py differs from the source snapshot.", + ) + args = parser.parse_args() + + matches, detail = snapshot_status() + if args.check: + if matches: + print(f"runtime model-variant snapshot is current: {detail}") + return 0 + print(f"runtime model-variant snapshot is stale: {detail}") + return 1 + + if matches: + print(f"runtime model-variant snapshot already current: {detail}") + return 0 + print(sync_snapshot()) + matches, detail = snapshot_status() + if not matches: + print(f"runtime model-variant snapshot verification failed: {detail}") + return 1 + print(f"runtime model-variant snapshot verified: {detail}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test_cognitive_leap_ultra_expert.py b/test_cognitive_leap_ultra_expert.py new file mode 100644 index 000000000..5b115428e --- /dev/null +++ b/test_cognitive_leap_ultra_expert.py @@ -0,0 +1,201 @@ +import torch +import torch.nn as nn +import sys +import os + +sys.path.append(os.path.join(os.getcwd(), 'source')) + +from model_variants import ChampionNetCognitiveLeapUltraExpert + +def test_cognitive_leap_ultra(): + print("Starting smoke test for Cognitive Leap Ultra Expert (v51) architecture...") + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"Using device: {device}") + + # Initialize model with 4 cores + model = ChampionNetCognitiveLeapUltraExpert( + latent_dim=64, core_dim=128, n_cycles=3, inner_steps=2, n_cores=4 + ).to(device) + model.train() + + # Force alpha to 1.0 for gradient visibility + for name, param in model.named_parameters(): + if name == 'layers.10.alpha': + param.data.fill_(1.0) + print(f"Set {name} to 1.0") + + dummy_input = torch.randn(2, 1, 128).to(device) + print(f"Dummy input shape: {dummy_input.shape}") + + # 1. Forward Pass (Training) + print("Running forward pass (train mode)...") + logits = model(dummy_input) + print(f"Output shape: {logits.shape}") + assert logits.shape == (2, 1, 10), f"Expected shape (2, 1, 10), got {logits.shape}" + + # Track diagnostics from the recursive head + head = model.layers[10] + ponder_cost = head.last_ponder_cost + consistency_loss = head.last_consistency_loss + gating_entropy = head.last_gating_entropy + print(f"Ponder cost (expected cycles): {ponder_cost.item()}") + print(f"Latent consistency loss: {consistency_loss.item()}") + print(f"Gating entropy: {gating_entropy.item()}") + assert ponder_cost.item() >= 1.0, "Ponder cost should be at least one cycle" + assert ponder_cost.item() <= 3.0, "Ponder cost cannot exceed n_cycles" + assert consistency_loss.item() > 0, "Consistency loss should be positive" + assert gating_entropy.item() > 0, "Gating entropy should be positive" + + # 2. Backward Pass + print("Running backward pass...") + target = torch.randint(0, 10, (2, 1)).to(device) + ce_loss = nn.functional.cross_entropy(logits.view(-1, 10), target.view(-1)) + + # Deep improvement supervision over the cached per-cycle decodes + dis_loss = model.deep_supervision_loss(target) + print(f"Deep supervision loss: {dis_loss.item()}") + assert dis_loss.item() > 0, "Deep supervision loss should be positive" + + # Include all paths in loss + total_loss = ce_loss + 0.05 * consistency_loss + 0.1 * dis_loss - 0.01 * gating_entropy + print(f"Total Loss value: {total_loss.item()}") + total_loss.backward() + + print("Checking gradients...") + # Check keys representing all core components of v51 + checks = { + 'thought_init': False, # Scratchpad initializer + 'answer_init': False, # Answer initializer + 'cores.0': False, # Recurrent cores + 'cores.3': False, + 'core_router': False, # Core gating router + 'cross_attn.qkv': False, # Cross attention QKV + 'cross_attn.proj': False, # Cross attention output projection + 'answer_update.0': False, # Outer answer refinement + 'cycle_embed': False, # Cycle identity conditioning + 'halt_head': False, # ACT halting head + 'decode_head': False, # Decoder + 'shared_up': False, # Base shared expert + } + + for name, param in model.named_parameters(): + for key in checks: + if key in name and param.grad is not None and param.grad.abs().sum() > 0: + checks[key] = True + + for key, found in checks.items(): + status = "OK" if found else "CRITICAL: NO GRADIENT!" + print(f" {key}: {status}") + + all_ok = all(checks.values()) + if not all_ok: + failed = [k for k, v in checks.items() if not v] + raise AssertionError(f"Missing gradients in: {failed}") + + print("All gradients verified successfully.") + + # 3. Test-time compute scaling + print("Checking test-time compute scaling...") + model.eval() + with torch.no_grad(): + out_shallow = model(dummy_input, reasoning_cycles=1) + out_deep = model(dummy_input, reasoning_cycles=8) + assert out_shallow.shape == out_deep.shape == (2, 1, 10) + diff = (out_shallow - out_deep).abs().max().item() + print(f"Max |logit| difference between 1 and 8 cycles: {diff}") + assert diff > 0, "Extra reasoning cycles should change the output" + + # 4. Confidence-based early exit (Entropy Exit) + print("Checking entropy-based early exit...") + head = model.layers[10] + + # Evaluate with very high entropy threshold (so it exits immediately) + with torch.no_grad(): + model(dummy_input, reasoning_cycles=8, adaptive_compute=True, exit_entropy_threshold=10.0) + cycles_used_early = head.last_cycles_used.item() + print(f"Cycles used with high entropy exit: {cycles_used_early} / 8 requested") + assert cycles_used_early < 8, "High entropy exit threshold should stop cycling early" + + # Evaluate with zero entropy threshold (so it never exits early by entropy) + with torch.no_grad(): + model( + dummy_input, + reasoning_cycles=8, + adaptive_compute=True, + exit_entropy_threshold=0.0, + exit_tol=0.0, + prediction_stability_patience=0, + ) + cycles_used_late = head.last_cycles_used.item() + print(f"Cycles used with zero thresholds: {cycles_used_late} / 8 requested") + assert cycles_used_late == 8, "Zero threshold should never exit early" + + print("Entropy-based early exit verified.") + print("\nSmoke test PASSED!") + + +def test_prediction_stability_verifier_controls_early_exit(): + torch.manual_seed(5102) + model = ChampionNetCognitiveLeapUltraExpert( + latent_dim=32, + core_dim=64, + n_cycles=3, + inner_steps=1, + n_cores=2, + dropout=0.0, + ).eval() + sample = torch.zeros(1, 1, 128) + head = model.layers[10] + + with torch.no_grad(): + stable_output = model( + sample, + reasoning_cycles=8, + adaptive_compute=True, + exit_tol=0.0, + exit_entropy_threshold=0.0, + prediction_stability_patience=2, + prediction_stability_tol=1.0, + prediction_stability_top_k=4, + ) + + assert stable_output.shape == (1, 1, 10) + assert head.last_cycles_used.item() == 2 + assert head.last_exit_reason == "prediction_stable" + assert head.last_prediction_streak.item() >= 2 + assert head.last_prediction_confidence_delta.item() >= 0.0 + assert head.last_prediction_topk_js_divergence.item() >= 0.0 + assert ( + head.last_prediction_topk_js_divergence_max.item() + >= head.last_prediction_topk_js_divergence.item() + ) + assert torch.isfinite(head.last_prediction_topk_js_divergence) + + identical = torch.tensor([[0.6, 0.3, 0.1]]) + assert torch.equal( + head._topk_js_divergence(identical, identical, top_k=2), + torch.zeros(1), + ) + + with torch.no_grad(): + full_output = model( + sample, + reasoning_cycles=8, + adaptive_compute=True, + exit_tol=0.0, + exit_entropy_threshold=0.0, + prediction_stability_patience=0, + ) + + assert full_output.shape == stable_output.shape + assert head.last_cycles_used.item() == 8 + assert head.last_exit_reason in {"max_cycles", "halt_mass"} + +if __name__ == "__main__": + try: + test_cognitive_leap_ultra() + except Exception as e: + print(f"Smoke test FAILED: {e}") + import traceback + traceback.print_exc() + sys.exit(1) diff --git a/test_runtime_model_variants_parity.py b/test_runtime_model_variants_parity.py new file mode 100644 index 000000000..3b78dc05b --- /dev/null +++ b/test_runtime_model_variants_parity.py @@ -0,0 +1,40 @@ +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parent + + +def test_runtime_model_variants_snapshot_matches_source(): + completed = subprocess.run( + [sys.executable, "source/sync_runtime_model_variants.py", "--check"], + cwd=ROOT, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +def test_runtime_model_variants_imports_without_source_tree(tmp_path: Path): + isolated = tmp_path / "runtime_python" + isolated.mkdir() + for name in ("model_variants.py", "run.py"): + (isolated / name).write_bytes((ROOT / "runtime_python" / name).read_bytes()) + + script = f""" +import sys +sys.path.insert(0, {str(isolated)!r}) +import model_variants +assert "cognitive_leap_ultra_expert" in model_variants.SUPPORTED_MODEL_SIZES +model = model_variants.build_model(model_size="cognitive_leap_ultra_expert") +assert model is not None +print(type(model).__name__) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=tmp_path, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr diff --git a/test_v51_materialize_demo.py b/test_v51_materialize_demo.py new file mode 100644 index 000000000..0eb5bf2be --- /dev/null +++ b/test_v51_materialize_demo.py @@ -0,0 +1,22 @@ +import os +import sys +from pathlib import Path + + +sys.path.insert(0, os.path.join(os.getcwd(), "source")) + +from materialize_v51_chat_demo import DEFAULT_ARTIFACT_DIR, build_demo_metadata + + +def test_v51_demo_metadata_uses_portable_project_paths(): + weights = DEFAULT_ARTIFACT_DIR / "cognitive_leap_ultra_v51_trained.pth" + metrics = DEFAULT_ARTIFACT_DIR / "benchmark_results.json" + + payload = build_demo_metadata(weights=weights, metrics_path=metrics) + + assert payload["checkpoint_path"] == "output/benchmark_v51_cognitive_leap_ultra/cognitive_leap_ultra_v51_trained.pth" + assert payload["benchmark_metrics"] == "output/benchmark_v51_cognitive_leap_ultra/benchmark_results.json" + assert payload["runtime_defaults"]["prediction_stability_patience"] == 2 + assert payload["runtime_defaults"]["prediction_stability_tol"] == 0.005 + assert "New folder (4)" not in payload["checkpoint_path"] + assert "New folder (4)" not in payload["benchmark_metrics"] From 883ad73ede2400e03597b59ee287564908a8f6a1 Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:38:02 +0100 Subject: [PATCH 3/6] feat: advance omni collective training experiments --- docs/omni_collective_v41_blueprint.md | 185 ++ source/benchmark_all_models_common.py | 31 +- source/export_omni_collective_v8_preview.py | 248 +++ source/omni_collective_v41_model.py | 204 +++ source/omni_collective_v42_model.py | 232 +++ source/prepare_omni_collective_v41.py | 457 +++++ source/prepare_omni_collective_v42.py | 513 ++++++ source/run_omni_collective_v41_background.py | 81 + .../run_train_omni_collective_v41_local.ps1 | 59 + .../run_train_omni_collective_v42_local.ps1 | 59 + source/train_omni_collective_v41.py | 1487 +++++++++++++++++ source/train_omni_collective_v42.py | 1201 +++++++++++++ source/train_omni_collective_v8.py | 18 +- test_omni_collective_v41.py | 15 + test_omni_collective_v42.py | 15 + test_omni_collective_v8.py | 19 +- test_prepare_omni_collective_v41.py | 91 + test_prepare_omni_collective_v42.py | 82 + test_train_omni_collective_v41.py | 262 +++ 19 files changed, 5244 insertions(+), 15 deletions(-) create mode 100644 docs/omni_collective_v41_blueprint.md create mode 100644 source/export_omni_collective_v8_preview.py create mode 100644 source/omni_collective_v41_model.py create mode 100644 source/omni_collective_v42_model.py create mode 100644 source/prepare_omni_collective_v41.py create mode 100644 source/prepare_omni_collective_v42.py create mode 100644 source/run_omni_collective_v41_background.py create mode 100644 source/run_train_omni_collective_v41_local.ps1 create mode 100644 source/run_train_omni_collective_v42_local.ps1 create mode 100644 source/train_omni_collective_v41.py create mode 100644 source/train_omni_collective_v42.py create mode 100644 test_omni_collective_v41.py create mode 100644 test_omni_collective_v42.py create mode 100644 test_prepare_omni_collective_v41.py create mode 100644 test_prepare_omni_collective_v42.py create mode 100644 test_train_omni_collective_v41.py diff --git a/docs/omni_collective_v41_blueprint.md b/docs/omni_collective_v41_blueprint.md new file mode 100644 index 000000000..80097e7a3 --- /dev/null +++ b/docs/omni_collective_v41_blueprint.md @@ -0,0 +1,185 @@ +# Omni Collective V41 Blueprint + +This is the first concrete preparation pass for `omni_collective_v41`, built from the finished local `v8` artifact and a small set of recent primary-source training ideas. + +## Why V41 Exists + +`v8` ended at `112,025,979` parameters and improved multimodal fusion, but its weakest metric was still `response_accuracy = 0.0935` in stage 2. The sample outputs also show three clear problems: + +- human communication is still too brittle and templated +- grounded routing and model/tool choice still fail on specialist prompts +- coding/problem-solving is better than earlier omni lines, but not strong enough to dominate the local stack + +So `v41` should not just be `v8 + more rows`. It should become: + +- smarter on hard reasoning and coding prompts +- more creative without becoming less grounded +- more natural and useful when talking to humans +- better at choosing the right route, specialist, or tool before answering + +## Target Envelope + +- target params: around `126M` +- keep the model compact enough for local deployment +- raise stage-2 response quality first, not just overall score +- challenge `v40_benchmax` on reasoning while keeping the broader multimodal specialist edge of the omni line + +## Architecture Plan + +The parameter increase should be moderate: + +- `text_hidden`: `272 -> 288` +- `fusion_hidden`: `1088 -> 1152` +- `expert_count`: `10 -> 12` +- `expert_hidden`: `1792 -> 1920` +- `deliberation_passes`: `10 -> 12` + +New `v41` heads: + +- `latent_plan_slots` + Distill short hidden plan states from explicit planning traces, then keep the plans internal at inference. +- `multi_token_prediction_aux` + Predict short future spans to improve long-range coherence and information density. +- `communication_polish_head` + Separate answer correctness from presentation quality so the model can solve first, then rewrite more clearly. +- `uncertainty_anchor_head` + Reward transparent uncertainty and missing-evidence statements instead of invented specifics. + +Two more explicit `v41` ideas should now be treated as first-class: + +- `reflection memory shards` + Keep compact bad-attempt -> reflection -> repaired-answer traces for coding and reasoning tasks. +- `reasoning budget labels` + Supervise when a task deserves a short, medium, or deep solve instead of globally increasing thinking depth. + +## New Training Ideas + +These are the main `v41` innovations, grounded in recent work but adapted to the Supermix local-first setup: + +1. `Latent Plan Distillation` + Train on explicit micro-plans and route choices, but compress them into hidden plan slots before deployment. + +2. `Reasoning Budget Curriculum` + Train short, medium, and long budget targets so `v41` learns when to think longer instead of always answering at one fixed depth. + +3. `Disagreement Harvesting` + Mine prompts where `v8`, `v40_benchmax`, and the Qwen teachers disagree. Turn those into compare-and-justify rows instead of taking a single teacher answer as truth. + +4. `Code Critique -> Repair` + Add bug reports, failing tests, stack traces, and patch-review examples where the model must diagnose, propose, and refine a repair. + +5. `Two-Pass Human Communication` + First solve the problem. Then run a short communication-polish rewrite trained to improve clarity, brevity, empathy, and structure without changing the facts. + +6. `Creativity Rescue Pairs` + Use pairs of dull-correct answers and vivid-but-still-grounded answers so creativity is learned as controlled improvement rather than free hallucination. + +7. `Route-Then-Answer Supervision` + Train on hidden or explicit route-selection labels before the final answer. This should improve Auto mode, collective mode, and specialist handoff decisions. + +8. `Reflection Memory` + Keep concise post-mortems for failed code, route, and reasoning attempts, then train on the repaired response. + +9. `Reasoning Budget Curriculum` + Train the same task under short, medium, and deep answer budgets so the model learns when extra deliberation is actually worth it. + +## Data Mix Changes + +The data should become denser and more varied, not just larger: + +- `high_density_reasoning_core` + A compact, hard reasoning set with fewer repetitive examples and higher per-row value. +- `code_critique_repair` + Traceback, failing-test, patch, and critique loops for coding and debugging. +- `reflection_memory_rows` + Failure -> short reflection -> improved retry traces. +- `human_communication_polish` + Rewrites, comparisons, explanations, de-escalation, uncertainty framing, and direct-but-clear human answers. +- `teacher_disagreement_hardcases` + Prompts where strong local teachers disagree on answer, route, or framing. +- `grounded_uncertainty_rows` + Current-fact and low-evidence prompts that explicitly train “what I know / what I don’t know / what I need”. +- `creative_constraint_pairs` + Creativity with fidelity, style control, and factual discipline. +- `route_then_answer_supervision` + Tool/model choice before answer. +- `reasoning_budget_curriculum` + Same task solved under short, medium, and deep answer budgets. +- `specialist_dense_topups` + Materials, protein, 3D, vision, and native-image deltas that keep the omni stack broad. + +## Training Recipe + +Recommended `v41` phases: + +1. `Stage 0: latent-plan bootstrap` + Distill short explicit plan traces, route labels, and critique stubs into hidden plan slots. + +2. `Stage 1: teacher-heavy SFT` + Mix `v8`, `v40_benchmax`, `qwen_v28`, `qwen_v30`, and `v7` with harder coding/problem-solving weighting and denser reasoning rows. + +3. `Stage 2: self-play preference loop` + Generate chosen vs rejected pairs using self-play and teacher disagreement mining. + +4. `Stage 3: code refine` + Train on critique -> repair traces and keep only objectively improved fixes. + +5. `Stage 4: communication polish` + Finish with human-facing clarity, compare/explain prompts, and grounded uncertainty. + +6. `Stage 5: self-correction finish` + Run deterministic self-correction on weak drafts with fair prompts and keep only improved outputs. + +## Runtime and Reliability Defaults + +`v8` was trained on CPU with no AMP, no compile, and no accumulation. `v41` should not repeat that if avoidable. + +Recommended defaults: + +- prefer `cuda` or `dml`, not CPU +- `amp=auto` +- `compile_model=true` +- `grad_accum_steps=4` +- `ema_decay=0.9993` +- `warmup_ratio=0.06` +- checkpoint every `100` optimizer steps +- always keep rolling stage checkpoints and stage-complete snapshots +- preview-benchmark stage-2 checkpoints during the run + +## Success Gates + +Minimum goals: + +- stage-2 `response_accuracy >= 0.145` +- stage-2 score `>= 0.455` +- common benchmark exact `>= 0.25` +- must beat `omni_collective_v8_preview` and `omni_collective_v7` +- stretch goal: approach or beat `v40_benchmax` on benchmark-heavy reasoning while remaining better on specialist and conversational tasks + +Promotion protocol: + +- benchmark at least `3` preview checkpoints during stage 2 +- do not promote a final model that regresses vision or domain accuracy by more than `0.02` vs `v8` +- require a clean win on the new communication, code-repair, and grounded-uncertainty eval packs before promotion + +## Primary-Source Inspirations + +These papers are being used as inspirations, not copied recipes: + +- DeepSeek-V3 Technical Report: https://arxiv.org/abs/2412.19437 +- Better & Faster Large Language Models via Multi-token Prediction: https://arxiv.org/abs/2404.19737 +- s1: Simple test-time scaling: https://arxiv.org/abs/2501.19393 +- LIMO: Less Is More for Reasoning: https://arxiv.org/abs/2502.03387 +- RefineCoder: https://arxiv.org/abs/2502.09183 +- Reflexion: https://arxiv.org/abs/2303.11366 +- Self-Play Fine-Tuning (SPIN): https://arxiv.org/abs/2401.01335 +- Self-Refine: https://arxiv.org/abs/2303.17651 +- Large Language Models have Intrinsic Self-Correction Ability: https://arxiv.org/abs/2406.15673 +- RLAIF: https://arxiv.org/abs/2309.00267 + +## Immediate Next Steps + +- add a `v41` prep script that emits a machine-readable blueprint from the latest `v8` summary +- build new dataset builders for disagreement mining, communication polish pairs, and route-then-answer supervision +- extend the omni model code with latent-plan, communication-polish, and uncertainty-anchor heads +- keep the `v8`-style resume and preview export path from the first training run onward diff --git a/source/benchmark_all_models_common.py b/source/benchmark_all_models_common.py index 3a9fc7da9..39d50c3d8 100644 --- a/source/benchmark_all_models_common.py +++ b/source/benchmark_all_models_common.py @@ -79,9 +79,6 @@ def token_f1(reference: object, prediction: object) -> float: from omni_collective_v6_model import OmniCollectiveEngineV6 from omni_collective_v7_model import OmniCollectiveEngineV7 from omni_collective_v8_model import OmniCollectiveEngineV8 - from omni_collective_v46_model import OmniCollectiveEngineV46 - from omni_collective_v42_model import OmniCollectiveEngineV42 - from omni_collective_v41_model import OmniCollectiveEngineV41 from protein_folding_model import ProteinFoldingEngine from three_d_generation_model import ThreeDGenerationEngine except ImportError: # pragma: no cover @@ -89,12 +86,34 @@ def token_f1(reference: object, prediction: object) -> float: from .omni_collective_v6_model import OmniCollectiveEngineV6 from .omni_collective_v7_model import OmniCollectiveEngineV7 from .omni_collective_v8_model import OmniCollectiveEngineV8 - from .omni_collective_v46_model import OmniCollectiveEngineV46 - from .omni_collective_v42_model import OmniCollectiveEngineV42 - from .omni_collective_v41_model import OmniCollectiveEngineV41 from .protein_folding_model import ProteinFoldingEngine from .three_d_generation_model import ThreeDGenerationEngine +# Optional V4x components +try: + from omni_collective_v46_model import OmniCollectiveEngineV46 +except ImportError: + try: + from .omni_collective_v46_model import OmniCollectiveEngineV46 + except ImportError: + OmniCollectiveEngineV46 = None + +try: + from omni_collective_v42_model import OmniCollectiveEngineV42 +except ImportError: + try: + from .omni_collective_v42_model import OmniCollectiveEngineV42 + except ImportError: + OmniCollectiveEngineV42 = None + +try: + from omni_collective_v41_model import OmniCollectiveEngineV41 +except ImportError: + try: + from .omni_collective_v41_model import OmniCollectiveEngineV41 + except ImportError: + OmniCollectiveEngineV41 = None + NUMBER_RE = re.compile(r"-?\d[\d,]*(?:\.\d+)?") FINAL_ANSWER_RE = re.compile(r"final answer\s*:\s*([^\n\r]+)", re.IGNORECASE) diff --git a/source/export_omni_collective_v8_preview.py b/source/export_omni_collective_v8_preview.py new file mode 100644 index 000000000..9e975f187 --- /dev/null +++ b/source/export_omni_collective_v8_preview.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import shutil +import sys +import time +import zipfile +from datetime import datetime +from pathlib import Path +from typing import Dict, Optional, Tuple + +import torch + +SOURCE_DIR = Path(__file__).resolve().parent +if str(SOURCE_DIR) not in sys.path: + sys.path.append(str(SOURCE_DIR)) + +try: + import train_omni_collective_v8 as train_v8 + from omni_collective_v8_model import OmniCollectiveEngineV8, OmniCollectiveNetV8 +except ImportError: # pragma: no cover + from . import train_omni_collective_v8 as train_v8 + from .omni_collective_v8_model import OmniCollectiveEngineV8, OmniCollectiveNetV8 + + +def _resolve_checkpoint_from_state(state_path: Path) -> Tuple[Dict[str, object], Path]: + state = json.loads(state_path.read_text(encoding="utf-8")) + checkpoint_path = Path(str(state["checkpoint_path"])).resolve() + return state, checkpoint_path + + +def _copy_snapshot(source_path: Path, destination_dir: Path) -> Path: + destination_dir.mkdir(parents=True, exist_ok=True) + snapshot_path = destination_dir / source_path.name + last_error: Optional[Exception] = None + for _ in range(8): + try: + shutil.copy2(source_path, snapshot_path) + return snapshot_path + except Exception as exc: # pragma: no cover - windows file races are timing-dependent + last_error = exc + time.sleep(0.5) + raise RuntimeError(f"Failed to snapshot checkpoint {source_path}: {last_error}") + + +def _select_state_dict(checkpoint: Dict[str, object]) -> Dict[str, torch.Tensor]: + for key in ("ema_state", "best_state", "model_state"): + value = checkpoint.get(key) + if isinstance(value, dict) and value: + return value + raise RuntimeError("Checkpoint did not contain a usable model state") + + +def export_preview( + *, + repo_root: Path, + output_dir: Path, + models_dir: Path, + images_dir: Path, + state_path: Path, + checkpoint_path: Optional[Path], + seed: int, + distill_limit: int, + teacher_model_limit: int, + device: str, +) -> Dict[str, object]: + if checkpoint_path is None: + state_payload, checkpoint_path = _resolve_checkpoint_from_state(state_path) + else: + state_payload = json.loads(state_path.read_text(encoding="utf-8")) if state_path.exists() else {} + checkpoint_path = checkpoint_path.resolve() + + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + artifact_dir = (output_dir / f"supermix_omni_collective_v8_preview_{stamp}").resolve() + artifact_dir.mkdir(parents=True, exist_ok=True) + checkpoint_snapshot = _copy_snapshot(checkpoint_path, artifact_dir) + checkpoint = torch.load(checkpoint_snapshot, map_location="cpu") + state_dict = _select_state_dict(checkpoint) + + _, full_rows, dataset_summary = train_v8.build_training_rows( + repo_root=repo_root, + models_dir=models_dir, + images_dir=images_dir, + seed=int(seed), + distill_limit=int(distill_limit), + teacher_model_limit=int(teacher_model_limit), + ) + train_rows, _ = train_v8.split_rows(full_rows, seed=int(seed) + 17) + vocab = train_v8.build_char_vocab([row.prompt for row in train_rows], min_frequency=1) + response_bank = sorted({row.response_text for row in full_rows if row.response_text}) + + model = OmniCollectiveNetV8( + vocab_size=max(len(vocab), 2), + num_intents=len(train_v8.OMNI_INTENTS_V2), + num_responses=max(len(response_bank), 1), + num_vision_classes=len(train_v8.SCIENCE_IMAGE_CLASSES), + num_domains=len(train_v8.OMNI_DOMAIN_LABELS_V2), + base_embed_dim=132, + text_hidden=272, + image_channels=56, + word_buckets=16384, + word_embed_dim=120, + deep_text_channels=384, + deep_image_channels=128, + fusion_hidden=1088, + memory_slots=28, + depth_steps=11, + expert_count=10, + expert_hidden=1792, + context_top_k=4, + expert_top_k=2, + ) + model.load_state_dict(state_dict) + model.eval() + + weights_path = artifact_dir / "omni_collective_v8_preview.pth" + meta_path = artifact_dir / "omni_collective_v8_preview_meta.json" + summary_path = artifact_dir / "omni_collective_v8_preview_summary.json" + zip_path = output_dir / f"supermix_omni_collective_v8_preview_{stamp}.zip" + + torch.save(model.state_dict(), weights_path) + parameter_count = int(sum(parameter.numel() for parameter in model.parameters())) + class_info = train_v8.ensure_base_images(images_dir) + + stage_resume_dir = checkpoint_path.parent + stage1_meta_path = stage_resume_dir / "stage1_complete.json" + stage1_meta = json.loads(stage1_meta_path.read_text(encoding="utf-8")) if stage1_meta_path.exists() else None + preview_state = { + "stage_name": str(checkpoint.get("stage_name") or ""), + "epoch": int(checkpoint.get("epoch") or 0), + "next_batch_index": int(checkpoint.get("next_batch_index") or 0), + "optimizer_steps_done": int(checkpoint.get("optimizer_steps_done") or 0), + "saved_at": str(checkpoint.get("saved_at") or ""), + } + meta = { + "architecture_version": 8, + "preview": True, + "preview_source": { + "state_path": str(state_path), + "checkpoint_path": str(checkpoint_path), + "checkpoint_snapshot": str(checkpoint_snapshot), + "state_payload": state_payload, + "checkpoint_preview_state": preview_state, + }, + "vocab": vocab, + "response_bank": response_bank, + "class_info": class_info, + "intent_labels": list(train_v8.OMNI_INTENTS_V2), + "domain_labels": list(train_v8.OMNI_DOMAIN_LABELS_V2), + "max_len": 384, + "image_size": 128, + "embed_dim": 132, + "text_hidden": 272, + "image_channels": 56, + "word_buckets": 16384, + "max_words": 84, + "word_embed_dim": 120, + "deep_text_channels": 384, + "deep_image_channels": 128, + "fusion_hidden": 1088, + "memory_slots": 28, + "depth_steps": 11, + "expert_count": 10, + "expert_hidden": 1792, + "context_top_k": 4, + "expert_top_k": 2, + "parameter_count": parameter_count, + "seed": int(seed), + "stage1": stage1_meta, + "stage2_preview": preview_state, + "dataset_summary": dataset_summary, + "deliberation_passes": 10, + "minimum_passes": 5, + "grounding_threshold": 0.54, + "prompt_understanding_mode": "all_model_multitype_grounded_consensus_math_protein_materials_three_d_conversation", + "notes": [ + "Preview export cut from the latest resumable stage2 checkpoint while the main training run continues.", + "This artifact is for benchmarking and inspection only; the full run remains the source of truth.", + ], + } + meta_path.write_text(json.dumps(meta, indent=2, ensure_ascii=True), encoding="utf-8") + + engine = OmniCollectiveEngineV8(weights_path=weights_path, meta_path=meta_path, device=torch.device(device)) + sample_prompts = [ + "Reply in exactly two sentences explaining why regression tests matter.", + "Solve 3*x + 7 = 19.", + "Which local model is best for benchmark-focused reasoning prompts?", + "Write a tiny OpenSCAD snippet for a centered cylinder with a hole.", + ] + summary = { + "artifact": zip_path.name, + "parameter_count": parameter_count, + "preview": True, + "dataset_summary": dataset_summary, + "preview_source": meta["preview_source"], + "sample_outputs": [{"prompt": prompt, "answer": engine.answer(prompt)} for prompt in sample_prompts], + } + summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=True), encoding="utf-8") + + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as archive: + archive.write(weights_path, arcname=weights_path.name) + archive.write(meta_path, arcname=meta_path.name) + archive.write(summary_path, arcname=summary_path.name) + archive.write(checkpoint_snapshot, arcname=checkpoint_snapshot.name) + + return { + "artifact_dir": str(artifact_dir), + "zip_path": str(zip_path), + "checkpoint_snapshot": str(checkpoint_snapshot), + "preview_batch_index": int(checkpoint.get("next_batch_index") or 0), + "preview_saved_at": str(checkpoint.get("saved_at") or ""), + "parameter_count": parameter_count, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Export a benchmarkable omni_collective_v8 preview from the latest resumable checkpoint.") + parser.add_argument("--repo_root", default=".") + parser.add_argument("--output_dir", default="output") + parser.add_argument("--models_dir", default=str(train_v8.DEFAULT_MODELS_DIR)) + parser.add_argument("--images_dir", default="output/science_vision_dataset/images") + parser.add_argument("--state_path", default="output/omni_collective_v8_train_state.json") + parser.add_argument("--checkpoint_path", default="") + parser.add_argument("--seed", type=int, default=403) + parser.add_argument("--distill_limit", type=int, default=160) + parser.add_argument("--teacher_model_limit", type=int, default=0) + parser.add_argument("--device", default="cpu") + args = parser.parse_args() + + result = export_preview( + repo_root=Path(args.repo_root).resolve(), + output_dir=Path(args.output_dir).resolve(), + models_dir=Path(args.models_dir).resolve(), + images_dir=Path(args.images_dir).resolve(), + state_path=Path(args.state_path).resolve(), + checkpoint_path=Path(args.checkpoint_path).resolve() if str(args.checkpoint_path).strip() else None, + seed=int(args.seed), + distill_limit=int(args.distill_limit), + teacher_model_limit=int(args.teacher_model_limit), + device=str(args.device), + ) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/source/omni_collective_v41_model.py b/source/omni_collective_v41_model.py new file mode 100644 index 000000000..c4a06b416 --- /dev/null +++ b/source/omni_collective_v41_model.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +import torch + +try: + from omni_collective_v8_model import OmniCollectiveEngineV8, OmniCollectiveNetV8, OmniPredictionV8, _prompt_variants_v8 + from math_equation_model import heuristic_intent, solve_intent + from image_feature_utils import describe_image_for_text + from image_recognition_model import CLASS_LABELS, SCIENCE_IMAGE_CLASSES +except ImportError: # pragma: no cover + from .omni_collective_v8_model import OmniCollectiveEngineV8, OmniCollectiveNetV8, OmniPredictionV8, _prompt_variants_v8 + from .math_equation_model import heuristic_intent, solve_intent + from .image_feature_utils import describe_image_for_text + from .image_recognition_model import CLASS_LABELS, SCIENCE_IMAGE_CLASSES + + +OmniCollectiveNetV41 = OmniCollectiveNetV8 + + +def _prompt_variants_v41( + prompt: str, + draft_answer: str = "", + *, + response_confidence: float = 0.0, + domain_confidence: float = 0.0, +) -> List[str]: + variants = list( + _prompt_variants_v8( + prompt, + draft_answer, + response_confidence=response_confidence, + domain_confidence=domain_confidence, + ) + ) + prompt_text = " ".join(str(prompt or "").split()) + lowered = prompt_text.lower() + extras = [ + ( + "Create a short hidden plan before answering.\n" + f"Request: {prompt_text}\n" + "Identify the route, the evidence you actually have, the likely failure mode, and the clearest final form of the answer." + ), + ( + "Run a reflection pass before finalizing the answer.\n" + f"Request: {prompt_text}\n" + "Look for the most likely mistake in your first instinct, then keep the supported core and repair the weak parts." + ), + ( + "Optimize the answer for a human reader.\n" + f"Request: {prompt_text}\n" + "Prefer directness, clarity, and useful structure while preserving the grounded substance." + ), + ] + if any(token in lowered for token in ("python", "code", "bug", "traceback", "test", "stack trace", "refactor", "patch")): + extras.append( + ( + "This is a coding or debugging request.\n" + f"Request: {prompt_text}\n" + "Simulate a code-review and repair pass before answering. Prefer exact fixes, concrete reasoning, and reproducible guidance." + ) + ) + if any(token in lowered for token in ("creative", "story", "poem", "vivid", "rewrite", "tone", "paragraph")): + extras.append( + ( + "This answer should be more vivid and human-friendly without becoming ungrounded.\n" + f"Request: {prompt_text}\n" + "Prefer controlled creativity over bland filler or invented specifics." + ) + ) + if any(token in lowered for token in ("which model", "auto mode", "collective", "agent mode", "route", "selector")): + extras.append( + ( + "Choose the route before answering.\n" + f"Request: {prompt_text}\n" + "Pick the best local specialist or generalist in your head first, then answer from that perspective." + ) + ) + + deduped: List[str] = [] + seen = set() + for item in variants + extras: + cooked = str(item).strip() + if cooked and cooked not in seen: + deduped.append(cooked) + seen.add(cooked) + return deduped + + +@dataclass +class OmniPredictionV41(OmniPredictionV8): + reasoning_mode: str = "latent_plan_reflective_multitype_consensus_v41" + + +class OmniCollectiveEngineV41(OmniCollectiveEngineV8): + def __init__(self, *, weights_path: Path, meta_path: Path, device: Optional[torch.device] = None) -> None: + super().__init__(weights_path=weights_path, meta_path=meta_path, device=device) + self.deliberation_passes = int(self.meta.get("deliberation_passes") or 12) + self.minimum_passes = int(self.meta.get("minimum_passes") or 6) + self.grounding_threshold = float(self.meta.get("grounding_threshold") or 0.56) + + def predict(self, prompt: str, image_path: Optional[str] = None) -> OmniPredictionV41: + prompt_text = str(prompt or "").strip() + image_tensor, has_image = self._encode_image(image_path) + raw_outputs = self._run_prompt(prompt_text, image_tensor, has_image) + response_conf, domain_conf = self._confidence_from_outputs(raw_outputs) + draft_idx = int(raw_outputs["response"].argmax(dim=0).item()) if self.responses else 0 + draft_answer = self.responses[draft_idx] if self.responses else "" + variants = _prompt_variants_v41( + prompt_text, + draft_answer, + response_confidence=response_conf, + domain_confidence=domain_conf, + ) + + weighted_outputs: List[tuple[float, Dict[str, torch.Tensor]]] = [] + combined = raw_outputs + for index, variant in enumerate(variants, start=1): + weight = 1.22 if index == 1 else 1.0 + lowered = variant.lower() + if "hidden plan" in lowered: + weight = 1.08 + elif "reflection pass" in lowered: + weight = 1.09 + elif "human reader" in lowered: + weight = 1.06 + elif "code-review and repair" in lowered: + weight = 1.08 + elif "choose the route before answering" in lowered: + weight = 1.07 + weighted_outputs.append((weight, self._run_prompt(variant, image_tensor, has_image))) + combined = self._combine_outputs(weighted_outputs) + response_conf, domain_conf = self._confidence_from_outputs(combined) + if index >= self.minimum_passes and response_conf >= self.grounding_threshold and domain_conf >= 0.48: + break + if index >= self.deliberation_passes: + break + + outputs = combined + intent_probs = torch.softmax(outputs["intent"], dim=0) + response_probs = torch.softmax(outputs["response"], dim=0) if self.responses else torch.ones(1) + vision_probs = torch.softmax(outputs["vision"], dim=0) + domain_probs = torch.softmax(outputs["domain"], dim=0) + + intent_idx = int(intent_probs.argmax(dim=0).item()) if intent_probs.numel() else 0 + response_idx = int(response_probs.argmax(dim=0).item()) if response_probs.numel() else 0 + domain_idx = int(domain_probs.argmax(dim=0).item()) if domain_probs.numel() else 0 + values, indices = torch.topk(vision_probs, k=min(3, len(SCIENCE_IMAGE_CLASSES))) + vision_key = SCIENCE_IMAGE_CLASSES[int(indices[0].item())] if SCIENCE_IMAGE_CLASSES else "" + top_vision = [ + {"concept": SCIENCE_IMAGE_CLASSES[int(index.item())], "confidence": round(float(value.item()), 4)} + for value, index in zip(values, indices) + if SCIENCE_IMAGE_CLASSES + ] + return OmniPredictionV41( + intent=self.intent_labels[min(intent_idx, len(self.intent_labels) - 1)] if self.intent_labels else "general", + response_text=self.responses[response_idx] if self.responses else "", + vision_concept=vision_key, + vision_confidence=float(values[0].item()) if values.numel() else 0.0, + top_vision=top_vision, + domain=self.domain_labels[min(domain_idx, len(self.domain_labels) - 1)] if self.domain_labels else "general", + domain_confidence=float(domain_probs[domain_idx].item()) if domain_probs.numel() else 0.0, + response_confidence=float(response_probs[response_idx].item()) if response_probs.numel() else 0.0, + reasoning_passes=len(weighted_outputs), + ) + + def answer(self, prompt: str, image_path: Optional[str] = None) -> str: + prompt_text = str(prompt or "").strip() + math_intent = heuristic_intent(prompt_text) + prediction = self.predict(prompt_text, image_path=image_path) + if math_intent: + try: + solved = solve_intent(prompt_text, math_intent) + return str(solved.get("response") or prediction.response_text or "Math intent detected.") + except Exception: + return prediction.response_text or "This looks like a math problem, but I could not solve it cleanly." + if image_path: + info = self.class_info.get(prediction.vision_concept, {}) + descriptor = describe_image_for_text(str(image_path)) + lines = [ + f"Image read: {CLASS_LABELS.get(prediction.vision_concept, prediction.vision_concept.replace('_', ' '))} ({prediction.vision_confidence * 100:.1f}% confidence).", + ] + if str(info.get("caption") or "").strip(): + lines.append(f"Likely concept: {str(info.get('caption')).strip()}.") + lines.append(f"Visual descriptor: {descriptor}.") + if prediction.response_text: + lines.append(prediction.response_text) + return "\n".join(lines) + if prediction.response_confidence < self.grounding_threshold and prediction.domain in {"knowledge", "general", "planning", "language", "model_selection"}: + if prediction.response_text: + return f"Best grounded answer from my local training: {prediction.response_text}" + return "I do not have enough grounded evidence in my local training for a confident answer." + return prediction.response_text or "I did not have a strong fused answer for that prompt." + + +__all__ = [ + "OmniCollectiveEngineV41", + "OmniCollectiveNetV41", + "OmniPredictionV41", + "_prompt_variants_v41", +] diff --git a/source/omni_collective_v42_model.py b/source/omni_collective_v42_model.py new file mode 100644 index 000000000..7621ea072 --- /dev/null +++ b/source/omni_collective_v42_model.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +import torch + +try: + from omni_collective_v41_model import OmniCollectiveEngineV41, OmniCollectiveNetV41, OmniPredictionV41, _prompt_variants_v41 + from image_recognition_model import SCIENCE_IMAGE_CLASSES +except ImportError: # pragma: no cover + from .omni_collective_v41_model import OmniCollectiveEngineV41, OmniCollectiveNetV41, OmniPredictionV41, _prompt_variants_v41 + from .image_recognition_model import SCIENCE_IMAGE_CLASSES + + +OmniCollectiveNetV42 = OmniCollectiveNetV41 + + +def _budget_hint_v42(prompt: str, *, response_confidence: float = 0.0, domain_confidence: float = 0.0) -> str: + lowered = str(prompt or "").lower() + deep_markers = ( + "benchmark", + "prove", + "derive", + "why does", + "gsm8k", + "mmlu", + "arc", + "hellaswag", + "debug", + "refactor", + "traceback", + "test failure", + "compare", + "choose", + ) + medium_markers = ( + "summarize", + "rewrite", + "plan", + "route", + "which model", + "auto mode", + "collective", + "image", + "video", + "audio", + "ocr", + ) + if response_confidence < 0.50 or domain_confidence < 0.44 or any(token in lowered for token in deep_markers): + return "deep" + if any(token in lowered for token in medium_markers): + return "medium" + return "short" + + +def _prompt_variants_v42( + prompt: str, + draft_answer: str = "", + *, + response_confidence: float = 0.0, + domain_confidence: float = 0.0, +) -> List[str]: + variants = list( + _prompt_variants_v41( + prompt, + draft_answer, + response_confidence=response_confidence, + domain_confidence=domain_confidence, + ) + ) + prompt_text = " ".join(str(prompt or "").split()) + lowered = prompt_text.lower() + budget = _budget_hint_v42( + prompt_text, + response_confidence=response_confidence, + domain_confidence=domain_confidence, + ) + extras = [ + ( + f"Choose a {budget} reasoning budget before answering.\n" + f"Request: {prompt_text}\n" + "Keep the answer direct, but spend more thinking only if the request is actually hard." + ), + ( + "Decide whether you need a verifier pass before finalizing.\n" + f"Request: {prompt_text}\n" + "If the first answer could fail on correctness, route choice, or code accuracy, mentally verify it before returning the final answer." + ), + ( + "Prefer the smallest evidence set that still justifies the answer.\n" + f"Request: {prompt_text}\n" + "Do not overuse context or explanation length when a shorter grounded answer is enough." + ), + ] + if any(token in lowered for token in ("benchmark", "boolq", "piqa", "hellaswag", "mmlu", "arc", "gsm8k")): + extras.append( + ( + "This is benchmark-style reasoning.\n" + f"Request: {prompt_text}\n" + "Use benchmark-specialist discipline first: identify the exact task type, eliminate weak options, then answer without filler." + ) + ) + if any(token in lowered for token in ("python", "code", "bug", "traceback", "refactor", "patch", "test")): + extras.append( + ( + "This is an agentic coding request.\n" + f"Request: {prompt_text}\n" + "Plan the edit, simulate a verifier or failing test, then return the repair-oriented answer." + ) + ) + if any(token in lowered for token in ("image", "video", "audio", "ocr", "diagram", "photo", "caption")): + extras.append( + ( + "This is a multimodal or omni-style request.\n" + f"Request: {prompt_text}\n" + "Think about the best grounded modality route first, then answer in a concise, evidence-aware form." + ) + ) + if any(token in lowered for token in ("long context", "many files", "huge", "large", "full log", "full transcript")): + extras.append( + ( + "This request may tempt overthinking.\n" + f"Request: {prompt_text}\n" + "Select only the most relevant evidence and stop once the answer is supported." + ) + ) + + deduped: List[str] = [] + seen = set() + for item in variants + extras: + cooked = str(item).strip() + if cooked and cooked not in seen: + deduped.append(cooked) + seen.add(cooked) + return deduped + + +@dataclass +class OmniPredictionV42(OmniPredictionV41): + reasoning_mode: str = "budgeted_route_first_verifier_polish_v42" + planned_budget: str = "short" + + +class OmniCollectiveEngineV42(OmniCollectiveEngineV41): + def __init__(self, *, weights_path: Path, meta_path: Path, device: Optional[torch.device] = None) -> None: + super().__init__(weights_path=weights_path, meta_path=meta_path, device=device) + self.deliberation_passes = int(self.meta.get("deliberation_passes") or 14) + self.minimum_passes = int(self.meta.get("minimum_passes") or 7) + self.grounding_threshold = float(self.meta.get("grounding_threshold") or 0.60) + + def predict(self, prompt: str, image_path: Optional[str] = None) -> OmniPredictionV42: + prompt_text = str(prompt or "").strip() + image_tensor, has_image = self._encode_image(image_path) + raw_outputs = self._run_prompt(prompt_text, image_tensor, has_image) + response_conf, domain_conf = self._confidence_from_outputs(raw_outputs) + draft_idx = int(raw_outputs["response"].argmax(dim=0).item()) if self.responses else 0 + draft_answer = self.responses[draft_idx] if self.responses else "" + planned_budget = _budget_hint_v42( + prompt_text, + response_confidence=response_conf, + domain_confidence=domain_conf, + ) + variants = _prompt_variants_v42( + prompt_text, + draft_answer, + response_confidence=response_conf, + domain_confidence=domain_conf, + ) + + weighted_outputs: List[tuple[float, Dict[str, torch.Tensor]]] = [] + combined = raw_outputs + for index, variant in enumerate(variants, start=1): + lowered = variant.lower() + weight = 1.20 if index == 1 else 1.0 + if "deep reasoning budget" in lowered: + weight = 1.10 + elif "medium reasoning budget" in lowered: + weight = 1.07 + elif "verifier pass" in lowered: + weight = 1.10 + elif "benchmark-style reasoning" in lowered: + weight = 1.09 + elif "agentic coding request" in lowered: + weight = 1.10 + elif "multimodal or omni-style request" in lowered: + weight = 1.08 + weighted_outputs.append((weight, self._run_prompt(variant, image_tensor, has_image))) + combined = self._combine_outputs(weighted_outputs) + response_conf, domain_conf = self._confidence_from_outputs(combined) + if index >= self.minimum_passes and response_conf >= self.grounding_threshold and domain_conf >= 0.50: + break + if index >= self.deliberation_passes: + break + + outputs = combined + intent_probs = torch.softmax(outputs["intent"], dim=0) + response_probs = torch.softmax(outputs["response"], dim=0) if self.responses else torch.ones(1) + vision_probs = torch.softmax(outputs["vision"], dim=0) + domain_probs = torch.softmax(outputs["domain"], dim=0) + + intent_idx = int(intent_probs.argmax(dim=0).item()) if intent_probs.numel() else 0 + response_idx = int(response_probs.argmax(dim=0).item()) if response_probs.numel() else 0 + domain_idx = int(domain_probs.argmax(dim=0).item()) if domain_probs.numel() else 0 + values, indices = torch.topk(vision_probs, k=min(3, len(SCIENCE_IMAGE_CLASSES))) + vision_key = SCIENCE_IMAGE_CLASSES[int(indices[0].item())] if SCIENCE_IMAGE_CLASSES else "" + top_vision = [ + {"concept": SCIENCE_IMAGE_CLASSES[int(index.item())], "confidence": round(float(value.item()), 4)} + for value, index in zip(values, indices) + if SCIENCE_IMAGE_CLASSES + ] + return OmniPredictionV42( + intent=self.intent_labels[min(intent_idx, len(self.intent_labels) - 1)] if self.intent_labels else "general", + response_text=self.responses[response_idx] if self.responses else "", + vision_concept=vision_key, + vision_confidence=float(values[0].item()) if values.numel() else 0.0, + top_vision=top_vision, + domain=self.domain_labels[min(domain_idx, len(self.domain_labels) - 1)] if self.domain_labels else "general", + domain_confidence=float(domain_probs[domain_idx].item()) if domain_probs.numel() else 0.0, + response_confidence=float(response_probs[response_idx].item()) if response_probs.numel() else 0.0, + reasoning_passes=len(weighted_outputs), + planned_budget=planned_budget, + ) + + +__all__ = [ + "OmniCollectiveEngineV42", + "OmniCollectiveNetV42", + "OmniPredictionV42", + "_prompt_variants_v42", +] diff --git a/source/prepare_omni_collective_v41.py b/source/prepare_omni_collective_v41.py new file mode 100644 index 000000000..6a6448225 --- /dev/null +++ b/source/prepare_omni_collective_v41.py @@ -0,0 +1,457 @@ +from __future__ import annotations + +import argparse +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + + +REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_OUTPUT_PATH = REPO_ROOT / "output" / "omni_collective_v41_blueprint.json" + + +RECENT_METHOD_REFERENCES: List[Dict[str, str]] = [ + { + "name": "DeepSeek-V3 Technical Report", + "url": "https://arxiv.org/abs/2412.19437", + "applied_as": "multi_token_prediction_auxiliary_head", + "why": "Add a compact next-span prediction loss so the model learns denser future-planning structure without a large parameter jump.", + }, + { + "name": "Better & Faster Large Language Models via Multi-token Prediction", + "url": "https://arxiv.org/abs/2404.19737", + "applied_as": "multi_token_prediction_training_recipe", + "why": "Use extra future-token heads during training to improve coding and algorithmic reasoning sample efficiency.", + }, + { + "name": "s1: Simple test-time scaling", + "url": "https://arxiv.org/abs/2501.19393", + "applied_as": "reasoning_budget_curriculum", + "why": "Train with short, medium, and long budget targets so v41 can spend more compute on hard prompts without bloating every answer.", + }, + { + "name": "LIMO: Less Is More for Reasoning", + "url": "https://arxiv.org/abs/2502.03387", + "applied_as": "high_density_reasoning_core_set", + "why": "Shift part of the reasoning mix away from sheer volume and toward a smaller, harder, cleaner curriculum.", + }, + { + "name": "RefineCoder", + "url": "https://arxiv.org/abs/2502.09183", + "applied_as": "code_critique_and_repair_loop", + "why": "Use critique-plus-repair traces for bug fixing, test repair, and patch refinement instead of one-shot code answers.", + }, + { + "name": "Reflexion: Language Agents with Verbal Reinforcement Learning", + "url": "https://arxiv.org/abs/2303.11366", + "applied_as": "verbal_reflection_and_repair_memory", + "why": "Store concise failure reflections for debugging and coding tasks instead of only training on final answers.", + }, + { + "name": "Self-Play Fine-Tuning (SPIN)", + "url": "https://arxiv.org/abs/2401.01335", + "applied_as": "iterative_self_play_preference_loop", + "why": "Generate self-play comparisons from v8, v40, and qwen teachers to improve helpfulness and problem-solving without requiring a giant new human-labeled set.", + }, + { + "name": "Self-Refine", + "url": "https://arxiv.org/abs/2303.17651", + "applied_as": "communication_polish_phase", + "why": "Teach v41 to critique and rewrite its own drafts into clearer human-facing answers.", + }, + { + "name": "Large Language Models have Intrinsic Self-Correction Ability", + "url": "https://arxiv.org/abs/2406.15673", + "applied_as": "fair_prompt_zero_temp_self_correction", + "why": "Use a dedicated self-correction pass with fair prompts and deterministic decoding to reduce hallucinations and repair weak drafts.", + }, + { + "name": "RLAIF vs. RLHF: Scaling Reinforcement Learning from Human Feedback with AI Feedback", + "url": "https://arxiv.org/abs/2309.00267", + "applied_as": "ai_feedback_preference_labels", + "why": "Use strong local teachers as AI preference labelers for communication, helpfulness, and harmlessness style comparisons.", + }, +] + + +def latest_v8_summary_path(repo_root: Path = REPO_ROOT) -> Path: + candidates = sorted( + repo_root.glob("output/supermix_omni_collective_v8_frontier_*/omni_collective_v8_frontier_summary.json"), + key=lambda item: (item.parent.name, item.stat().st_mtime), + reverse=True, + ) + if not candidates: + raise FileNotFoundError("No omni_collective_v8 frontier summary found under output/.") + return candidates[0].resolve() + + +def latest_common_benchmark_summary_path(repo_root: Path = REPO_ROOT) -> Optional[Path]: + candidates = sorted( + repo_root.glob("output/benchmark_all_models_common_plus_summary_*.json"), + key=lambda item: item.stat().st_mtime, + reverse=True, + ) + return candidates[0].resolve() if candidates else None + + +def _round_to_million(value: int) -> int: + return int(round(float(value) / 1_000_000.0) * 1_000_000) + + +def _safe_float(value: Any) -> float: + try: + return float(value) + except Exception: + return 0.0 + + +def _stage_metrics(summary: Dict[str, Any], stage_name: str) -> Dict[str, float]: + metrics = summary.get(stage_name, {}).get("val_metrics", {}) + if not isinstance(metrics, dict): + return {} + return {str(key): _safe_float(value) for key, value in metrics.items()} + + +def _recommended_parameter_target(current_parameter_count: int) -> int: + growth_floor = current_parameter_count + 8_000_000 + growth_soft = current_parameter_count + max(12_000_000, int(current_parameter_count * 0.125)) + growth_cap = current_parameter_count + 20_000_000 + cooked = min(max(growth_soft, growth_floor), growth_cap) + return _round_to_million(cooked) + + +def _extract_frontier_scores(repo_root: Path) -> Dict[str, Any]: + benchmark_path = latest_common_benchmark_summary_path(repo_root) + if benchmark_path is None: + return { + "benchmark_summary_path": None, + "best_overall": None, + "best_fusion": None, + } + payload = load_summary(benchmark_path) + rows = list(payload.get("summary_rows") or []) + if not rows: + return { + "benchmark_summary_path": str(benchmark_path), + "best_overall": None, + "best_fusion": None, + } + best_overall = max(rows, key=lambda item: float(item.get("overall_exact") or 0.0)) + fusion_rows = [row for row in rows if str(row.get("family") or "") == "fusion"] + best_fusion = max(fusion_rows, key=lambda item: float(item.get("overall_exact") or 0.0)) if fusion_rows else None + return { + "benchmark_summary_path": str(benchmark_path), + "best_overall": { + "model": str(best_overall.get("model") or ""), + "family": str(best_overall.get("family") or ""), + "overall_exact": _safe_float(best_overall.get("overall_exact") or 0.0), + }, + "best_fusion": ( + { + "model": str(best_fusion.get("model") or ""), + "family": str(best_fusion.get("family") or ""), + "overall_exact": _safe_float(best_fusion.get("overall_exact") or 0.0), + } + if best_fusion is not None + else None + ), + } + + +def build_v41_blueprint(summary: Dict[str, Any], summary_path: Optional[Path] = None) -> Dict[str, Any]: + parameter_count = int(summary.get("parameter_count") or 0) + stage1_rows = int(summary.get("dataset_summary", {}).get("stage1_rows") or 0) + stage2_rows = int(summary.get("dataset_summary", {}).get("stage2_rows") or 0) + stage2_metrics = _stage_metrics(summary, "stage2") + stage1_metrics = _stage_metrics(summary, "stage1") + response_accuracy = stage2_metrics.get("response_accuracy", 0.0) + intent_accuracy = stage2_metrics.get("intent_accuracy", 0.0) + vision_accuracy = stage2_metrics.get("vision_accuracy", 0.0) + domain_accuracy = stage2_metrics.get("domain_accuracy", 0.0) + teacher_keys = list(summary.get("dataset_summary", {}).get("teacher_league", {}).get("teacher_keys", [])) + frontier_scores = _extract_frontier_scores(REPO_ROOT) + + target_parameter_count = _recommended_parameter_target(parameter_count) + current_text_hidden = 272 + current_fusion_hidden = 1088 + current_expert_count = 10 + current_expert_hidden = 1792 + + return { + "family": "omni_collective_v41", + "prepared_at": datetime.now(timezone.utc).isoformat(), + "derived_from": { + "artifact": str(summary.get("artifact") or ""), + "summary_path": str(summary_path.resolve()) if summary_path else "", + "parameter_count": parameter_count, + "stage1_rows": stage1_rows, + "stage2_rows": stage2_rows, + "stage1_best_score": _safe_float(summary.get("stage1", {}).get("best_score")), + "stage2_best_score": _safe_float(summary.get("stage2", {}).get("best_score")), + }, + "observed_gaps": [ + { + "name": "response_quality", + "current": response_accuracy, + "target": max(0.145, round(response_accuracy + 0.05, 4)), + "why": "v8 still underperforms most on response selection and human-facing answer quality.", + }, + { + "name": "grounded_tool_and_model_choice", + "current": domain_accuracy, + "target": max(0.72, round(domain_accuracy + 0.04, 4)), + "why": "Sample outputs show routing failures on benchmark, materials, and grounded-fact prompts.", + }, + { + "name": "creative_communication", + "current": intent_accuracy, + "target": max(0.72, round(intent_accuracy + 0.03, 4)), + "why": "v41 should explain, compare, and reassure more naturally instead of defaulting to template fragments.", + }, + { + "name": "vision_and_specialist_fusion", + "current": vision_accuracy, + "target": max(0.73, round(vision_accuracy + 0.04, 4)), + "why": "v8 improved multimodal accuracy, but v41 still needs cleaner cross-domain fusion between text, vision, materials, protein, and 3D prompts.", + }, + ], + "teacher_strategy": { + "primary_teachers": [ + "omni_collective_v8", + "v40_benchmax", + "qwen_v28", + "qwen_v30", + "omni_collective_v7", + ], + "secondary_teachers": [ + key + for key in teacher_keys + if key not in {"v40_benchmax", "omni_collective_v7", "qwen_v28", "qwen_v30"} + ][:8], + "gemma4_slice": { + "status": "include_small_grounded_slice_if_local_weights_exist", + "target_rows": 24, + "note": "Keep Gemma 4 bounded and quality-gated rather than letting it dominate the teacher mix.", + }, + "disagreement_mining": { + "enabled": True, + "focus": [ + "benchmark_reasoning", + "coding_repairs", + "materials_grounding", + "human_communication", + ], + "note": "Save prompts where v8, v40, and qwen disagree, then build compare-and-justify rows instead of taking only the winner.", + }, + }, + "current_frontier": frontier_scores, + "architecture": { + "parameter_count_target": target_parameter_count, + "current_parameter_count": parameter_count, + "growth_vs_v8": target_parameter_count - parameter_count, + "text_hidden": {"from": current_text_hidden, "to": 288}, + "fusion_hidden": {"from": current_fusion_hidden, "to": 1152}, + "expert_count": {"from": current_expert_count, "to": 12}, + "expert_hidden": {"from": current_expert_hidden, "to": 1920}, + "expert_top_k": 2, + "deliberation_passes": {"from": 10, "to": 12}, + "min_deliberation_passes": 6, + "new_heads": [ + { + "name": "latent_plan_slots", + "purpose": "Distill short hidden plan states from explicit draft-plan traces, then keep the plans internal at inference time.", + }, + { + "name": "multi_token_prediction_aux", + "purpose": "Predict short future spans to increase information density and long-range coherence.", + }, + { + "name": "communication_polish_head", + "purpose": "Separate answer correctness from answer presentation so the model learns concise, human-friendly rewrites.", + }, + { + "name": "uncertainty_anchor_head", + "purpose": "Teach the model to declare uncertainty and missing evidence instead of inventing specifics.", + }, + ], + "novel_v41_ideas": [ + "Route-first then answer: supervise a hidden route/tool choice before the final answer for better Auto and collective behavior.", + "Disagreement harvest replay: convert teacher disagreement into compare-and-justify rows instead of only winner-take-all distillation.", + "Two-pass communication finish: solve first, then learn a short rewrite pass that preserves facts while improving tone and clarity.", + "Creativity rescue pairs: train on correct-but-flat vs vivid-and-still-grounded answer pairs so creativity does not come from hallucination.", + "Reflection memory shards: keep tiny verbal post-mortems for failed code, reasoning, and route decisions, then train on improved retries.", + "Reasoning-budget labels: supervise when the model should think short, medium, or long instead of only increasing deliberation globally.", + ], + }, + "data_program": [ + { + "name": "latent_plan_bootstrap", + "target_rows": 1800, + "source_style": "plan-first then answer rows with hidden plan targets", + "goal": "Teach v41 to infer structure and route choice before generating the final answer.", + }, + { + "name": "high_density_reasoning_core", + "target_rows": 3200, + "source_style": "LIMO-style compact hard reasoning set", + "goal": "Raise problem-solving quality without flooding the mix with repetitive low-signal reasoning rows.", + }, + { + "name": "code_critique_repair", + "target_rows": 2800, + "source_style": "traceback, failing-test, patch, and critique->repair examples", + "goal": "Improve coding, debugging, and repository-grounded problem solving.", + }, + { + "name": "reflection_memory_rows", + "target_rows": 1500, + "source_style": "bad attempt, short reflection, repaired answer", + "goal": "Turn failures into compact reusable repair knowledge instead of only supervising the winning answer.", + }, + { + "name": "human_communication_polish", + "target_rows": 6200, + "source_style": "rewrite, summarize, compare, explain, reassure, and uncertainty-calibrated replies", + "goal": "Make v41 more useful and easier to talk to without losing directness.", + }, + { + "name": "teacher_disagreement_hardcases", + "target_rows": 1800, + "source_style": "v8/v40/qwen disagreement mining", + "goal": "Expose the model to exactly the prompts where current strong local teachers diverge.", + }, + { + "name": "grounded_uncertainty_rows", + "target_rows": 900, + "source_style": "current-fact uncertainty, partial evidence, and missing-tool disclaimers", + "goal": "Reduce hallucinated specifics and improve transparent communication.", + }, + { + "name": "creative_constraint_pairs", + "target_rows": 2200, + "source_style": "creative answer pairs scored for vividness, usefulness, and factual discipline", + "goal": "Boost creativity without dropping answer discipline.", + }, + { + "name": "route_then_answer_supervision", + "target_rows": 1200, + "source_style": "model selection, tool choice, specialist handoff, then final answer", + "goal": "Make Auto and collective mode smarter because the base model reasons about routing explicitly.", + }, + { + "name": "reasoning_budget_curriculum", + "target_rows": 1400, + "source_style": "same task solved under short, medium, and deep reasoning budgets", + "goal": "Teach the model when extra deliberation is worthwhile instead of overthinking every prompt.", + }, + { + "name": "specialist_dense_topups", + "target_rows": 1600, + "source_style": "materials, protein, 3D, native image, and science vision deltas", + "goal": "Keep the specialist fusion edge while shifting more capacity toward coding and human communication.", + }, + ], + "training_program": [ + { + "stage": "stage0_latent_plan_bootstrap", + "recipe": "Distill short explicit plans, route choices, and critique stubs into latent slots before normal SFT.", + "inspiration": "DeepSeek-V3 + quiet planning style", + }, + { + "stage": "stage1_teacher_sft", + "recipe": "Teacher-heavy mixed SFT with harder coding/problem-solving weighting and a smaller, denser reasoning core.", + "inspiration": "LIMO-style compact reasoning curriculum", + }, + { + "stage": "stage2_self_play_preference", + "recipe": "Generate self-play pairs from v8/v40/qwen-style answers, then train chosen-vs-rejected preference objectives.", + "inspiration": "SPIN", + }, + { + "stage": "stage3_code_refine", + "recipe": "Run critique->repair loops on code, tests, and debugging prompts and keep only verifiably improved completions.", + "inspiration": "RefineCoder, Reflexion, and code self-reflection", + }, + { + "stage": "stage4_communication_polish", + "recipe": "Short final polish phase on human communication, grounded uncertainty, and compare/explain prompts using self-refine and AI feedback labels.", + "inspiration": "Self-Refine + RLAIF", + }, + { + "stage": "stage5_self_correction_finish", + "recipe": "Run deterministic self-correction on weak drafts with fair prompts and keep only improved final answers.", + "inspiration": "Intrinsic Self-Correction", + }, + ], + "runtime_recommendations": { + "device": "prefer_cuda_or_dml_never_cpu_if_avoidable", + "amp": "auto", + "amp_dtype": "bfloat16", + "compile_model": True, + "compile_mode": "reduce-overhead", + "grad_accum_steps": 4, + "ema_decay": 0.9993, + "warmup_ratio": 0.06, + "min_lr_scale": 0.03, + "checkpoint_every_optimizer_steps": 100, + "keep_last_stage_checkpoints": 3, + "save_stage_complete_weights": True, + }, + "success_gates": { + "response_accuracy_min": max(0.145, round(response_accuracy + 0.05, 4)), + "stage2_score_min": max(0.455, round(_safe_float(summary.get("stage2", {}).get("best_score")) + 0.03, 4)), + "common_benchmark_exact_min": 0.25, + "must_beat": ["omni_collective_v8_preview", "omni_collective_v7"], + "stretch_goal": "challenge_v40_benchmax_on_benchmark_reasoning_without losing multimodal utility", + }, + "promotion_protocol": { + "required_preview_benchmarks": 3, + "preview_selection_rule": "Promote only the best preview checkpoint by weighted score across benchmark exact, coding repair, communication, and grounded uncertainty packs.", + "must_not_regress": [ + "stage2 vision_accuracy vs v8 by more than 0.02", + "stage2 domain_accuracy vs v8 by more than 0.02", + ], + "must_beat_frontier_if_available": frontier_scores.get("best_fusion"), + "human_style_gate": { + "required": True, + "description": "v41 must beat v8 on the communication-polish eval pack before promotion, even if the raw benchmark score is close.", + }, + }, + "recent_method_references": RECENT_METHOD_REFERENCES, + "implementation_checklist": [ + "Freeze the latest v8 summary, benchmark outputs, and hard-failure prompts as v41 design inputs.", + "Create new dataset builders for disagreement mining, route-then-answer rows, communication polish pairs, latent-plan rows, and code-repair rows.", + "Add latent plan slots, multi-token prediction auxiliary loss, and uncertainty-anchor head to the v41 model.", + "Reuse the improved v8 checkpoint/resume path from the start so long v41 runs survive reboots and disk pressure.", + "Benchmark preview checkpoints during stage2 rather than waiting until the end of the full run.", + "Run the new communication, coding-repair, and grounded-uncertainty eval packs before any promotion decision.", + ], + } + + +def load_summary(path: Path) -> Dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare a machine-readable blueprint for omni_collective_v41.") + parser.add_argument("--summary", default="", help="Optional path to a v8 frontier summary JSON. Defaults to the latest local v8 summary.") + parser.add_argument("--output", default=str(DEFAULT_OUTPUT_PATH), help="Where to write the v41 blueprint JSON.") + parser.add_argument("--stdout", action="store_true", help="Print the blueprint to stdout as well.") + args = parser.parse_args() + + summary_path = Path(args.summary).resolve() if str(args.summary).strip() else latest_v8_summary_path(REPO_ROOT) + output_path = Path(args.output).resolve() + + summary = load_summary(summary_path) + blueprint = build_v41_blueprint(summary, summary_path=summary_path) + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(blueprint, indent=2, ensure_ascii=True), encoding="utf-8") + if args.stdout: + print(json.dumps(blueprint, indent=2, ensure_ascii=True)) + + +if __name__ == "__main__": + main() diff --git a/source/prepare_omni_collective_v42.py b/source/prepare_omni_collective_v42.py new file mode 100644 index 000000000..20cd1366c --- /dev/null +++ b/source/prepare_omni_collective_v42.py @@ -0,0 +1,513 @@ +from __future__ import annotations + +import argparse +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +try: + from prepare_omni_collective_v41 import latest_common_benchmark_summary_path, load_summary +except ImportError: # pragma: no cover + from .prepare_omni_collective_v41 import latest_common_benchmark_summary_path, load_summary + + +REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_OUTPUT_PATH = REPO_ROOT / "output" / "omni_collective_v42_blueprint.json" +V40_MANIFEST_PATH = Path(__file__).resolve().parent / "v40_benchmax_manifest.json" + + +RECENT_METHOD_REFERENCES_V42: List[Dict[str, str]] = [ + { + "name": "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate", + "url": "https://arxiv.org/abs/2504.19874", + "applied_as": "teacher_generation_kv_cache_compression", + "why": "Treat TurboQuant as a runtime enabler for longer teacher traces and cheaper verifier passes, not as a weight-training substitute.", + }, + { + "name": "Small Language Models Need Strong Verifiers to Self-Correct Reasoning", + "url": "https://arxiv.org/abs/2404.17140", + "applied_as": "verifier_then_repair_rows", + "why": "Add explicit critique and repair supervision instead of assuming the student will self-correct from final answers alone.", + }, + { + "name": "Self-Enhanced Reasoning Training", + "url": "https://arxiv.org/abs/2502.12744", + "applied_as": "latent_reasoning_bootstrap", + "why": "Use self-generated reasoning paths selectively to activate latent reasoning without requiring giant chain-of-thought corpora everywhere.", + }, + { + "name": "Lost at the Beginning of Reasoning", + "url": "https://arxiv.org/abs/2506.22058", + "applied_as": "opening_step_quality_filtering", + "why": "Prefer high-quality first steps and route choices so the model does not waste budget on weak starts.", + }, + { + "name": "Learning from Partial Chain-of-Thought via Truncated-Reasoning Self-Distillation", + "url": "https://arxiv.org/abs/2603.13274", + "applied_as": "partial_reasoning_distillation", + "why": "Distill useful partial traces so v42 can stop earlier on easy problems without losing answer quality.", + }, + { + "name": "google/gemma-4-31B-it", + "url": "https://huggingface.co/google/gemma-4-31B-it", + "applied_as": "grounded_multimodal_teacher", + "why": "Use Gemma 4 as a bounded teacher for grounded multimodal answers, image reasoning, and careful presentation.", + }, + { + "name": "Qwen/Qwen3.5-397B-A17B", + "url": "https://huggingface.co/Qwen/Qwen3.5-397B-A17B", + "applied_as": "long_context_tool_use_teacher", + "why": "Use Qwen 3.5 for long-context, tool-use, and reasoning traces, especially where v41 is weak against v40_benchmax.", + }, + { + "name": "Qwen/Qwen3-Coder-Next", + "url": "https://huggingface.co/Qwen/Qwen3-Coder-Next", + "applied_as": "agentic_coding_teacher", + "why": "Use Coder-Next for code-edit, test-repair, and long-horizon debugging traces instead of only one-shot code answers.", + }, + { + "name": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "url": "https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct", + "applied_as": "omni_audio_video_teacher", + "why": "Use the omni model for audio-video alignment, captioning, OCR-style prompts, and broader multimodal route supervision.", + }, +] + + +def latest_v41_summary_path(repo_root: Path = REPO_ROOT) -> Path: + candidates = sorted( + repo_root.glob("output/supermix_omni_collective_v41_frontier_*/omni_collective_v41_frontier_summary.json"), + key=lambda item: (item.parent.name, item.stat().st_mtime), + reverse=True, + ) + if not candidates: + raise FileNotFoundError("No omni_collective_v41 frontier summary found under output/.") + return candidates[0].resolve() + + +def latest_v41_zip_path(repo_root: Path = REPO_ROOT) -> Path: + candidates = sorted( + repo_root.glob("output/supermix_omni_collective_v41_frontier_*.zip"), + key=lambda item: (item.name, item.stat().st_mtime), + reverse=True, + ) + if not candidates: + raise FileNotFoundError("No omni_collective_v41 frontier zip found under output/.") + return candidates[0].resolve() + + +def _safe_float(value: Any) -> float: + try: + return float(value) + except Exception: + return 0.0 + + +def _stage_metrics(summary: Dict[str, Any], stage_name: str) -> Dict[str, float]: + metrics = summary.get(stage_name, {}).get("val_metrics", {}) + if not isinstance(metrics, dict): + return {} + return {str(key): _safe_float(value) for key, value in metrics.items()} + + +def _round_to_million(value: int) -> int: + return int(round(float(value) / 1_000_000.0) * 1_000_000) + + +def _recommended_parameter_target(current_parameter_count: int) -> int: + growth_floor = current_parameter_count + 10_000_000 + growth_soft = current_parameter_count + max(16_000_000, int(current_parameter_count * 0.14)) + growth_cap = current_parameter_count + 24_000_000 + return _round_to_million(min(max(growth_soft, growth_floor), growth_cap)) + + +def _extract_benchmark_context(repo_root: Path) -> Dict[str, Any]: + benchmark_path = latest_common_benchmark_summary_path(repo_root) + if benchmark_path is None: + return { + "benchmark_summary_path": None, + "v40_benchmax": None, + "omni_collective_v41": None, + "best_overall": None, + } + payload = load_summary(benchmark_path) + rows = list(payload.get("summary_rows") or []) + indexed = {str(row.get("model") or ""): row for row in rows} + best_overall = max(rows, key=lambda item: float(item.get("overall_exact") or 0.0)) if rows else None + + def _shape(model_key: str) -> Optional[Dict[str, Any]]: + row = indexed.get(model_key) + if row is None: + return None + return { + "model": model_key, + "family": str(row.get("family") or ""), + "overall_exact": _safe_float(row.get("overall_exact") or 0.0), + } + + return { + "benchmark_summary_path": str(benchmark_path.resolve()), + "v40_benchmax": _shape("v40_benchmax"), + "omni_collective_v41": _shape("omni_collective_v41"), + "best_overall": ( + { + "model": str(best_overall.get("model") or ""), + "family": str(best_overall.get("family") or ""), + "overall_exact": _safe_float(best_overall.get("overall_exact") or 0.0), + } + if best_overall is not None + else None + ), + } + + +def _load_v40_manifest(path: Path = V40_MANIFEST_PATH) -> Dict[str, Any]: + if not path.exists(): + return {} + return json.loads(path.read_text(encoding="utf-8")) + + +def build_v42_blueprint(summary: Dict[str, Any], summary_path: Optional[Path] = None) -> Dict[str, Any]: + parameter_count = int(summary.get("parameter_count") or 0) + stage1_rows = int(summary.get("dataset_summary", {}).get("stage1_rows") or 0) + stage2_rows = int(summary.get("dataset_summary", {}).get("stage2_rows") or 0) + stage1_metrics = _stage_metrics(summary, "stage1") + stage2_metrics = _stage_metrics(summary, "stage2") + response_accuracy = stage2_metrics.get("response_accuracy", 0.0) + intent_accuracy = stage2_metrics.get("intent_accuracy", 0.0) + vision_accuracy = stage2_metrics.get("vision_accuracy", 0.0) + domain_accuracy = stage2_metrics.get("domain_accuracy", 0.0) + current_benchmark = _extract_benchmark_context(REPO_ROOT) + v40_manifest = _load_v40_manifest() + v40_exact = _safe_float((current_benchmark.get("v40_benchmax") or {}).get("overall_exact")) + v41_exact = _safe_float((current_benchmark.get("omni_collective_v41") or {}).get("overall_exact")) + benchmark_gap = round(max(0.0, v40_exact - v41_exact), 4) + target_parameter_count = _recommended_parameter_target(parameter_count) + + return { + "family": "omni_collective_v42", + "prepared_at": datetime.now(timezone.utc).isoformat(), + "as_of_date": "2026-04-10", + "derived_from": { + "artifact": str(summary.get("artifact") or ""), + "summary_path": str(summary_path.resolve()) if summary_path else "", + "parameter_count": parameter_count, + "stage1_rows": stage1_rows, + "stage2_rows": stage2_rows, + "stage1_best_score": _safe_float(summary.get("stage1", {}).get("best_score")), + "stage2_best_score": _safe_float(summary.get("stage2", {}).get("best_score")), + }, + "baseline_findings": { + "common_benchmark_exact": { + "v40_benchmax": v40_exact, + "omni_collective_v41": v41_exact, + "gap_to_close": benchmark_gap, + }, + "v41_strengths_to_keep": [ + "route-then-answer behavior", + "communication polish", + "coding repair supervision", + "grounded uncertainty responses", + ], + "v40_strengths_to_import": [ + "failure-cluster curriculum", + "benchmark hard-example replay", + "budget-aware verifier prompts", + "consensus-teacher sampling", + ], + }, + "observed_gaps": [ + { + "name": "benchmark_reasoning", + "current": v41_exact, + "target": max(0.25, round(v41_exact + 0.08, 4)), + "why": "v42 needs to recover the benchmark-specialist edge from v40_benchmax while keeping v41's broader multimodal utility.", + }, + { + "name": "response_quality", + "current": response_accuracy, + "target": max(0.16, round(response_accuracy + 0.05, 4)), + "why": "v41 improved presentation but still underperforms on response selection and final answer quality.", + }, + { + "name": "agentic_coding_and_verification", + "current": domain_accuracy, + "target": max(0.73, round(domain_accuracy + 0.04, 4)), + "why": "v42 should route to better verifier-first coding behavior and retain stronger repair traces from Qwen3-Coder-Next style supervision.", + }, + { + "name": "multimodal_grounding", + "current": vision_accuracy, + "target": max(0.48, round(vision_accuracy + 0.08, 4)), + "why": "Gemma 4 and Qwen3-Omni expose better grounded multimodal behaviors than the current v41 prep mix.", + }, + { + "name": "diversity_and_global_coverage", + "current": intent_accuracy, + "target": max(0.64, round(intent_accuracy + 0.05, 4)), + "why": "Qwen3.5 explicitly broadens long-context, multilingual, and tool-use coverage that the current v41 dataset underrepresents.", + }, + ], + "teacher_strategy": { + "primary_teachers": [ + "omni_collective_v41", + "v40_benchmax", + "google/gemma-4-31B-it", + "Qwen/Qwen3.5-397B-A17B", + "Qwen/Qwen3-Coder-Next", + "Qwen/Qwen3-Omni-30B-A3B-Instruct", + ], + "bounded_teachers": [ + { + "model": "google/gemma-4-E4B-it", + "role": "small grounded multimodal comparisons", + "row_budget": 160, + }, + { + "model": "google/gemma-4-26B-A4B-it", + "role": "mid-sized multimodal and instruction style comparisons", + "row_budget": 220, + }, + ], + "teacher_roles": [ + { + "model": "v40_benchmax", + "role": "benchmark hard-case teacher", + "why": "Use it for benchmark replay, verifier rows, and failure-cluster curriculum.", + }, + { + "model": "google/gemma-4-31B-it", + "role": "grounded multimodal teacher", + "why": "Use it for image-grounded comparisons, safer multimodal phrasing, and controlled explanation style.", + }, + { + "model": "Qwen/Qwen3.5-397B-A17B", + "role": "long-context reasoning and tool-use teacher", + "why": "Use it for long-context summaries, tool traces, reasoning budgets, and route-then-answer rows.", + }, + { + "model": "Qwen/Qwen3-Coder-Next", + "role": "agentic coding teacher", + "why": "Use it for code-edit planning, test repair, and repository-grounded debugging traces.", + }, + { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "role": "audio-video multimodal teacher", + "why": "Use it for caption, OCR-style, video, and cross-modal route supervision.", + }, + ], + "teacher_mixture_rules": [ + "Never let any single hosted teacher dominate the distilled mix.", + "Keep benchmark replay anchored to v40_benchmax plus local verifier rows.", + "Keep Gemma 4 bounded to grounded multimodal slices instead of generic chat domination.", + "Use Qwen3-Coder-Next only on coding-heavy prompts so the rest of the mix stays broad.", + ], + "turboquant_runtime_note": { + "enabled": True, + "purpose": "Use TurboQuant-style KV compression or equivalent when generating long teacher traces so the data program can afford deeper verifier passes.", + "training_implication": "Supervise budget choice and evidence compression, not the quantization algorithm itself.", + }, + }, + "architecture": { + "parameter_count_target": target_parameter_count, + "current_parameter_count": parameter_count, + "growth_vs_v41": target_parameter_count - parameter_count, + "text_hidden": {"from": 288, "to": 304}, + "fusion_hidden": {"from": 1152, "to": 1216}, + "expert_count": {"from": 12, "to": 14}, + "expert_hidden": {"from": 1920, "to": 2048}, + "expert_top_k": 2, + "deliberation_passes": {"from": 12, "to": 14}, + "min_deliberation_passes": 7, + "new_heads": [ + { + "name": "budget_router_head", + "purpose": "Predict whether a prompt needs short, medium, or deep reasoning before the model spends extra compute.", + }, + { + "name": "teacher_consensus_head", + "purpose": "Separate answer confidence from agreement confidence so the model can learn when to distrust a single teacher.", + }, + { + "name": "cache_budget_head", + "purpose": "Teach the model how much evidence to keep active and when compressed context is enough.", + }, + { + "name": "verifier_gate_head", + "purpose": "Learn when a second-pass verifier or repair step is worth the latency on coding and reasoning prompts.", + }, + ], + "novel_v42_ideas": [ + "Benchmark bridge curriculum: replay v40-style hard clusters directly inside the broader v41 omni scaffold.", + "Teacher role specialization: use Gemma 4, Qwen 3.5, Qwen3-Coder-Next, and Qwen3-Omni for sharply different data slices instead of generic all-purpose distillation.", + "TurboQuant-inspired context economy: train explicit evidence-budget labels so long-context behavior improves without inflating every answer.", + "Partial reasoning distillation: keep good early and mid reasoning segments, not only full-length traces.", + "Verifier-first coding: train code repair rows where a verifier note or failing test arrives before the final patch.", + "Diversity expansion by scenario, modality, and user goal rather than only adding more generic chat volume.", + ], + }, + "data_program": [ + { + "name": "v40_benchmark_bridge", + "target_rows": 4200, + "source_style": "v40 hard-example replay, failure clusters, and verifier-guided benchmark repair", + "goal": "Close the measured v40-v41 benchmark gap without collapsing the omni model into a narrow benchmark specialist.", + }, + { + "name": "gemma4_grounded_multimodal", + "target_rows": 2600, + "source_style": "bounded Gemma 4 image-text and grounded explanation slices", + "goal": "Lift multimodal grounding and safer, clearer explanation quality.", + }, + { + "name": "qwen35_long_context_and_tool_use", + "target_rows": 3800, + "source_style": "Qwen3.5 long-context, tool-use, and reasoning-budget traces", + "goal": "Improve route choice, long-context stability, and tool-aware reasoning.", + }, + { + "name": "qwen3_coder_next_agentic_repairs", + "target_rows": 3400, + "source_style": "agentic coding, repository-grounded edits, failing tests, and patch repair traces", + "goal": "Make coding behavior more deliberate and more verifiable than v41's mostly answer-level repair mix.", + }, + { + "name": "qwen3_omni_audio_video_alignment", + "target_rows": 1800, + "source_style": "audio, video, OCR-style, and cross-modal route supervision", + "goal": "Improve multimodal specialist routing and richer omni prompt understanding.", + }, + { + "name": "verifier_guided_self_correction", + "target_rows": 2600, + "source_style": "draft, verifier note, repaired answer", + "goal": "Teach v42 to repair weak drafts with targeted feedback rather than only rerolling responses.", + }, + { + "name": "partial_reasoning_distill", + "target_rows": 2100, + "source_style": "high-quality partial chains, truncated reasoning, and early-stop successful traces", + "goal": "Keep strong reasoning while reducing wasteful overthinking on easy or medium tasks.", + }, + { + "name": "global_diversity_and_style", + "target_rows": 4800, + "source_style": "broader users, locales, tones, professions, and request styles", + "goal": "Make the dataset larger and more diverse in a meaningful way instead of only duplicating similar prompts.", + }, + { + "name": "cache_budget_and_context_selection", + "target_rows": 1600, + "source_style": "compressed evidence selection, context budget labels, and when-to-stop traces", + "goal": "Translate TurboQuant-style efficiency lessons into trainable budget and evidence behavior.", + }, + { + "name": "specialist_dense_topups_v42", + "target_rows": 2400, + "source_style": "materials, protein, 3D, science vision, and model-selection top-ups", + "goal": "Preserve the specialist fusion edge while expanding coding, benchmark, and multimodal coverage.", + }, + ], + "training_program": [ + { + "stage": "stage0_route_budget_bootstrap", + "recipe": "Train route choice, verifier choice, and short-medium-deep budget prediction before heavier SFT.", + "inspiration": "reasoning budget control plus route-first supervision", + }, + { + "stage": "stage1_teacher_sft", + "recipe": "Mix v41 carryover rows with v40 hard examples and the new Gemma/Qwen role-specific teacher slices.", + "inspiration": "teacher specialization instead of flat all-model distillation", + }, + { + "stage": "stage2_verifier_and_repair", + "recipe": "Keep only repairs that beat the draft under verifier heuristics or benchmark-style exact checks.", + "inspiration": "strong verifier self-correction", + }, + { + "stage": "stage3_partial_reasoning_distill", + "recipe": "Distill good partial traces and explicit early-stop examples so v42 learns when enough reasoning is enough.", + "inspiration": "TRSD and early-step quality filtering", + }, + { + "stage": "stage4_agentic_coding_finish", + "recipe": "Run code-edit, test-repair, and repo-grounded critique->patch loops with Coder-Next style traces.", + "inspiration": "agentic coding teachers and repair-first data", + }, + { + "stage": "stage5_multimodal_omni_finish", + "recipe": "Short final phase on Gemma 4 and Qwen3-Omni slices for grounded image/audio/video route supervision.", + "inspiration": "teacher-role specialization", + }, + ], + "runtime_recommendations": { + "device": "prefer_cuda_or_dml_never_cpu_if_avoidable", + "amp": "auto", + "amp_dtype": "bfloat16", + "compile_model": True, + "compile_mode": "reduce-overhead", + "grad_accum_steps": 6, + "ema_decay": 0.9994, + "warmup_ratio": 0.07, + "min_lr_scale": 0.025, + "checkpoint_every_optimizer_steps": 100, + "keep_last_stage_checkpoints": 4, + "save_stage_complete_weights": True, + "teacher_generation_acceleration": "TurboQuant-style KV compression or equivalent cache reduction for long teacher passes", + }, + "success_gates": { + "response_accuracy_min": max(0.16, round(response_accuracy + 0.05, 4)), + "stage2_score_min": max(0.39, round(_safe_float(summary.get("stage2", {}).get("best_score")) + 0.03, 4)), + "common_benchmark_exact_min": max(0.245, round(v40_exact + 0.002, 4)), + "must_beat": ["omni_collective_v41"], + "stretch_goal": "beat_v40_benchmax_on_common_benchmark_exact_while_preserving_omni_utility", + }, + "promotion_protocol": { + "required_preview_benchmarks": 4, + "preview_selection_rule": "Promote only checkpoints that beat v41 overall and stay within multimodal regression limits while closing the v40 benchmark gap.", + "must_not_regress": [ + "stage2 vision_accuracy vs v41 by more than 0.02", + "stage2 domain_accuracy vs v41 by more than 0.02", + "coding repair pack vs v41 by more than 0.01", + ], + "hard_gate": "No promotion if common benchmark exact remains below v40_benchmax.", + }, + "recent_method_references": RECENT_METHOD_REFERENCES_V42, + "v40_manifest_snapshot": { + "objective": str(v40_manifest.get("objective") or ""), + "innovation_flags": dict(v40_manifest.get("innovation_flags") or {}), + "novel_features": list(v40_manifest.get("novel_features") or []), + }, + "current_benchmark_context": current_benchmark, + "implementation_checklist": [ + "Freeze the latest v41 frontier summary and v40 benchmark summary as the v42 design inputs.", + "Add new v42 dataset builders for benchmark bridge, teacher-role routing, verifier repair, TurboQuant-style budget rows, and diversity expansion.", + "Keep the v41 data pipeline as the stable base, then add v42 row groups on top instead of rewriting the whole training loop.", + "Use Gemma 4 and Qwen 3.5 family models only in bounded, role-specific slices.", + "Evaluate preview checkpoints on the common benchmark plus coding-repair and multimodal route packs before promotion.", + ], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare a machine-readable blueprint for omni_collective_v42.") + parser.add_argument("--summary", default="", help="Optional path to a v41 frontier summary JSON. Defaults to the latest local v41 summary.") + parser.add_argument("--output", default=str(DEFAULT_OUTPUT_PATH), help="Where to write the v42 blueprint JSON.") + parser.add_argument("--stdout", action="store_true", help="Print the blueprint to stdout as well.") + args = parser.parse_args() + + summary_path = Path(args.summary).resolve() if str(args.summary).strip() else latest_v41_summary_path(REPO_ROOT) + output_path = Path(args.output).resolve() + + summary = load_summary(summary_path) + blueprint = build_v42_blueprint(summary, summary_path=summary_path) + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(blueprint, indent=2, ensure_ascii=True), encoding="utf-8") + if args.stdout: + print(json.dumps(blueprint, indent=2, ensure_ascii=True)) + + +if __name__ == "__main__": + main() diff --git a/source/run_omni_collective_v41_background.py b/source/run_omni_collective_v41_background.py new file mode 100644 index 000000000..2bcde9e41 --- /dev/null +++ b/source/run_omni_collective_v41_background.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import argparse +import io +import os +import sys +import traceback +from pathlib import Path +from typing import Iterable, TextIO + +import train_omni_collective_v41 + + +class TeeStream(io.TextIOBase): + def __init__(self, streams: Iterable[TextIO]) -> None: + super().__init__() + self._streams = [stream for stream in streams if stream is not None] + + def write(self, s: str) -> int: # pragma: no cover + text = str(s) + for stream in self._streams: + stream.write(text) + stream.flush() + return len(text) + + def flush(self) -> None: # pragma: no cover + for stream in self._streams: + stream.flush() + + def isatty(self) -> bool: # pragma: no cover + return False + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Background launcher for omni_collective_v41 with in-process log teeing.") + parser.add_argument("--out-log", required=True) + parser.add_argument("--err-log", required=True) + parser.add_argument("--worker-pid-file", required=True) + parser.add_argument("train_args", nargs=argparse.REMAINDER) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + out_path = Path(args.out_log).resolve() + err_path = Path(args.err_log).resolve() + pid_path = Path(args.worker_pid_file).resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + err_path.parent.mkdir(parents=True, exist_ok=True) + pid_path.parent.mkdir(parents=True, exist_ok=True) + + out_handle = out_path.open("a", encoding="utf-8", buffering=1, newline="\n") + err_handle = err_path.open("a", encoding="utf-8", buffering=1, newline="\n") + pid_path.write_text(str(os.getpid()), encoding="ascii") + + original_stdout = sys.stdout + original_stderr = sys.stderr + sys.stdout = TeeStream([original_stdout, out_handle]) + sys.stderr = TeeStream([original_stderr, err_handle]) + try: + train_args = list(args.train_args) + if train_args and train_args[0] == "--": + train_args = train_args[1:] + sys.argv = ["train_omni_collective_v41.py"] + train_args + train_omni_collective_v41.main() + except Exception: + traceback.print_exc() + raise + finally: + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + sys.stdout = original_stdout + sys.stderr = original_stderr + out_handle.close() + err_handle.close() + + +if __name__ == "__main__": + main() diff --git a/source/run_train_omni_collective_v41_local.ps1 b/source/run_train_omni_collective_v41_local.ps1 new file mode 100644 index 000000000..4257d1bcb --- /dev/null +++ b/source/run_train_omni_collective_v41_local.ps1 @@ -0,0 +1,59 @@ +Set-StrictMode -Version Latest +$ErrorActionPreference = "Stop" + +$repoRoot = Split-Path -Parent $PSScriptRoot +Set-Location $repoRoot + +$logDir = Join-Path $repoRoot "output" +New-Item -ItemType Directory -Force -Path $logDir | Out-Null + +$timestamp = Get-Date -Format "yyyyMMdd_HHmmss" +$outLog = Join-Path $logDir "omni_collective_v41_train_${timestamp}.out.log" +$errLog = Join-Path $logDir "omni_collective_v41_train_${timestamp}.err.log" +$pidPath = Join-Path $logDir "omni_collective_v41_train.pid" +$workerPidPath = Join-Path $logDir "omni_collective_v41_train.worker.pid" + +if (Test-Path $pidPath) { + Remove-Item -LiteralPath $pidPath -Force -ErrorAction SilentlyContinue +} +if (Test-Path $workerPidPath) { + Remove-Item -LiteralPath $workerPidPath -Force -ErrorAction SilentlyContinue +} + +$python = Join-Path $repoRoot ".venv-dml\Scripts\python.exe" +if (-not (Test-Path $python)) { + throw "Python executable not found at $python" +} + +$pythonArgs = @( + "-u" + "source\train_omni_collective_v41.py" + "--train_frontier" +) + $args + +$proc = Start-Process -FilePath $python -ArgumentList $pythonArgs -WorkingDirectory $repoRoot -RedirectStandardOutput $outLog -RedirectStandardError $errLog -PassThru +$launcherPid = $proc.Id +$workerPid = $launcherPid + +for ($attempt = 0; $attempt -lt 90; $attempt++) { + Start-Sleep -Seconds 1 + $child = Get-CimInstance Win32_Process -Filter "ParentProcessId = $launcherPid" -ErrorAction SilentlyContinue | + Where-Object { $_.Name -like 'python*.exe' } | + Select-Object -First 1 + if ($child -and $child.ProcessId) { + $workerPid = [int]$child.ProcessId + break + } +} + +Set-Content -Path $pidPath -Value $launcherPid -Encoding ascii +Set-Content -Path $workerPidPath -Value $workerPid -Encoding ascii + +[pscustomobject]@{ + launcher_pid = $launcherPid + worker_pid = $workerPid + out_log = $outLog + err_log = $errLog + pid_file = $pidPath + worker_pid_file = $workerPidPath +} | ConvertTo-Json -Depth 4 diff --git a/source/run_train_omni_collective_v42_local.ps1 b/source/run_train_omni_collective_v42_local.ps1 new file mode 100644 index 000000000..34c2b0418 --- /dev/null +++ b/source/run_train_omni_collective_v42_local.ps1 @@ -0,0 +1,59 @@ +Set-StrictMode -Version Latest +$ErrorActionPreference = "Stop" + +$repoRoot = Split-Path -Parent $PSScriptRoot +Set-Location $repoRoot + +$logDir = Join-Path $repoRoot "output" +New-Item -ItemType Directory -Force -Path $logDir | Out-Null + +$timestamp = Get-Date -Format "yyyyMMdd_HHmmss" +$outLog = Join-Path $logDir "omni_collective_v42_train_${timestamp}.out.log" +$errLog = Join-Path $logDir "omni_collective_v42_train_${timestamp}.err.log" +$pidPath = Join-Path $logDir "omni_collective_v42_train.pid" +$workerPidPath = Join-Path $logDir "omni_collective_v42_train.worker.pid" + +if (Test-Path $pidPath) { + Remove-Item -LiteralPath $pidPath -Force -ErrorAction SilentlyContinue +} +if (Test-Path $workerPidPath) { + Remove-Item -LiteralPath $workerPidPath -Force -ErrorAction SilentlyContinue +} + +$python = Join-Path $repoRoot ".venv-dml\Scripts\python.exe" +if (-not (Test-Path $python)) { + throw "Python executable not found at $python" +} + +$pythonArgs = @( + "-u" + "source\train_omni_collective_v42.py" + "--train_frontier" +) + $args + +$proc = Start-Process -FilePath $python -ArgumentList $pythonArgs -WorkingDirectory $repoRoot -RedirectStandardOutput $outLog -RedirectStandardError $errLog -PassThru +$launcherPid = $proc.Id +$workerPid = $launcherPid + +for ($attempt = 0; $attempt -lt 90; $attempt++) { + Start-Sleep -Seconds 1 + $child = Get-CimInstance Win32_Process -Filter "ParentProcessId = $launcherPid" -ErrorAction SilentlyContinue | + Where-Object { $_.Name -like 'python*.exe' } | + Select-Object -First 1 + if ($child -and $child.ProcessId) { + $workerPid = [int]$child.ProcessId + break + } +} + +Set-Content -Path $pidPath -Value $launcherPid -Encoding ascii +Set-Content -Path $workerPidPath -Value $workerPid -Encoding ascii + +[pscustomobject]@{ + launcher_pid = $launcherPid + worker_pid = $workerPid + out_log = $outLog + err_log = $errLog + pid_file = $pidPath + worker_pid_file = $workerPidPath +} | ConvertTo-Json -Depth 4 diff --git a/source/train_omni_collective_v41.py b/source/train_omni_collective_v41.py new file mode 100644 index 000000000..491bb7703 --- /dev/null +++ b/source/train_omni_collective_v41.py @@ -0,0 +1,1487 @@ +from __future__ import annotations + +import argparse +import json +import random +import shutil +import zipfile +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import torch + +SOURCE_DIR = Path(__file__).resolve().parent +REPO_ROOT = SOURCE_DIR.parent +if str(SOURCE_DIR) not in __import__("sys").path: + __import__("sys").path.append(str(SOURCE_DIR)) + +try: + from image_recognition_model import SCIENCE_IMAGE_CLASSES + from omni_collective_model import OMNI_DOMAIN_LABELS_V2, OMNI_INTENTS_V2, build_char_vocab + from omni_collective_v41_model import OmniCollectiveEngineV41, OmniCollectiveNetV41 + from omni_training_runtime import maybe_compile_model, resolve_training_runtime + from prepare_omni_collective_v41 import build_v41_blueprint, latest_v8_summary_path, load_summary + from train_omni_collective_v2 import DEFAULT_MODELS_DIR, OmniRow, _normalize, split_rows + from train_omni_collective_v4 import _load_expanded_state_from_zip + from train_omni_collective_v8 import ( + _train_stage_resumable_v8, + build_training_rows as build_training_rows_v8, + ) +except ImportError: # pragma: no cover + from .image_recognition_model import SCIENCE_IMAGE_CLASSES + from .omni_collective_model import OMNI_DOMAIN_LABELS_V2, OMNI_INTENTS_V2, build_char_vocab + from .omni_collective_v41_model import OmniCollectiveEngineV41, OmniCollectiveNetV41 + from .omni_training_runtime import maybe_compile_model, resolve_training_runtime + from .prepare_omni_collective_v41 import build_v41_blueprint, latest_v8_summary_path, load_summary + from .train_omni_collective_v2 import DEFAULT_MODELS_DIR, OmniRow, _normalize, split_rows + from .train_omni_collective_v4 import _load_expanded_state_from_zip + from .train_omni_collective_v8 import ( + _train_stage_resumable_v8, + build_training_rows as build_training_rows_v8, + ) + + +DEFAULT_OUTPUT_ROOT = REPO_ROOT / "output" / "omni_collective_v41_prep" +DEFAULT_TRAIN_OUTPUT_DIR = REPO_ROOT / "output" +DEFAULT_BASE_ZIP = REPO_ROOT / "output" / "supermix_omni_collective_v8_frontier_20260408.zip" +V41_MODEL_CONFIG: Dict[str, Any] = { + "family": "omni_collective_v41", + "base_model_family": "omni_collective_v8", + "text_hidden": 288, + "fusion_hidden": 1152, + "expert_count": 12, + "expert_hidden": 1920, + "expert_top_k": 2, + "deliberation_passes": 12, + "minimum_passes": 6, + "new_heads": [ + "latent_plan_slots", + "multi_token_prediction_aux", + "communication_polish_head", + "uncertainty_anchor_head", + ], +} + + +def _safe_domain(domain: str) -> str: + cooked = str(domain or "general") + if cooked in OMNI_DOMAIN_LABELS_V2: + return cooked + if cooked in {"reasoning", "materials", "protein"}: + return "knowledge" + if cooked == "comparison": + return "general" + return "general" + + +def _safe_intent(intent: str, *, domain: str) -> str: + cooked = str(intent or "general") + if cooked in OMNI_INTENTS_V2: + return cooked + domain_key = _safe_domain(domain) + if domain_key in {"coding", "model_selection", "image_prompt", "spatial_3d", "math", "vision", "language", "knowledge"}: + return "vision" if domain_key == "vision" else domain_key + if cooked == "reasoning": + return "comparison" + return "general" + + +def _eval_payload( + prompt: str, + *, + expected: str, + focus: str, + metric: str, + source: str, +) -> Dict[str, str]: + return { + "prompt": _normalize(prompt, 420), + "expected": _normalize(expected, 520), + "focus": str(focus), + "metric": str(metric), + "source": str(source), + } + + +def _row( + prompt: str, + response: str, + *, + intent: str = "general", + domain: str = "general", + source: str, +) -> OmniRow: + safe_domain = _safe_domain(str(domain)) + return OmniRow( + prompt=_normalize(str(prompt), 420), + intent=_safe_intent(str(intent), domain=safe_domain), + response_text=_normalize(str(response), 520), + domain=safe_domain, + source=str(source), + ) + + +def _row_payload(row: OmniRow) -> Dict[str, Any]: + return { + "prompt": row.prompt, + "intent": row.intent, + "response_text": row.response_text, + "domain": row.domain, + "source": row.source, + "image_path": row.image_path, + "vision_label": row.vision_label, + } + + +def _write_json(path: Path, payload: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8") + + +def _write_jsonl(path: Path, rows: Sequence[OmniRow]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + lines = [json.dumps(_row_payload(row), ensure_ascii=True) for row in rows] + path.write_text(("\n".join(lines) + "\n") if lines else "", encoding="utf-8") + + +def _read_jsonl_rows(path: Path) -> List[OmniRow]: + rows: List[OmniRow] = [] + if not path.exists(): + return rows + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line: + continue + payload = json.loads(line) + rows.append( + OmniRow( + prompt=str(payload.get("prompt") or ""), + intent=_safe_intent(str(payload.get("intent") or "general"), domain=str(payload.get("domain") or "general")), + response_text=str(payload.get("response_text") or ""), + domain=_safe_domain(str(payload.get("domain") or "general")), + source=str(payload.get("source") or "v41_prep"), + image_path=str(payload.get("image_path")) if payload.get("image_path") else None, + vision_label=str(payload.get("vision_label")) if payload.get("vision_label") else None, + ) + ) + return rows + + +def _dedupe_rows(rows: Sequence[OmniRow]) -> List[OmniRow]: + deduped: List[OmniRow] = [] + seen: set[tuple[str, str, str, str]] = set() + for row in rows: + key = ( + str(row.prompt).strip().lower(), + str(row.response_text).strip().lower(), + str(row.domain).strip().lower(), + str(row.source).strip().lower(), + ) + if key in seen: + continue + seen.add(key) + deduped.append(row) + return deduped + + +def _seeded_take(rows: Iterable[OmniRow], *, seed: int, limit: int) -> List[OmniRow]: + cooked = list(rows) + rng = random.Random(int(seed)) + rng.shuffle(cooked) + return cooked[: max(0, int(limit))] + + +def _communication_polish_rows_v41(*, seed: int, limit: int = 180) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + + clarity_pairs = [ + ( + "Rewrite this answer so it is direct and human-friendly without losing the core meaning.\n" + "Draft: The present implementation exhibits a non-trivial probability of behavioral regression under future edits.", + "This implementation is likely to break when people change it later, so it needs regression tests.", + "general", + "communication_polish_v41::clarity_rewrite", + ), + ( + "Rewrite this answer so it sounds more grounded and less like filler.\n" + "Draft: There are many possible avenues worth considering depending on your broader strategic context.", + "The best next step depends on your goal, so first decide whether you care most about speed, quality, or cost.", + "general", + "communication_polish_v41::clarity_rewrite", + ), + ( + "Rewrite this answer so it stays professional but easier to read.\n" + "Draft: We regret the inconvenience and are endeavoring to resolve the issue in due course.", + "We know this is disruptive, and we are working on a fix now.", + "general", + "communication_polish_v41::clarity_rewrite", + ), + ] + for prompt, response, domain, source in clarity_pairs: + rows.append(_row(prompt, response, intent="general", domain=domain, source=source)) + + uncertainty_pairs = [ + ( + "The assistant cannot verify a current fact locally. Give the best grounded reply.", + "I cannot verify that current fact from local information alone, so I should avoid guessing. I can help you reason from what is already available or tell you what source would be needed to confirm it.", + "knowledge", + "communication_polish_v41::grounded_uncertainty", + ), + ( + "Write a short answer that is honest about uncertainty but still useful.\n" + "Question: Is this company still using the same pricing plan today?", + "I cannot confirm today's pricing from local information alone. The safe answer is to check the current vendor pricing page before making a decision.", + "knowledge", + "communication_polish_v41::grounded_uncertainty", + ), + ( + "Give a concise response that refuses to invent missing evidence.\n" + "Question: Which benchmark score did the unpublished model reach?", + "I do not have evidence for that unpublished score, so I should not make one up. If you have the summary or benchmark output, I can interpret it.", + "knowledge", + "communication_polish_v41::grounded_uncertainty", + ), + ] + for prompt, response, domain, source in uncertainty_pairs: + rows.append(_row(prompt, response, intent="knowledge", domain=domain, source=source)) + + compare_pairs = [ + ( + "Answer this like a strong human assistant.\n" + "Question: Compare two ways to solve a bug and recommend one.", + "State both approaches briefly, name the main tradeoff, then recommend one based on speed, risk, and maintainability.", + "reasoning", + "communication_polish_v41::compare_recommend", + ), + ( + "Rewrite this so the recommendation is explicit.\n" + "Draft: Option A has lower risk. Option B may be faster in the short term.", + "I would choose option A because it is safer and easier to maintain, even if option B might be slightly faster today.", + "reasoning", + "communication_polish_v41::compare_recommend", + ), + ( + "Write a compact answer that compares two local models for coding help.", + "Summarize the strengths of both models, then clearly say which one you would use for the user's specific coding task.", + "model_selection", + "communication_polish_v41::compare_recommend", + ), + ] + for prompt, response, domain, source in compare_pairs: + rows.append(_row(prompt, response, intent="general", domain=domain, source=source)) + + coding_pairs = [ + ( + "Answer this coding question more clearly.\n" + "Question: Why do regression tests matter after a refactor?", + "Regression tests make refactors safer because they tell you when behavior changed in a way you did not intend.", + "coding", + "communication_polish_v41::coding_explain", + ), + ( + "Give a concise coding explanation a human teammate would appreciate.\n" + "Question: Why does a failing test with a stack trace help debugging?", + "A failing test narrows the problem and the stack trace points you to the exact path that broke, so you spend less time guessing.", + "coding", + "communication_polish_v41::coding_explain", + ), + ( + "Rewrite this code-review comment so it is direct, respectful, and actionable.\n" + "Draft: This is bad and should probably be redone.", + "This path is fragile because it mixes state updates and file writes. Please split the logic and add a regression test for the failure case.", + "coding", + "communication_polish_v41::coding_explain", + ), + ] + for prompt, response, domain, source in coding_pairs: + rows.append(_row(prompt, response, intent="coding", domain=domain, source=source)) + + creative_pairs = [ + ( + "Give a creative but grounded answer to this request.\n" + "Question: Write a vivid one-paragraph description of a storm-battered lighthouse at night.", + "The lighthouse leaned into the wind like a stubborn white nail driven into black stone, its lamp cutting brief gold arcs across rain that looked almost solid in the dark.", + "creative", + "communication_polish_v41::creative_grounded", + ), + ( + "Rewrite this answer so it is more vivid without becoming purple or fake.\n" + "Draft: The room was messy and tense.", + "The room felt lived-in rather than ruined, with half-finished notes on the table and the kind of silence that made every small movement sound deliberate.", + "creative", + "communication_polish_v41::creative_grounded", + ), + ( + "Write a creative answer that still respects the request exactly.\n" + "Question: Give me two sentences about why curiosity matters.", + "Curiosity keeps your mind from going still when the world gets complicated. It turns confusion into a trail of questions you can actually follow.", + "creative", + "communication_polish_v41::creative_grounded", + ), + ] + for prompt, response, domain, source in creative_pairs: + rows.append(_row(prompt, response, intent="general", domain=domain, source=source)) + + human_collab_pairs = [ + ( + "Write a better collaborator-style answer.\n" + "Question: The build failed again. What do we do now?", + "First capture the exact failing step, then decide whether the issue is code, environment, or packaging. After that, fix the smallest reproducible cause before trying the whole build again.", + "reasoning", + "communication_polish_v41::human_collab", + ), + ( + "Answer like a pragmatic engineer talking to a teammate.\n" + "Question: The run crashed during checkpoint save. What should happen next?", + "Do not restart from zero. Check whether the latest checkpoint or temp checkpoint is valid, free enough disk space, and resume from the newest safe state.", + "reasoning", + "communication_polish_v41::human_collab", + ), + ( + "Write a direct but calm response.\n" + "Question: I think I lost a lot of work.", + "Check what is actually gone before assuming the worst. Look for the last good commit, any saved checkpoints, and any local backup branches first.", + "general", + "communication_polish_v41::human_collab", + ), + ] + for prompt, response, domain, source in human_collab_pairs: + rows.append(_row(prompt, response, intent="general", domain=domain, source=source)) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _latent_plan_rows_v41(*, seed: int, limit: int = 120) -> Tuple[List[OmniRow], Dict[str, int]]: + rows = [ + _row( + ( + "Plan internally before answering.\n" + "Request: Compare two debugging approaches for an intermittent file-write failure.\n" + "Hidden plan target: identify reproducibility, isolate IO path, compare fast patch vs durable fix, recommend durable fix.\n" + "Return the final answer only." + ), + "Start by reproducing the failure consistently, then isolate the exact file-write path and compare a quick retry patch against a safer atomic-write fix. I would choose the durable atomic-write fix because it addresses the root cause instead of hiding it.", + intent="reasoning", + domain="reasoning", + source="latent_plan_v41::debugging", + ), + _row( + ( + "Route first, then answer.\n" + "Request: Which local model should handle a benchmark-style reasoning prompt?\n" + "Hidden plan target: classify request as benchmark reasoning, compare omni vs specialist, choose benchmark specialist, answer directly.\n" + "Return the final answer only." + ), + "For benchmark-style reasoning prompts, v40_benchmax is the best local choice because it is the specialist benchmark-maximization line rather than the broader multimodal omni model.", + intent="model_selection", + domain="model_selection", + source="latent_plan_v41::route_selection", + ), + _row( + ( + "Think in a short hidden plan, then answer.\n" + "Request: Explain to a non-technical user why a saved checkpoint matters.\n" + "Hidden plan target: define checkpoint, explain recovery value, keep tone plain, avoid jargon.\n" + "Return the final answer only." + ), + "A checkpoint is a saved snapshot of the work so far. If a run crashes or the computer restarts, it lets you continue from a recent point instead of starting again from the beginning.", + intent="general", + domain="general", + source="latent_plan_v41::human_explain", + ), + _row( + ( + "Use a hidden plan before writing code.\n" + "Request: Write an OpenSCAD snippet for a centered cylinder with a hole.\n" + "Hidden plan target: choose difference, keep both cylinders centered, preserve valid syntax.\n" + "Return the final answer only." + ), + "difference(){ cylinder(h=20, r=12, center=true, $fn=64); cylinder(h=22, r=4, center=true, $fn=48); }", + intent="knowledge", + domain="spatial_3d", + source="latent_plan_v41::openscad", + ), + ] + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _code_critique_repair_rows_v41(*, seed: int, limit: int = 160) -> Tuple[List[OmniRow], Dict[str, int]]: + rows = [ + _row( + ( + "A checkpoint save crashed during atomic rename on Windows.\n" + "Traceback: PermissionError: [WinError 5] Access is denied during replace into stage1_progress.pt\n" + "Task: Explain the likely root cause and the right durable fix." + ), + "The likely issue is a Windows file-lock or antivirus race during the final rename. The durable fix is to keep the temp-file recovery path, retry or fall back safely, and load the newest valid checkpoint rather than crashing the whole run.", + intent="coding", + domain="coding", + source="code_repair_v41::checkpoint_write", + ), + _row( + ( + "A reviewer says this patch is too vague.\n" + "Draft comment: This code is bad and should be redone.\n" + "Rewrite it into a useful code-review comment." + ), + "This path is fragile because it mixes state updates and file writes in one step. Please split the responsibilities and add a regression test for the failed-write case.", + intent="coding", + domain="coding", + source="code_repair_v41::review_rewrite", + ), + _row( + ( + "A bug report says the app selector routed a benchmark prompt to vision mode.\n" + "Task: Give the most likely root cause and the first fix to try." + ), + "The routing heuristics likely over-weighted a vision keyword and skipped the stronger benchmark-reasoning match. The first fix is to add an explicit benchmark-reasoning route check before generic vision fallbacks.", + intent="coding", + domain="coding", + source="code_repair_v41::routing_bug", + ), + _row( + ( + "A failing test says an OpenSCAD snippet is not centered.\n" + "Bad code: difference(){cylinder(h=20, r=12); translate([0,0,-1]) cylinder(h=22, r=4);} \n" + "Task: Return a corrected version." + ), + "difference(){ cylinder(h=20, r=12, center=true, $fn=64); cylinder(h=22, r=4, center=true, $fn=48); }", + intent="coding", + domain="coding", + source="code_repair_v41::openscad_fix", + ), + _row( + ( + "A stack trace shows OSError: [Errno 28] No space left on device during checkpoint save.\n" + "Task: Tell the operator exactly what to do next." + ), + "Do not restart from zero. Free disk space, keep the newest valid checkpoint or temp checkpoint, and resume from that saved stage state instead of throwing away the run.", + intent="coding", + domain="coding", + source="code_repair_v41::disk_full_resume", + ), + ] + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _reasoning_budget_rows_v41(*, seed: int, limit: int = 180) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + specs = [ + ( + "Explain why regression tests matter after a refactor.", + { + "short": "Regression tests catch unintended behavior changes quickly.", + "medium": "Regression tests matter after a refactor because structure can change without you noticing a behavioral break. They give you a fast safety signal that the code still does what it did before.", + "deep": "Regression tests matter after a refactor because refactors change structure, dependencies, and control flow even when the goal is to preserve behavior. A strong regression suite converts that promise into something checkable: it exposes unintended changes early, narrows debugging when something breaks, and gives you confidence to keep improving the design instead of freezing the code out of fear.", + }, + "coding", + ), + ( + "Compare two ways to recover from a failed training checkpoint save.", + { + "short": "Use the newest valid checkpoint if possible; otherwise restart from the last safe stage checkpoint.", + "medium": "The first option is to recover from the newest valid temp or final checkpoint, which preserves progress and is usually the right choice. The fallback is to restart from the last safe stage checkpoint if the newest file is corrupted.", + "deep": "There are two realistic recovery paths. If the newest temp or final checkpoint can still be loaded, resume from that state because it preserves optimizer and stage progress. If the newest save is corrupted, fall back to the last durable stage checkpoint and treat the failed save as an operational reliability issue, not a reason to throw away the entire run. The right recommendation is almost always to validate the checkpoint lineage first, then restart from the newest safe state.", + }, + "reasoning", + ), + ] + for prompt, answers, domain in specs: + for budget, response in answers.items(): + rows.append( + _row( + ( + f"Reasoning budget: {budget}.\n" + f"Request: {prompt}\n" + "Return an answer that matches the requested depth." + ), + response, + intent="reasoning", + domain=domain, + source=f"reasoning_budget_v41::{budget}", + ) + ) + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _default_disagreement_repairs() -> List[Dict[str, str]]: + return [ + { + "prompt": "Reply in exactly two sentences explaining why regression tests matter.", + "bad_answer": "Best grounded answer from my local training: [Bot responding in a professional and direct manner, addressing the user's situation directly and offering relevant engagement or assistance regarding a delayed flight.]", + "good_answer": "Regression tests catch behavior changes that a refactor accidentally introduces. They give you a fast safety check so you can improve code without guessing whether it still works.", + "domain": "coding", + "intent": "coding", + }, + { + "prompt": "What should a grounded assistant do when it cannot verify a current fact locally?", + "bad_answer": "Short brief (February 4, 2026): AP reported IOC discussions about shifting future Winter Games earlier in the calendar to reduce climate-related snow reliability risks.", + "good_answer": "State clearly that the fact cannot be verified from local information, avoid inventing specifics, and say what source or tool would be needed to confirm it.", + "domain": "knowledge", + "intent": "knowledge", + }, + { + "prompt": "Which local model is best for benchmark-focused reasoning prompts?", + "bad_answer": "Best grounded answer from my local training: Upload an image first, then I can identify the science concept and explain the visual clues.", + "good_answer": "For benchmark-focused reasoning prompts, v40_benchmax is the strongest local choice because it is the benchmark-specialized line rather than a broader multimodal generalist.", + "domain": "model_selection", + "intent": "model_selection", + }, + { + "prompt": "Generate a CIF-style seed for a thermoelectric oxide candidate with a low thermal conductivity target.", + "bad_answer": "Best grounded answer from my local training: Short guide: Memoize expensive calculations, keep component props stable, and avoid unnecessary parent re-renders.", + "good_answer": "A grounded seed would start from a layered or distorted perovskite-style oxide candidate, note the low thermal conductivity target explicitly, and avoid pretending to provide fully validated ab initio properties.", + "domain": "materials", + "intent": "knowledge", + }, + { + "prompt": "Write a tiny OpenSCAD snippet for a centered cylinder with a hole.", + "bad_answer": "difference(){cylinder(h=20, r=12, $fn=64); translate([0,0,-1]) cylinder(h=22, r=4, $fn=48);}", + "good_answer": "difference(){ cylinder(h=20, r=12, center=true, $fn=64); cylinder(h=22, r=4, center=true, $fn=48); }", + "domain": "spatial_3d", + "intent": "knowledge", + }, + ] + + +def _teacher_disagreement_rows_v41( + *, + summary: Dict[str, Any], + seed: int, + limit: int = 120, +) -> Tuple[List[OmniRow], Dict[str, Any]]: + teacher_keys = [ + str(item) + for item in summary.get("dataset_summary", {}).get("teacher_league", {}).get("teacher_keys", []) + if str(item).strip() + ] + teacher_a = teacher_keys[0] if teacher_keys else "v40_benchmax" + teacher_b = teacher_keys[1] if len(teacher_keys) > 1 else "qwen_v30" + teacher_c = teacher_keys[2] if len(teacher_keys) > 2 else "omni_collective_v8" + + rows: List[OmniRow] = [] + for item in _default_disagreement_repairs(): + prompt = item["prompt"] + bad_answer = item["bad_answer"] + good_answer = item["good_answer"] + domain = item["domain"] + intent = item["intent"] + rows.append( + _row( + ( + "Two strong local teachers disagree on the answer. Return the most grounded final answer.\n" + f"Request: {prompt}\n" + f"Draft from {teacher_a}: {bad_answer}\n" + f"Draft from {teacher_b}: {good_answer}" + ), + good_answer, + intent=intent, + domain=domain, + source="teacher_disagreement_v41::resolve", + ) + ) + rows.append( + _row( + ( + "Choose the better route before answering.\n" + f"Request: {prompt}\n" + f"Wrong route from {teacher_a}: {bad_answer}\n" + f"Better route from {teacher_c}: {good_answer}\n" + "Return the final answer only." + ), + good_answer, + intent=intent, + domain=domain, + source="teacher_disagreement_v41::route_then_answer", + ) + ) + rows.append( + _row( + ( + "Critique the weak draft implicitly and answer correctly.\n" + f"Request: {prompt}\n" + f"Weak draft: {bad_answer}\n" + "Provide the corrected grounded answer." + ), + good_answer, + intent=intent, + domain=domain, + source="teacher_disagreement_v41::repair", + ) + ) + + sample_outputs = summary.get("sample_outputs") or [] + for sample in sample_outputs: + prompt = _normalize(str(sample.get("prompt") or ""), 260) + answer = _normalize(str(sample.get("answer") or ""), 420) + if len(prompt) < 12 or len(answer) < 12: + continue + rows.append( + _row( + ( + "A local model answered this prompt, but another strong teacher may disagree.\n" + f"Request: {prompt}\n" + f"Current answer: {answer}\n" + "Return a more grounded, more useful final answer if the current one is weak." + ), + answer, + intent="general", + domain="general", + source="teacher_disagreement_v41::stability", + ) + ) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, { + "teacher_pair": [teacher_a, teacher_b, teacher_c], + "row_count": len(selected), + "source_counts": dict(sorted(counts.items())), + } + + +def _promotion_eval_pack_v41() -> List[Dict[str, str]]: + return [ + _eval_payload( + "Reply in exactly two sentences explaining why regression tests matter.", + expected="Two grounded sentences that clearly explain safety against unintended behavior changes.", + focus="communication_and_coding", + metric="clarity_and_correctness", + source="promotion_eval_v41::regression_tests", + ), + _eval_payload( + "Which local model is best for benchmark-focused reasoning prompts?", + expected="Should explicitly prefer v40_benchmax or clearly justify a stronger newer benchmark specialist if one exists.", + focus="routing_and_model_selection", + metric="grounded_route_choice", + source="promotion_eval_v41::benchmark_route", + ), + _eval_payload( + "A checkpoint save failed on Windows with Access is denied. What should happen next?", + expected="Should recommend validating newest checkpoint or temp checkpoint, then resuming instead of restarting from zero.", + focus="coding_and_operations", + metric="repair_quality", + source="promotion_eval_v41::checkpoint_recovery", + ), + _eval_payload( + "What should a grounded assistant do when it cannot verify a current fact locally?", + expected="Must avoid guessing, state uncertainty clearly, and say what evidence or source is needed.", + focus="groundedness", + metric="uncertainty_honesty", + source="promotion_eval_v41::uncertainty", + ), + _eval_payload( + "Write a vivid but grounded paragraph about a storm-battered lighthouse at night.", + expected="Should be creative and readable without fabricating irrelevant facts or collapsing into bland filler.", + focus="creative_communication", + metric="creative_quality", + source="promotion_eval_v41::creative_grounded", + ), + _eval_payload( + "Generate a CIF-style seed for a thermoelectric oxide candidate with low thermal conductivity target.", + expected="Should stay within bounded materials language and avoid pretending to provide fully validated physical properties.", + focus="materials_grounding", + metric="specialist_groundedness", + source="promotion_eval_v41::materials", + ), + ] + + +def _bundle_rows(rows: Sequence[OmniRow]) -> Dict[str, Any]: + counts: Dict[str, int] = {} + domains: Dict[str, int] = {} + intents: Dict[str, int] = {} + for row in rows: + counts[row.source] = counts.get(row.source, 0) + 1 + domains[row.domain] = domains.get(row.domain, 0) + 1 + intents[row.intent] = intents.get(row.intent, 0) + 1 + return { + "row_count": len(rows), + "source_counts": dict(sorted(counts.items())), + "domain_counts": dict(sorted(domains.items())), + "intent_counts": dict(sorted(intents.items())), + } + + +def assemble_v41_training_rows( + *, + base_stage1_rows: Sequence[OmniRow], + base_full_rows: Sequence[OmniRow], + base_summary: Dict[str, Any], + prep_root: Path, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + prep_root = prep_root.resolve() + communication_rows = _read_jsonl_rows(prep_root / "v41_communication_polish_rows.jsonl") + latent_plan_rows = _read_jsonl_rows(prep_root / "v41_latent_plan_rows.jsonl") + code_repair_rows = _read_jsonl_rows(prep_root / "v41_code_critique_repair_rows.jsonl") + reasoning_budget_rows = _read_jsonl_rows(prep_root / "v41_reasoning_budget_rows.jsonl") + disagreement_rows = _read_jsonl_rows(prep_root / "v41_teacher_disagreement_rows.jsonl") + + added_rows = list(communication_rows) + list(latent_plan_rows) + list(code_repair_rows) + list(reasoning_budget_rows) + list(disagreement_rows) + full_rows = list(base_full_rows) + added_rows + stage1_rows = [row for row in full_rows if row.domain not in {"vision", "spatial_3d", "video"}] + + source_counts = dict(base_summary.get("source_counts") or {}) + for row in added_rows: + source_counts[row.source] = source_counts.get(row.source, 0) + 1 + + summary = { + "stage1_rows": len(stage1_rows), + "stage2_rows": len(full_rows), + "pre_dedupe_rows": len(full_rows), + "source_counts": dict(sorted(source_counts.items())), + "base_stage1_rows": len(base_stage1_rows), + "base_stage2_rows": len(base_full_rows), + "v41_added_rows": len(added_rows), + "v41_prep_root": str(prep_root), + "v41_row_groups": { + "communication_polish": _bundle_rows(communication_rows), + "latent_plan": _bundle_rows(latent_plan_rows), + "code_critique_repair": _bundle_rows(code_repair_rows), + "reasoning_budget": _bundle_rows(reasoning_budget_rows), + "teacher_disagreement": _bundle_rows(disagreement_rows), + }, + "base_summary": base_summary, + } + return stage1_rows, full_rows, summary + + +def build_training_rows_v41( + *, + repo_root: Path, + models_dir: Path, + images_dir: Path, + summary_path: Path, + output_root: Path, + seed: int = 41, + communication_limit: int = 180, + disagreement_limit: int = 120, + base_distill_limit: int = 0, + base_teacher_model_limit: int = 0, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + summary = load_summary(summary_path) + prep_payload = build_v41_prep_pack( + summary_path=summary_path, + output_root=output_root, + seed=seed, + communication_limit=communication_limit, + disagreement_limit=disagreement_limit, + ) + frozen_dataset_summary = summary.get("dataset_summary") if isinstance(summary.get("dataset_summary"), dict) else None + base_stage1_rows, base_full_rows, base_summary = build_training_rows_v8( + repo_root=repo_root, + models_dir=models_dir, + images_dir=images_dir, + seed=seed, + distill_limit=base_distill_limit, + teacher_model_limit=base_teacher_model_limit, + frozen_dataset_summary=frozen_dataset_summary, + ) + stage1_rows, full_rows, merged_summary = assemble_v41_training_rows( + base_stage1_rows=base_stage1_rows, + base_full_rows=base_full_rows, + base_summary=base_summary, + prep_root=Path(prep_payload["output_root"]), + ) + merged_summary["prep_payload"] = prep_payload + return stage1_rows, full_rows, merged_summary + + +def build_training_rows_v41_dry_run( + *, + summary_path: Path, + output_root: Path, + seed: int = 41, + communication_limit: int = 180, + disagreement_limit: int = 120, +) -> Dict[str, Any]: + summary = load_summary(summary_path) + prep_payload = build_v41_prep_pack( + summary_path=summary_path, + output_root=output_root, + seed=seed, + communication_limit=communication_limit, + disagreement_limit=disagreement_limit, + ) + prep_root = Path(prep_payload["output_root"]).resolve() + + communication_rows = _read_jsonl_rows(prep_root / "v41_communication_polish_rows.jsonl") + latent_plan_rows = _read_jsonl_rows(prep_root / "v41_latent_plan_rows.jsonl") + code_repair_rows = _read_jsonl_rows(prep_root / "v41_code_critique_repair_rows.jsonl") + reasoning_budget_rows = _read_jsonl_rows(prep_root / "v41_reasoning_budget_rows.jsonl") + disagreement_rows = _read_jsonl_rows(prep_root / "v41_teacher_disagreement_rows.jsonl") + + added_rows = list(communication_rows) + list(latent_plan_rows) + list(code_repair_rows) + list(reasoning_budget_rows) + list(disagreement_rows) + added_stage1_rows = [row for row in added_rows if row.domain not in {"vision", "spatial_3d", "video"}] + frozen_dataset_summary = dict(summary.get("dataset_summary") or {}) + base_stage1_count = int(frozen_dataset_summary.get("stage1_rows") or 0) + base_stage2_count = int(frozen_dataset_summary.get("stage2_rows") or 0) + source_counts = dict(frozen_dataset_summary.get("source_counts") or {}) + for row in added_rows: + source_counts[row.source] = source_counts.get(row.source, 0) + 1 + + dry_run_summary = { + "dry_run": True, + "base_mode": "frozen_counts_plus_v41_rows", + "source_summary": str(summary_path.resolve()), + "prep_payload": prep_payload, + "base_stage1_rows": base_stage1_count, + "base_stage2_rows": base_stage2_count, + "added_stage1_rows": len(added_stage1_rows), + "added_stage2_rows": len(added_rows), + "estimated_stage1_rows": base_stage1_count + len(added_stage1_rows), + "estimated_stage2_rows": base_stage2_count + len(added_rows), + "source_counts": dict(sorted(source_counts.items())), + "v41_row_groups": { + "communication_polish": _bundle_rows(communication_rows), + "latent_plan": _bundle_rows(latent_plan_rows), + "code_critique_repair": _bundle_rows(code_repair_rows), + "reasoning_budget": _bundle_rows(reasoning_budget_rows), + "teacher_disagreement": _bundle_rows(disagreement_rows), + }, + } + dry_run_path = Path(output_root).resolve() / "omni_collective_v41_dry_run_summary.json" + _write_json(dry_run_path, dry_run_summary) + return dry_run_summary | {"summary_path": str(dry_run_path)} + + +def build_v41_prep_pack( + *, + summary_path: Path, + output_root: Path, + seed: int = 41, + communication_limit: int = 180, + disagreement_limit: int = 120, +) -> Dict[str, Any]: + summary = load_summary(summary_path) + blueprint = build_v41_blueprint(summary, summary_path=summary_path) + communication_rows, communication_summary = _communication_polish_rows_v41( + seed=seed + 11, + limit=communication_limit, + ) + latent_plan_rows, latent_plan_summary = _latent_plan_rows_v41(seed=seed + 17, limit=max(24, communication_limit // 3)) + code_repair_rows, code_repair_summary = _code_critique_repair_rows_v41(seed=seed + 23, limit=max(30, disagreement_limit // 2)) + reasoning_budget_rows, reasoning_budget_summary = _reasoning_budget_rows_v41(seed=seed + 27, limit=max(24, communication_limit // 2)) + disagreement_rows, disagreement_summary = _teacher_disagreement_rows_v41( + summary=summary, + seed=seed + 29, + limit=disagreement_limit, + ) + promotion_eval_pack = _promotion_eval_pack_v41() + + output_root = output_root.resolve() + blueprint_path = output_root / "omni_collective_v41_blueprint.json" + communication_path = output_root / "v41_communication_polish_rows.jsonl" + latent_plan_path = output_root / "v41_latent_plan_rows.jsonl" + code_repair_path = output_root / "v41_code_critique_repair_rows.jsonl" + reasoning_budget_path = output_root / "v41_reasoning_budget_rows.jsonl" + disagreement_path = output_root / "v41_teacher_disagreement_rows.jsonl" + promotion_eval_path = output_root / "v41_promotion_eval_pack.json" + prep_summary_path = output_root / "omni_collective_v41_prep_summary.json" + + _write_json(blueprint_path, blueprint) + _write_jsonl(communication_path, communication_rows) + _write_jsonl(latent_plan_path, latent_plan_rows) + _write_jsonl(code_repair_path, code_repair_rows) + _write_jsonl(reasoning_budget_path, reasoning_budget_rows) + _write_jsonl(disagreement_path, disagreement_rows) + _write_json(promotion_eval_path, {"items": promotion_eval_pack}) + + payload = { + "family": "omni_collective_v41", + "source_summary": str(summary_path.resolve()), + "output_root": str(output_root), + "model_config": dict(V41_MODEL_CONFIG), + "communication_polish": { + "row_count": len(communication_rows), + "source_counts": communication_summary, + "path": str(communication_path), + }, + "latent_plan": { + "row_count": len(latent_plan_rows), + "source_counts": latent_plan_summary, + "path": str(latent_plan_path), + }, + "code_critique_repair": { + "row_count": len(code_repair_rows), + "source_counts": code_repair_summary, + "path": str(code_repair_path), + }, + "reasoning_budget": { + "row_count": len(reasoning_budget_rows), + "source_counts": reasoning_budget_summary, + "path": str(reasoning_budget_path), + }, + "teacher_disagreement": disagreement_summary | {"path": str(disagreement_path)}, + "promotion_eval_pack": { + "item_count": len(promotion_eval_pack), + "path": str(promotion_eval_path), + }, + "blueprint_path": str(blueprint_path), + } + _write_json(prep_summary_path, payload) + return payload + + +def _write_json_atomic_v41(path: Path, payload: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_suffix(path.suffix + ".tmp") + temp_path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8") + temp_path.replace(path) + + +def _run_state_path_v41(output_dir: Path) -> Path: + return output_dir / "omni_collective_v41_train_state.json" + + +def _stage_resume_dir_v41( + output_dir: Path, + *, + seed: int, + distill_limit: int, + teacher_model_limit: int, + smoke_train: bool, +) -> Path: + teacher_tag = int(teacher_model_limit) if int(teacher_model_limit) > 0 else 0 + mode_tag = "smoke" if smoke_train else "frontier" + return output_dir / "omni_v41_stage_resume" / f"{mode_tag}_seed_{int(seed)}_distill_{int(distill_limit)}_teacherlimit_{teacher_tag}" + + +def _write_run_state_v41(path: Path, payload: Dict[str, Any]) -> None: + cooked = dict(payload) + cooked["updated_at"] = datetime.now().isoformat() + _write_json_atomic_v41(path, cooked) + + +def _smoke_training_rows_v41( + *, + stage1_rows: Sequence[OmniRow], + full_rows: Sequence[OmniRow], + seed: int, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + stage1_priority = [row for row in stage1_rows if "_v41::" in row.source or "_v41" in row.source] + full_priority = [row for row in full_rows if "_v41::" in row.source or "_v41" in row.source] + base_stage1 = [row for row in stage1_rows if row not in stage1_priority] + multimodal = [row for row in full_rows if row.domain in {"vision", "spatial_3d", "video"}] + + selected_stage1 = _dedupe_rows( + _seeded_take(stage1_priority, seed=seed + 1, limit=22) + + _seeded_take(base_stage1, seed=seed + 2, limit=18) + ) + selected_full = _dedupe_rows( + list(selected_stage1) + + _seeded_take(full_priority, seed=seed + 3, limit=10) + + _seeded_take(multimodal, seed=seed + 4, limit=10) + ) + selected_stage1 = [row for row in selected_full if row.domain not in {"vision", "spatial_3d", "video"}] + summary = { + "mode": "smoke", + "stage1_rows": len(selected_stage1), + "stage2_rows": len(selected_full), + "v41_priority_rows": len([row for row in selected_full if "_v41::" in row.source or "_v41" in row.source]), + "multimodal_rows": len([row for row in selected_full if row.domain in {"vision", "spatial_3d", "video"}]), + "source_counts": _bundle_rows(selected_full)["source_counts"], + } + return selected_stage1, selected_full, summary + + +def train_model_v41( + *, + repo_root: Path, + output_dir: Path, + models_dir: Path, + base_zip: Path, + images_dir: Path, + summary_path: Path, + prep_output_root: Path, + image_size: int, + batch_size: int, + stage1_epochs: int, + stage2_epochs: int, + stage1_lr: float, + stage2_lr: float, + seed: int, + communication_limit: int, + disagreement_limit: int, + base_distill_limit: int, + base_teacher_model_limit: int, + requested_device: str, + amp_mode: str, + amp_dtype: str, + compile_model: bool, + compile_mode: str, + grad_accum_steps: int, + ema_decay: float, + warmup_steps: int, + warmup_ratio: float, + min_lr_scale: float, + smoke_train: bool = False, +) -> Dict[str, Any]: + torch.manual_seed(int(seed)) + random.seed(int(seed)) + output_dir.mkdir(parents=True, exist_ok=True) + stage_resume_dir = _stage_resume_dir_v41( + output_dir, + seed=int(seed), + distill_limit=int(base_distill_limit), + teacher_model_limit=int(base_teacher_model_limit), + smoke_train=bool(smoke_train), + ) + stage_resume_dir.mkdir(parents=True, exist_ok=True) + run_state_path = _run_state_path_v41(output_dir) + _write_run_state_v41( + run_state_path, + { + "status": "building_dataset", + "stage": "dataset", + "resume_dir": str(stage_resume_dir), + "smoke_train": bool(smoke_train), + }, + ) + + stage1_rows, full_rows, dataset_summary = build_training_rows_v41( + repo_root=repo_root, + models_dir=models_dir, + images_dir=images_dir, + summary_path=summary_path, + output_root=prep_output_root, + seed=seed, + communication_limit=communication_limit, + disagreement_limit=disagreement_limit, + base_distill_limit=base_distill_limit, + base_teacher_model_limit=base_teacher_model_limit, + ) + if smoke_train: + stage1_rows, full_rows, smoke_summary = _smoke_training_rows_v41( + stage1_rows=stage1_rows, + full_rows=full_rows, + seed=seed, + ) + dataset_summary = dict(dataset_summary) + dataset_summary["smoke_subset"] = smoke_summary + dataset_summary["stage1_rows"] = len(stage1_rows) + dataset_summary["stage2_rows"] = len(full_rows) + + print(json.dumps({"event": "dataset_built", "summary": dataset_summary}, ensure_ascii=True), flush=True) + _write_json_atomic_v41(stage_resume_dir / "dataset_summary.json", {"dataset_summary": dataset_summary}) + _write_run_state_v41( + run_state_path, + { + "status": "dataset_built", + "stage": "dataset", + "resume_dir": str(stage_resume_dir), + "dataset_summary": dataset_summary, + "smoke_train": bool(smoke_train), + }, + ) + + train_rows, val_rows = split_rows(full_rows, seed=seed + 17) + train_stage1 = [row for row in train_rows if row.domain not in {"vision", "spatial_3d", "video"}] + vocab = build_char_vocab([row.prompt for row in train_rows], min_frequency=1) + response_bank = sorted({row.response_text for row in full_rows if row.response_text}) + print( + json.dumps( + { + "event": "label_space", + "train_rows": len(train_rows), + "val_rows": len(val_rows), + "response_bank": len(response_bank), + "vocab_size": len(vocab), + "smoke_train": bool(smoke_train), + }, + ensure_ascii=True, + ), + flush=True, + ) + _write_run_state_v41( + run_state_path, + { + "status": "label_space_ready", + "stage": "label_space", + "resume_dir": str(stage_resume_dir), + "train_rows": len(train_rows), + "val_rows": len(val_rows), + "response_bank": len(response_bank), + "vocab_size": len(vocab), + "smoke_train": bool(smoke_train), + }, + ) + + max_len = 420 + max_words = 96 + word_buckets = 18432 + runtime = resolve_training_runtime( + repo_root=repo_root, + requested_device=requested_device, + amp_mode=amp_mode, + amp_dtype=amp_dtype, + compile_requested=compile_model, + compile_mode=compile_mode, + grad_accum_steps=grad_accum_steps, + ema_decay=ema_decay, + warmup_steps=warmup_steps, + warmup_ratio=warmup_ratio, + min_lr_scale=min_lr_scale, + batch_size=batch_size, + ) + print(json.dumps({"event": "runtime_config", "runtime": runtime.to_payload()}, ensure_ascii=True), flush=True) + model = OmniCollectiveNetV41( + vocab_size=max(len(vocab), 2), + num_intents=len(OMNI_INTENTS_V2), + num_responses=max(len(response_bank), 1), + num_vision_classes=len(SCIENCE_IMAGE_CLASSES), + num_domains=len(OMNI_DOMAIN_LABELS_V2), + base_embed_dim=140, + text_hidden=int(V41_MODEL_CONFIG["text_hidden"]), + image_channels=60, + word_buckets=word_buckets, + word_embed_dim=128, + deep_text_channels=416, + deep_image_channels=136, + fusion_hidden=int(V41_MODEL_CONFIG["fusion_hidden"]), + memory_slots=32, + depth_steps=12, + expert_count=int(V41_MODEL_CONFIG["expert_count"]), + expert_hidden=int(V41_MODEL_CONFIG["expert_hidden"]), + context_top_k=4, + expert_top_k=int(V41_MODEL_CONFIG["expert_top_k"]), + ).to(runtime.device) + warm_start = _load_expanded_state_from_zip(model, base_zip) + forward_model, runtime = maybe_compile_model(model, runtime) + print(json.dumps({"event": "warm_start", "info": warm_start}, ensure_ascii=True), flush=True) + _write_run_state_v41( + run_state_path, + { + "status": "warm_start_complete", + "stage": "warm_start", + "resume_dir": str(stage_resume_dir), + "warm_start": warm_start, + "runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + }, + ) + + progress_every = 1 if smoke_train else None + checkpoint_every = 4 if smoke_train else None + stage1 = _train_stage_resumable_v8( + model=model, + forward_model=forward_model, + train_rows=train_stage1, + val_rows=val_rows, + vocab=vocab, + response_bank=response_bank, + image_size=image_size, + max_len=max_len, + max_words=max_words, + word_buckets=word_buckets, + batch_size=batch_size, + learning_rate=stage1_lr, + epochs=stage1_epochs, + seed=seed + 101, + runtime=runtime, + loss_weights={"intent": 0.56, "response": 1.08, "domain": 0.76, "vision": 0.60}, + balance_weight=0.042, + stage_name="stage1", + stage_dir=stage_resume_dir, + run_state_path=run_state_path, + progress_every_batches=progress_every, + checkpoint_every_batches=checkpoint_every, + grad_accum_steps=grad_accum_steps, + ) + stage2 = _train_stage_resumable_v8( + model=model, + forward_model=forward_model, + train_rows=train_rows, + val_rows=val_rows, + vocab=vocab, + response_bank=response_bank, + image_size=image_size, + max_len=max_len, + max_words=max_words, + word_buckets=word_buckets, + batch_size=max(4, int(batch_size) // 2), + learning_rate=stage2_lr, + epochs=stage2_epochs, + seed=seed + 151, + runtime=runtime, + loss_weights={"intent": 0.52, "response": 1.12, "domain": 0.82, "vision": 1.08}, + balance_weight=0.058, + stage_name="stage2", + stage_dir=stage_resume_dir, + run_state_path=run_state_path, + progress_every_batches=progress_every, + checkpoint_every_batches=checkpoint_every, + grad_accum_steps=grad_accum_steps, + ) + + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + mode_tag = "smoke" if smoke_train else "frontier" + artifact_dir = output_dir / f"supermix_omni_collective_v41_{mode_tag}_{stamp}" + artifact_dir.mkdir(parents=True, exist_ok=True) + weights_path = artifact_dir / f"omni_collective_v41_{mode_tag}.pth" + meta_path = artifact_dir / f"omni_collective_v41_{mode_tag}_meta.json" + summary_out_path = artifact_dir / f"omni_collective_v41_{mode_tag}_summary.json" + zip_path = output_dir / f"supermix_omni_collective_v41_{mode_tag}_{stamp}.zip" + desktop_zip_path = models_dir / zip_path.name + + torch.save(model.state_dict(), weights_path) + parameter_count = int(sum(parameter.numel() for parameter in model.parameters())) + meta = { + "architecture_version": 41, + "family": "omni_collective_v41", + "smoke_train": bool(smoke_train), + "vocab": vocab, + "response_bank": response_bank, + "class_info": {}, + "intent_labels": list(OMNI_INTENTS_V2), + "domain_labels": list(OMNI_DOMAIN_LABELS_V2), + "max_len": max_len, + "image_size": int(image_size), + "embed_dim": 140, + "text_hidden": int(V41_MODEL_CONFIG["text_hidden"]), + "image_channels": 60, + "word_buckets": word_buckets, + "max_words": max_words, + "word_embed_dim": 128, + "deep_text_channels": 416, + "deep_image_channels": 136, + "fusion_hidden": int(V41_MODEL_CONFIG["fusion_hidden"]), + "memory_slots": 32, + "depth_steps": 12, + "expert_count": int(V41_MODEL_CONFIG["expert_count"]), + "expert_hidden": int(V41_MODEL_CONFIG["expert_hidden"]), + "context_top_k": 4, + "expert_top_k": int(V41_MODEL_CONFIG["expert_top_k"]), + "parameter_count": parameter_count, + "warm_start": warm_start, + "stage1": stage1, + "stage2": stage2, + "seed": int(seed), + "training_runtime": runtime.to_payload(), + "deliberation_passes": int(V41_MODEL_CONFIG["deliberation_passes"]), + "minimum_passes": int(V41_MODEL_CONFIG["minimum_passes"]), + "grounding_threshold": 0.58, + "prompt_understanding_mode": "hidden_plan_reflective_route_first_communication_polish_code_repair_v41", + "notes": [ + "v41 expands v8 with hidden planning, reflection-style prompt variants, communication-polish supervision, and stronger coding-repair emphasis.", + "The training scaffold keeps the frozen v8 base rows and adds higher-density v41 rows for disagreement resolution, grounded uncertainty, reasoning budgets, and human-facing rewrites.", + ], + } + _write_json_atomic_v41(meta_path, meta) + sample_outputs: List[Dict[str, str]] = [] + if not smoke_train: + engine = OmniCollectiveEngineV41(weights_path=weights_path, meta_path=meta_path, device=runtime.device) + sample_prompts = [ + "Reply in exactly two sentences explaining why regression tests matter.", + "Rewrite this answer so it is direct and professional: The present implementation exhibits a non-trivial probability of behavioral regression under future edits.", + "Two strong local teachers disagree. Choose the best route before answering: Which local model is best for benchmark-focused reasoning prompts?", + "A checkpoint save failed with Access is denied on Windows. What should happen next?", + "Write a vivid but grounded paragraph about a storm-battered lighthouse at night.", + ] + sample_outputs = [{"prompt": prompt, "answer": engine.answer(prompt)} for prompt in sample_prompts] + summary = { + "artifact": zip_path.name, + "family": "omni_collective_v41", + "smoke_train": bool(smoke_train), + "parameter_count": parameter_count, + "dataset_summary": dataset_summary, + "warm_start": warm_start, + "stage1": stage1, + "stage2": stage2, + "training_runtime": runtime.to_payload(), + "sample_outputs": sample_outputs, + "notes": [ + "v41 is tuned to be more direct with humans, stronger at route-then-answer behavior, and more repair-oriented on coding and debugging tasks.", + "This scaffold is the smoke-train entry point for future full v41 runs on top of the finished v8 base artifact.", + "Smoke packaging skips post-train sample inference so quick validation runs do not stall on CPU-only environments.", + ], + } + _write_json_atomic_v41(summary_out_path, summary) + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as archive: + archive.write(weights_path, arcname=weights_path.name) + archive.write(meta_path, arcname=meta_path.name) + archive.write(summary_out_path, arcname=summary_out_path.name) + if not smoke_train: + models_dir.mkdir(parents=True, exist_ok=True) + shutil.copy2(zip_path, desktop_zip_path) + _write_run_state_v41( + run_state_path, + { + "status": "complete", + "stage": "done", + "resume_dir": str(stage_resume_dir), + "zip_path": str(zip_path), + "desktop_zip_path": str(desktop_zip_path) if not smoke_train else None, + "artifact_dir": str(artifact_dir), + "parameter_count": parameter_count, + "stage1_best_score": float(stage1["best_score"]), + "stage2_best_score": float(stage2["best_score"]), + "runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + }, + ) + return { + "zip_path": str(zip_path), + "desktop_zip_path": str(desktop_zip_path) if not smoke_train else None, + "artifact_dir": str(artifact_dir), + "parameter_count": parameter_count, + "stage1_val": stage1["val_metrics"], + "stage2_val": stage2["val_metrics"], + "warm_start": warm_start, + "dataset_summary": dataset_summary, + "training_runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Prepare and smoke-train the omni_collective_v41 continuation scaffold." + ) + parser.add_argument("--summary", default="", help="Optional v8 frontier summary JSON. Defaults to the latest local v8 summary.") + parser.add_argument("--output_root", default=str(DEFAULT_OUTPUT_ROOT), help="Directory for the generated v41 prep pack.") + parser.add_argument("--output_dir", default=str(DEFAULT_TRAIN_OUTPUT_DIR), help="Directory for v41 train artifacts.") + parser.add_argument("--models_dir", default=str(DEFAULT_MODELS_DIR), help="Models directory for dry-run dataset assembly and optional publish copy.") + parser.add_argument("--base_zip", default=str(DEFAULT_BASE_ZIP), help="Base v8 artifact zip used for warm start.") + parser.add_argument("--images_dir", default="output/science_vision_dataset/images", help="Image directory for dry-run dataset assembly.") + parser.add_argument("--seed", type=int, default=41) + parser.add_argument("--communication_limit", type=int, default=180) + parser.add_argument("--disagreement_limit", type=int, default=120) + parser.add_argument("--base_distill_limit", type=int, default=0) + parser.add_argument("--base_teacher_model_limit", type=int, default=0) + parser.add_argument("--image_size", type=int, default=96) + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--stage1_epochs", type=int, default=1) + parser.add_argument("--stage2_epochs", type=int, default=1) + parser.add_argument("--stage1_lr", type=float, default=0.00022) + parser.add_argument("--stage2_lr", type=float, default=0.00010) + parser.add_argument("--device", default="auto") + parser.add_argument("--amp", default="off") + parser.add_argument("--amp_dtype", default="auto") + parser.add_argument("--compile_model", action="store_true") + parser.add_argument("--compile_mode", default="reduce-overhead") + parser.add_argument("--grad_accum_steps", type=int, default=1) + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--warmup_steps", type=int, default=0) + parser.add_argument("--warmup_ratio", type=float, default=0.05) + parser.add_argument("--min_lr_scale", type=float, default=0.05) + parser.add_argument("--dry_run_dataset", action="store_true", help="Build and summarize the merged v41 dry-run dataset on top of the frozen v8 base.") + parser.add_argument("--train_smoke", action="store_true", help="Run a tiny resumable v41 smoke train on top of the finished v8 base artifact.") + parser.add_argument("--train_frontier", action="store_true", help="Run the full resumable v41 frontier training job.") + parser.add_argument("--stdout", action="store_true", help="Print the prep or smoke summary JSON.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + summary_path = Path(args.summary).resolve() if str(args.summary).strip() else latest_v8_summary_path(REPO_ROOT) + payload = build_v41_prep_pack( + summary_path=summary_path, + output_root=Path(args.output_root), + seed=int(args.seed), + communication_limit=int(args.communication_limit), + disagreement_limit=int(args.disagreement_limit), + ) + if args.dry_run_dataset: + dry_run_summary = build_training_rows_v41_dry_run( + summary_path=summary_path, + output_root=Path(args.output_root), + seed=int(args.seed), + communication_limit=int(args.communication_limit), + disagreement_limit=int(args.disagreement_limit), + ) + payload = payload | { + "dry_run_dataset": dry_run_summary + } + if args.train_smoke: + smoke_result = train_model_v41( + repo_root=REPO_ROOT, + output_dir=Path(args.output_dir).resolve(), + models_dir=Path(args.models_dir).resolve(), + base_zip=Path(args.base_zip).resolve(), + images_dir=Path(args.images_dir).resolve(), + summary_path=summary_path, + prep_output_root=Path(args.output_root).resolve(), + image_size=int(args.image_size), + batch_size=int(args.batch_size), + stage1_epochs=int(args.stage1_epochs), + stage2_epochs=int(args.stage2_epochs), + stage1_lr=float(args.stage1_lr), + stage2_lr=float(args.stage2_lr), + seed=int(args.seed), + communication_limit=int(args.communication_limit), + disagreement_limit=int(args.disagreement_limit), + base_distill_limit=int(args.base_distill_limit), + base_teacher_model_limit=int(args.base_teacher_model_limit), + requested_device=str(args.device), + amp_mode=str(args.amp), + amp_dtype=str(args.amp_dtype), + compile_model=bool(args.compile_model), + compile_mode=str(args.compile_mode), + grad_accum_steps=int(args.grad_accum_steps), + ema_decay=float(args.ema_decay), + warmup_steps=int(args.warmup_steps), + warmup_ratio=float(args.warmup_ratio), + min_lr_scale=float(args.min_lr_scale), + smoke_train=True, + ) + payload = payload | {"smoke_train": smoke_result} + if args.train_frontier: + frontier_result = train_model_v41( + repo_root=REPO_ROOT, + output_dir=Path(args.output_dir).resolve(), + models_dir=Path(args.models_dir).resolve(), + base_zip=Path(args.base_zip).resolve(), + images_dir=Path(args.images_dir).resolve(), + summary_path=summary_path, + prep_output_root=Path(args.output_root).resolve(), + image_size=int(args.image_size), + batch_size=int(args.batch_size), + stage1_epochs=int(args.stage1_epochs), + stage2_epochs=int(args.stage2_epochs), + stage1_lr=float(args.stage1_lr), + stage2_lr=float(args.stage2_lr), + seed=int(args.seed), + communication_limit=int(args.communication_limit), + disagreement_limit=int(args.disagreement_limit), + base_distill_limit=int(args.base_distill_limit), + base_teacher_model_limit=int(args.base_teacher_model_limit), + requested_device=str(args.device), + amp_mode=str(args.amp), + amp_dtype=str(args.amp_dtype), + compile_model=bool(args.compile_model), + compile_mode=str(args.compile_mode), + grad_accum_steps=int(args.grad_accum_steps), + ema_decay=float(args.ema_decay), + warmup_steps=int(args.warmup_steps), + warmup_ratio=float(args.warmup_ratio), + min_lr_scale=float(args.min_lr_scale), + smoke_train=False, + ) + payload = payload | {"train_frontier": frontier_result} + if args.stdout: + print(json.dumps(payload, indent=2, ensure_ascii=True)) + + +if __name__ == "__main__": + main() diff --git a/source/train_omni_collective_v42.py b/source/train_omni_collective_v42.py new file mode 100644 index 000000000..ef1510f1d --- /dev/null +++ b/source/train_omni_collective_v42.py @@ -0,0 +1,1201 @@ +from __future__ import annotations + +import argparse +import json +import random +import shutil +import zipfile +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Sequence, Tuple + +import torch + +SOURCE_DIR = Path(__file__).resolve().parent +REPO_ROOT = SOURCE_DIR.parent +if str(SOURCE_DIR) not in __import__("sys").path: + __import__("sys").path.append(str(SOURCE_DIR)) + +try: + from image_recognition_model import SCIENCE_IMAGE_CLASSES + from omni_collective_model import OMNI_DOMAIN_LABELS_V2, OMNI_INTENTS_V2, build_char_vocab + from omni_collective_v42_model import OmniCollectiveEngineV42, OmniCollectiveNetV42 + from omni_training_runtime import maybe_compile_model, resolve_training_runtime + from prepare_omni_collective_v42 import build_v42_blueprint, latest_v41_summary_path, latest_v41_zip_path, load_summary + from train_omni_collective_v2 import DEFAULT_MODELS_DIR, OmniRow, split_rows + from train_omni_collective_v4 import _load_expanded_state_from_zip + from train_omni_collective_v8 import _train_stage_resumable_v8 + from train_omni_collective_v41 import ( + _bundle_rows, + _code_critique_repair_rows_v41, + _communication_polish_rows_v41, + _dedupe_rows, + _eval_payload, + _latent_plan_rows_v41, + _promotion_eval_pack_v41, + _read_jsonl_rows, + _reasoning_budget_rows_v41, + _row, + _seeded_take, + _teacher_disagreement_rows_v41, + _write_json, + _write_json_atomic_v41, + _write_jsonl, + build_training_rows_v41, + build_training_rows_v41_dry_run, + ) +except ImportError: # pragma: no cover + from .image_recognition_model import SCIENCE_IMAGE_CLASSES + from .omni_collective_model import OMNI_DOMAIN_LABELS_V2, OMNI_INTENTS_V2, build_char_vocab + from .omni_collective_v42_model import OmniCollectiveEngineV42, OmniCollectiveNetV42 + from .omni_training_runtime import maybe_compile_model, resolve_training_runtime + from .prepare_omni_collective_v42 import build_v42_blueprint, latest_v41_summary_path, latest_v41_zip_path, load_summary + from .train_omni_collective_v2 import DEFAULT_MODELS_DIR, OmniRow, split_rows + from .train_omni_collective_v4 import _load_expanded_state_from_zip + from .train_omni_collective_v8 import _train_stage_resumable_v8 + from .train_omni_collective_v41 import ( + _bundle_rows, + _code_critique_repair_rows_v41, + _communication_polish_rows_v41, + _dedupe_rows, + _eval_payload, + _latent_plan_rows_v41, + _promotion_eval_pack_v41, + _read_jsonl_rows, + _reasoning_budget_rows_v41, + _row, + _seeded_take, + _teacher_disagreement_rows_v41, + _write_json, + _write_json_atomic_v41, + _write_jsonl, + build_training_rows_v41, + build_training_rows_v41_dry_run, + ) + + +DEFAULT_OUTPUT_ROOT = REPO_ROOT / "output" / "omni_collective_v42_prep" +DEFAULT_TRAIN_OUTPUT_DIR = REPO_ROOT / "output" +V42_MODEL_CONFIG: Dict[str, Any] = { + "family": "omni_collective_v42", + "base_model_family": "omni_collective_v41", + "text_hidden": 304, + "fusion_hidden": 1216, + "expert_count": 14, + "expert_hidden": 2048, + "expert_top_k": 2, + "deliberation_passes": 14, + "minimum_passes": 7, + "new_heads": [ + "budget_router_head", + "teacher_consensus_head", + "cache_budget_head", + "verifier_gate_head", + ], +} + + +def _cleanup_smoke_checkpoint_temps(stage_resume_dir: Path) -> None: + for name in ("stage1_progress.pt.tmp", "stage2_progress.pt.tmp"): + path = stage_resume_dir / name + if path.exists(): + path.unlink() + + +def _cleanup_completed_smoke_stage(stage_resume_dir: Path, stage_name: str) -> None: + progress_path = stage_resume_dir / f"{stage_name}_progress.pt" + temp_path = stage_resume_dir / f"{stage_name}_progress.pt.tmp" + if progress_path.exists(): + progress_path.unlink() + if temp_path.exists(): + temp_path.unlink() + + +def _benchmark_bridge_rows_v42(*, seed: int, limit: int = 220) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + + route_rows = [ + ( + "Route first, then answer.\n" + "Request: Which local model should handle a common benchmark reasoning prompt if I care about exact score first?\n" + "Return the final answer only.", + "Use v40_benchmax first for benchmark-heavy exact reasoning, then let the broader omni model handle follow-up explanation or multimodal context.", + "model_selection", + "benchmark_bridge_v42::route_selection", + ), + ( + "Answer like a benchmark-focused reviewer.\n" + "Question: Why should v42 inherit v40 hard-example replay instead of relying only on v41 communication polish?", + "Because v41 improved presentation and routing, but v40 still wins on raw benchmark exactness. Hard-example replay keeps the failure patterns that matter for benchmark tasks instead of training only on better-looking answers.", + "reasoning", + "benchmark_bridge_v42::bridge_rationale", + ), + ( + "Two local lines disagree.\n" + "Question: v40_benchmax gives a terse exact answer, while v41 gives a broader answer with extra explanation. What should v42 learn?", + "v42 should keep the exact benchmark answer from the benchmark specialist, then add only the minimum extra explanation that does not change correctness.", + "reasoning", + "benchmark_bridge_v42::disagreement_resolution", + ), + ] + for prompt, response, domain, source in route_rows: + rows.append(_row(prompt, response, intent="general", domain=domain, source=source)) + + benchmark_buckets = [ + ("BoolQ", "read the question literally, then answer yes or no with the shortest justified wording"), + ("PIQA", "prefer concrete physical plausibility over surface similarity"), + ("HellaSwag", "finish the scenario by matching causal and stylistic continuity"), + ("MMLU", "identify the discipline first, then eliminate options before answering"), + ("ARC-Challenge", "favor grounded science reasoning and explicit elimination"), + ("GSM8K", "compute the answer directly and verify the arithmetic before finalizing"), + ] + for benchmark_name, strategy in benchmark_buckets: + prompt = ( + "Give a compact training answer for benchmark behavior.\n" + f"Question: How should v42 approach {benchmark_name} style questions?\n" + "Return the final answer only." + ) + response = f"For {benchmark_name} tasks, {strategy}." + rows.append(_row(prompt, response, intent="comparison", domain="reasoning", source=f"benchmark_bridge_v42::{benchmark_name.lower()}")) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _teacher_role_rows_v42(*, seed: int, limit: int = 260) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + teacher_rows = [ + ( + "Which teacher should supervise a grounded image-and-text explanation task?", + "Prefer google/gemma-4-31B-it for grounded multimodal explanation slices, then keep only answers that stay concrete and evidence-aware.", + "model_selection", + "teacher_roles_v42::gemma4_grounded", + ), + ( + "Which teacher should supervise a long-context reasoning and tool-use example?", + "Prefer Qwen/Qwen3.5-397B-A17B because it is the strongest current teacher in this mix for long-context reasoning, tool use, and route-aware problem solving.", + "model_selection", + "teacher_roles_v42::qwen35_long_context", + ), + ( + "Which teacher should supervise a failing-test to patch-repair example?", + "Prefer Qwen/Qwen3-Coder-Next for patch, test-repair, and repo-grounded coding trajectories, then verify the repair before keeping it.", + "model_selection", + "teacher_roles_v42::coder_next", + ), + ( + "Which teacher should supervise audio, video, or OCR-style route examples?", + "Prefer Qwen/Qwen3-Omni-30B-A3B-Instruct for audio-video and cross-modal route supervision, then distill only the clean grounded answer form into v42.", + "model_selection", + "teacher_roles_v42::qwen_omni", + ), + ] + for prompt, response, domain, source in teacher_rows: + rows.append(_row(prompt, response, intent="model_selection", domain=domain, source=source)) + + use_cases = [ + ("A user uploads an image and wants a grounded description with uncertainty when evidence is weak.", "google/gemma-4-31B-it", "grounded multimodal"), + ("A user pastes a huge log file and asks for the minimal root-cause summary plus next action.", "Qwen/Qwen3.5-397B-A17B", "long-context reasoning"), + ("A user has a failing test and wants a concrete patch plus why the failure happened.", "Qwen/Qwen3-Coder-Next", "agentic coding"), + ("A user asks which local specialist should handle an audio transcription plus answer step.", "Qwen/Qwen3-Omni-30B-A3B-Instruct", "omni routing"), + ] + for request, teacher, focus in use_cases: + prompt = ( + "Choose the teacher and explain the training reason.\n" + f"Request: {request}\n" + "Return the final answer only." + ) + response = f"Use {teacher} here because the slice is mainly about {focus}, not generic chat style." + rows.append(_row(prompt, response, intent="comparison", domain="model_selection", source="teacher_roles_v42::teacher_choice")) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _verifier_repair_rows_v42(*, seed: int, limit: int = 220) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + repair_examples = [ + ( + "Repair the answer after a verifier note.\n" + "Question: Which local model should handle benchmark-focused reasoning prompts?\n" + "Draft: omni_collective_v41 is always the best choice because it is the newest omni model.\n" + "Verifier note: The draft ignores that v40_benchmax currently scores higher on the common benchmark suite.", + "Use v40_benchmax first for benchmark-focused reasoning because it currently outperforms omni_collective_v41 on the saved common benchmark comparison.", + "teacher_verifier_v42::benchmark_repair", + ), + ( + "Repair the answer after a verifier note.\n" + "Question: Should a long transcript always get the deepest reasoning budget?\n" + "Draft: Yes, use the deepest budget for every long transcript.\n" + "Verifier note: The draft overgeneralizes and wastes compute on easy summarization tasks.", + "No. Use a deeper budget only when the transcript also requires hard reasoning or multi-step verification. Straight summarization should keep a medium budget.", + "teacher_verifier_v42::budget_repair", + ), + ( + "Repair this coding answer after a verifier note.\n" + "Question: How should a model respond to a failing test after a code change?\n" + "Draft: Rewrite the whole file and try again.\n" + "Verifier note: The draft is too destructive and ignores the failing assertion.", + "Start from the failing assertion, identify the smallest behavior change that caused it, patch that path, and rerun the targeted test before broader edits.", + "teacher_verifier_v42::coding_repair", + ), + ( + "Repair this multimodal answer after a verifier note.\n" + "Question: The image evidence is weak. How should the answer handle that?\n" + "Draft: State the most likely object as a fact.\n" + "Verifier note: The draft hides uncertainty.", + "Give the most likely reading, state that confidence is limited, and avoid inventing details the image does not support.", + "teacher_verifier_v42::multimodal_repair", + ), + ] + for prompt, response, source in repair_examples: + rows.append(_row(prompt, response, intent="comparison", domain="reasoning", source=source)) + + verifier_styles = [ + ("incorrect route choice", "switch to the best local specialist before answering"), + ("weak first reasoning step", "restart from the cleanest supported first step"), + ("too much filler", "compress the answer without dropping the evidence"), + ("unsafe uncertainty handling", "state what is missing and stop guessing"), + ] + for problem, remedy in verifier_styles: + prompt = ( + "Write the repaired answer after verifier feedback.\n" + f"Verifier issue: {problem}.\n" + "Return the final answer only." + ) + response = f"Repair the response by applying the smallest fix that solves the {problem}, then {remedy}." + rows.append(_row(prompt, response, intent="comparison", domain="general", source="teacher_verifier_v42::repair_rule")) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _turboquant_budget_rows_v42(*, seed: int, limit: int = 180) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + budget_examples = [ + ( + "Choose the reasoning and evidence budget.\n" + "Request: Summarize a very long build log into the one failing step and the next fix.\n" + "Return the final answer only.", + "Use a medium reasoning budget and a compressed evidence set. Keep only the failing step, the relevant stack trace line, and the next concrete action.", + "turboquant_budget_v42::compressed_log_summary", + ), + ( + "Choose the reasoning and evidence budget.\n" + "Request: Solve a GSM8K-style arithmetic problem.\n" + "Return the final answer only.", + "Use a deep reasoning budget for the calculation itself, but keep the final answer concise and verify the arithmetic before stopping.", + "turboquant_budget_v42::math_budget", + ), + ( + "Choose the reasoning and evidence budget.\n" + "Request: Compare two model options for a coding task.\n" + "Return the final answer only.", + "Use a medium budget: identify the task, compare the main tradeoff, and recommend one model without overexplaining.", + "turboquant_budget_v42::model_compare_budget", + ), + ( + "Answer with context-economy discipline.\n" + "Question: What should v42 learn from TurboQuant without pretending to train quantization itself?", + "v42 should learn evidence selection and budget control. TurboQuant mainly makes long teacher and verifier runs cheaper; the trainable behavior is when to use more or less context, not how to quantize weights.", + "turboquant_budget_v42::training_translation", + ), + ] + for prompt, response, source in budget_examples: + rows.append(_row(prompt, response, intent="general", domain="reasoning", source=source)) + + contexts = [ + ("a giant transcript where only three decisions matter", "shortlist the three decisions and ignore the rest"), + ("a multi-file debugging session with one repeated failure", "track the repeated failure path and drop unrelated logs"), + ("a benchmark question with four plausible options", "keep the elimination evidence for the surviving option only"), + ("a multimodal request with ambiguous image evidence", "keep the visible cues and explicitly note the uncertainty"), + ] + for context, rule in contexts: + prompt = ( + "Write the evidence-budget rule.\n" + f"Situation: {context}.\n" + "Return the final answer only." + ) + response = f"Keep only the evidence that changes the answer: {rule}." + rows.append(_row(prompt, response, intent="comparison", domain="reasoning", source="turboquant_budget_v42::evidence_rule")) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _diversity_rows_v42(*, seed: int, limit: int = 260) -> Tuple[List[OmniRow], Dict[str, int]]: + rows: List[OmniRow] = [] + personas = [ + ("a warehouse operations lead", "planning", "The scanner rollout failed on one shift. What should the team do next?"), + ("a clinical lab analyst", "knowledge", "How do I explain a low-confidence image result without overstating it?"), + ("a civil engineer", "reasoning", "Compare two material options and recommend one for durability."), + ("a game developer", "coding", "A recent refactor broke input handling. What is the safest fix path?"), + ("a teacher", "language", "Rewrite this explanation so a teenager can follow it without losing the science."), + ("a founder", "comparison", "Which local model should I use for a product support workflow with screenshots?"), + ("a researcher", "knowledge", "Summarize this new model family shift in one direct paragraph."), + ("a 3D designer", "spatial_3d", "Which local specialist should handle an OpenSCAD or geometry prompt?"), + ] + styles = [ + ("keep it direct and practical", "Focus on the next concrete step and avoid vague filler."), + ("be calm and transparent about uncertainty", "State what is known, what is unclear, and the safest next move."), + ("compare options and recommend one", "Name the tradeoff, then make the recommendation explicit."), + ] + for persona, domain, question in personas: + for style, guidance in styles: + prompt = ( + f"Answer for {persona}.\n" + f"Question: {question}\n" + f"Style requirement: {style}." + ) + response = f"{guidance} Tailor the wording to {persona} while staying concise." + rows.append(_row(prompt, response, intent="general", domain=domain, source="diversity_mix_v42::persona_style")) + + selected = _seeded_take(rows, seed=seed, limit=limit) + counts: Dict[str, int] = {} + for row in selected: + counts[row.source] = counts.get(row.source, 0) + 1 + return selected, dict(sorted(counts.items())) + + +def _promotion_eval_pack_v42() -> List[Dict[str, str]]: + evals = list(_promotion_eval_pack_v41()) + evals.extend( + [ + _eval_payload( + "Which local model should handle a benchmark-style reasoning prompt if exact score is the main goal?", + expected="v40_benchmax", + focus="benchmark_bridge", + metric="contains", + source="promotion_eval_v42::benchmark_route", + ), + _eval_payload( + "Translate TurboQuant into a v42 training implication in one sentence.", + expected="budget control and evidence compression for longer teacher traces, not quantization as a direct training target", + focus="budget_reasoning", + metric="contains", + source="promotion_eval_v42::turboquant_translation", + ), + _eval_payload( + "Which teacher is the best fit for failing-test and patch-repair traces?", + expected="Qwen/Qwen3-Coder-Next", + focus="teacher_roles", + metric="contains", + source="promotion_eval_v42::coder_teacher", + ), + _eval_payload( + "A large transcript only needs the one failing step and the next action. Which reasoning budget should v42 prefer?", + expected="medium", + focus="budget_reasoning", + metric="contains", + source="promotion_eval_v42::budget_choice", + ), + ] + ) + return evals + + +def assemble_v42_training_rows( + *, + base_stage1_rows: Sequence[OmniRow], + base_full_rows: Sequence[OmniRow], + base_summary: Dict[str, Any], + prep_root: Path, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + prep_root = prep_root.resolve() + benchmark_rows = _read_jsonl_rows(prep_root / "v42_benchmark_bridge_rows.jsonl") + teacher_rows = _read_jsonl_rows(prep_root / "v42_teacher_role_rows.jsonl") + verifier_rows = _read_jsonl_rows(prep_root / "v42_verifier_repair_rows.jsonl") + budget_rows = _read_jsonl_rows(prep_root / "v42_turboquant_budget_rows.jsonl") + diversity_rows = _read_jsonl_rows(prep_root / "v42_diversity_rows.jsonl") + + added_rows = list(benchmark_rows) + list(teacher_rows) + list(verifier_rows) + list(budget_rows) + list(diversity_rows) + pre_dedupe_rows = len(base_full_rows) + len(added_rows) + full_rows = _dedupe_rows(list(base_full_rows) + added_rows) + stage1_rows = [row for row in full_rows if row.domain not in {"vision", "spatial_3d", "video"}] + + summary = { + "stage1_rows": len(stage1_rows), + "stage2_rows": len(full_rows), + "pre_dedupe_rows": pre_dedupe_rows, + "source_counts": _bundle_rows(full_rows)["source_counts"], + "base_stage1_rows": len(base_stage1_rows), + "base_stage2_rows": len(base_full_rows), + "v42_added_rows": len(added_rows), + "v42_prep_root": str(prep_root), + "v42_row_groups": { + "benchmark_bridge": _bundle_rows(benchmark_rows), + "teacher_roles": _bundle_rows(teacher_rows), + "verifier_repair": _bundle_rows(verifier_rows), + "turboquant_budget": _bundle_rows(budget_rows), + "diversity_mix": _bundle_rows(diversity_rows), + }, + "base_summary": base_summary, + } + return stage1_rows, full_rows, summary + + +def build_training_rows_v42( + *, + repo_root: Path, + models_dir: Path, + images_dir: Path, + summary_path: Path, + output_root: Path, + seed: int = 42, + benchmark_limit: int = 220, + teacher_route_limit: int = 260, + verifier_limit: int = 220, + budget_limit: int = 180, + diversity_limit: int = 260, + base_communication_limit: int = 220, + base_disagreement_limit: int = 140, + base_distill_limit: int = 0, + base_teacher_model_limit: int = 0, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + prep_payload = build_v42_prep_pack( + summary_path=summary_path, + output_root=output_root, + seed=seed, + benchmark_limit=benchmark_limit, + teacher_route_limit=teacher_route_limit, + verifier_limit=verifier_limit, + budget_limit=budget_limit, + diversity_limit=diversity_limit, + ) + base_stage1_rows, base_full_rows, base_summary = build_training_rows_v41( + repo_root=repo_root, + models_dir=models_dir, + images_dir=images_dir, + summary_path=summary_path, + output_root=Path(prep_payload["v41_base_output_root"]), + seed=seed, + communication_limit=base_communication_limit, + disagreement_limit=base_disagreement_limit, + base_distill_limit=base_distill_limit, + base_teacher_model_limit=base_teacher_model_limit, + ) + stage1_rows, full_rows, merged_summary = assemble_v42_training_rows( + base_stage1_rows=base_stage1_rows, + base_full_rows=base_full_rows, + base_summary=base_summary, + prep_root=Path(prep_payload["output_root"]), + ) + merged_summary["prep_payload"] = prep_payload + return stage1_rows, full_rows, merged_summary + + +def build_training_rows_v42_dry_run( + *, + summary_path: Path, + output_root: Path, + seed: int = 42, + benchmark_limit: int = 220, + teacher_route_limit: int = 260, + verifier_limit: int = 220, + budget_limit: int = 180, + diversity_limit: int = 260, + base_communication_limit: int = 220, + base_disagreement_limit: int = 140, +) -> Dict[str, Any]: + prep_payload = build_v42_prep_pack( + summary_path=summary_path, + output_root=output_root, + seed=seed, + benchmark_limit=benchmark_limit, + teacher_route_limit=teacher_route_limit, + verifier_limit=verifier_limit, + budget_limit=budget_limit, + diversity_limit=diversity_limit, + ) + base_dry_run = build_training_rows_v41_dry_run( + summary_path=summary_path, + output_root=Path(prep_payload["v41_base_output_root"]), + seed=seed, + communication_limit=base_communication_limit, + disagreement_limit=base_disagreement_limit, + ) + prep_root = Path(prep_payload["output_root"]).resolve() + benchmark_rows = _read_jsonl_rows(prep_root / "v42_benchmark_bridge_rows.jsonl") + teacher_rows = _read_jsonl_rows(prep_root / "v42_teacher_role_rows.jsonl") + verifier_rows = _read_jsonl_rows(prep_root / "v42_verifier_repair_rows.jsonl") + budget_rows = _read_jsonl_rows(prep_root / "v42_turboquant_budget_rows.jsonl") + diversity_rows = _read_jsonl_rows(prep_root / "v42_diversity_rows.jsonl") + + added_rows = list(benchmark_rows) + list(teacher_rows) + list(verifier_rows) + list(budget_rows) + list(diversity_rows) + added_stage1_rows = [row for row in added_rows if row.domain not in {"vision", "spatial_3d", "video"}] + source_counts = dict(base_dry_run.get("source_counts") or {}) + for row in added_rows: + source_counts[row.source] = source_counts.get(row.source, 0) + 1 + + dry_run_summary = { + "dry_run": True, + "base_mode": "frozen_v41_plus_v42_rows", + "source_summary": str(summary_path.resolve()), + "prep_payload": prep_payload, + "base_stage1_rows": int(base_dry_run.get("estimated_stage1_rows") or 0), + "base_stage2_rows": int(base_dry_run.get("estimated_stage2_rows") or 0), + "added_stage1_rows": len(added_stage1_rows), + "added_stage2_rows": len(added_rows), + "estimated_stage1_rows": int(base_dry_run.get("estimated_stage1_rows") or 0) + len(added_stage1_rows), + "estimated_stage2_rows": int(base_dry_run.get("estimated_stage2_rows") or 0) + len(added_rows), + "source_counts": dict(sorted(source_counts.items())), + "v42_row_groups": { + "benchmark_bridge": _bundle_rows(benchmark_rows), + "teacher_roles": _bundle_rows(teacher_rows), + "verifier_repair": _bundle_rows(verifier_rows), + "turboquant_budget": _bundle_rows(budget_rows), + "diversity_mix": _bundle_rows(diversity_rows), + }, + "v41_base_dry_run": base_dry_run, + } + dry_run_path = Path(output_root).resolve() / "omni_collective_v42_dry_run_summary.json" + _write_json(dry_run_path, dry_run_summary) + return dry_run_summary | {"summary_path": str(dry_run_path)} + + +def build_v42_prep_pack( + *, + summary_path: Path, + output_root: Path, + seed: int = 42, + benchmark_limit: int = 220, + teacher_route_limit: int = 260, + verifier_limit: int = 220, + budget_limit: int = 180, + diversity_limit: int = 260, +) -> Dict[str, Any]: + summary = load_summary(summary_path) + blueprint = build_v42_blueprint(summary, summary_path=summary_path) + benchmark_rows, benchmark_summary = _benchmark_bridge_rows_v42(seed=seed + 11, limit=benchmark_limit) + teacher_rows, teacher_summary = _teacher_role_rows_v42(seed=seed + 17, limit=teacher_route_limit) + verifier_rows, verifier_summary = _verifier_repair_rows_v42(seed=seed + 23, limit=verifier_limit) + budget_rows, budget_summary = _turboquant_budget_rows_v42(seed=seed + 29, limit=budget_limit) + diversity_rows, diversity_summary = _diversity_rows_v42(seed=seed + 31, limit=diversity_limit) + promotion_eval_pack = _promotion_eval_pack_v42() + + output_root = output_root.resolve() + output_root.mkdir(parents=True, exist_ok=True) + blueprint_path = output_root / "omni_collective_v42_blueprint.json" + benchmark_path = output_root / "v42_benchmark_bridge_rows.jsonl" + teacher_path = output_root / "v42_teacher_role_rows.jsonl" + verifier_path = output_root / "v42_verifier_repair_rows.jsonl" + budget_path = output_root / "v42_turboquant_budget_rows.jsonl" + diversity_path = output_root / "v42_diversity_rows.jsonl" + promotion_eval_path = output_root / "v42_promotion_eval_pack.json" + prep_summary_path = output_root / "omni_collective_v42_prep_summary.json" + v41_base_output_root = (output_root / "v41_base").resolve() + + _write_json(blueprint_path, blueprint) + _write_jsonl(benchmark_path, benchmark_rows) + _write_jsonl(teacher_path, teacher_rows) + _write_jsonl(verifier_path, verifier_rows) + _write_jsonl(budget_path, budget_rows) + _write_jsonl(diversity_path, diversity_rows) + _write_json(promotion_eval_path, {"evaluations": promotion_eval_pack}) + + prep_summary = { + "family": "omni_collective_v42", + "prepared_from": str(summary_path.resolve()), + "output_root": str(output_root), + "v41_base_output_root": str(v41_base_output_root), + "row_groups": { + "benchmark_bridge": {"row_count": len(benchmark_rows), "source_counts": benchmark_summary, "path": str(benchmark_path)}, + "teacher_roles": {"row_count": len(teacher_rows), "source_counts": teacher_summary, "path": str(teacher_path)}, + "verifier_repair": {"row_count": len(verifier_rows), "source_counts": verifier_summary, "path": str(verifier_path)}, + "turboquant_budget": {"row_count": len(budget_rows), "source_counts": budget_summary, "path": str(budget_path)}, + "diversity_mix": {"row_count": len(diversity_rows), "source_counts": diversity_summary, "path": str(diversity_path)}, + }, + "promotion_eval_pack_path": str(promotion_eval_path), + "total_new_rows": len(benchmark_rows) + len(teacher_rows) + len(verifier_rows) + len(budget_rows) + len(diversity_rows), + "blueprint_path": str(blueprint_path), + } + _write_json(prep_summary_path, prep_summary) + return prep_summary | {"summary_path": str(prep_summary_path)} + + +def _run_state_path_v42(output_dir: Path) -> Path: + return output_dir / "omni_collective_v42_train_state.json" + + +def _stage_resume_dir_v42( + output_dir: Path, + *, + seed: int, + distill_limit: int, + teacher_model_limit: int, + smoke_train: bool, +) -> Path: + teacher_tag = int(teacher_model_limit) if int(teacher_model_limit) > 0 else 0 + mode_tag = "smoke" if smoke_train else "frontier" + return output_dir / "omni_v42_stage_resume" / f"{mode_tag}_seed_{int(seed)}_distill_{int(distill_limit)}_teacherlimit_{teacher_tag}" + + +def _write_run_state_v42(path: Path, payload: Dict[str, Any]) -> None: + cooked = dict(payload) + cooked["updated_at"] = datetime.now().isoformat() + _write_json_atomic_v41(path, cooked) + + +def _smoke_training_rows_v42( + *, + stage1_rows: Sequence[OmniRow], + full_rows: Sequence[OmniRow], + seed: int, +) -> Tuple[List[OmniRow], List[OmniRow], Dict[str, Any]]: + stage1_priority = [row for row in stage1_rows if "_v42::" in row.source or "_v42" in row.source] + full_priority = [row for row in full_rows if "_v42::" in row.source or "_v42" in row.source] + base_stage1 = [row for row in stage1_rows if row not in stage1_priority] + multimodal = [row for row in full_rows if row.domain in {"vision", "spatial_3d", "video"}] + + selected_stage1 = _dedupe_rows( + _seeded_take(stage1_priority, seed=seed + 1, limit=26) + + _seeded_take(base_stage1, seed=seed + 2, limit=22) + ) + selected_full = _dedupe_rows( + list(selected_stage1) + + _seeded_take(full_priority, seed=seed + 3, limit=14) + + _seeded_take(multimodal, seed=seed + 4, limit=10) + ) + selected_stage1 = [row for row in selected_full if row.domain not in {"vision", "spatial_3d", "video"}] + summary = { + "mode": "smoke", + "stage1_rows": len(selected_stage1), + "stage2_rows": len(selected_full), + "v42_priority_rows": len([row for row in selected_full if "_v42::" in row.source or "_v42" in row.source]), + "multimodal_rows": len([row for row in selected_full if row.domain in {"vision", "spatial_3d", "video"}]), + "source_counts": _bundle_rows(selected_full)["source_counts"], + } + return selected_stage1, selected_full, summary + + +def train_model_v42( + *, + repo_root: Path, + output_dir: Path, + models_dir: Path, + base_zip: Path, + images_dir: Path, + summary_path: Path, + prep_output_root: Path, + image_size: int, + batch_size: int, + stage1_epochs: int, + stage2_epochs: int, + stage1_lr: float, + stage2_lr: float, + seed: int, + benchmark_limit: int, + teacher_route_limit: int, + verifier_limit: int, + budget_limit: int, + diversity_limit: int, + base_communication_limit: int, + base_disagreement_limit: int, + base_distill_limit: int, + base_teacher_model_limit: int, + requested_device: str, + amp_mode: str, + amp_dtype: str, + compile_model: bool, + compile_mode: str, + grad_accum_steps: int, + ema_decay: float, + warmup_steps: int, + warmup_ratio: float, + min_lr_scale: float, + smoke_train: bool = False, +) -> Dict[str, Any]: + torch.manual_seed(int(seed)) + random.seed(int(seed)) + output_dir.mkdir(parents=True, exist_ok=True) + stage_resume_dir = _stage_resume_dir_v42( + output_dir, + seed=int(seed), + distill_limit=int(base_distill_limit), + teacher_model_limit=int(base_teacher_model_limit), + smoke_train=bool(smoke_train), + ) + stage_resume_dir.mkdir(parents=True, exist_ok=True) + run_state_path = _run_state_path_v42(output_dir) + _write_run_state_v42( + run_state_path, + { + "status": "building_dataset", + "stage": "dataset", + "resume_dir": str(stage_resume_dir), + "smoke_train": bool(smoke_train), + }, + ) + + stage1_rows, full_rows, dataset_summary = build_training_rows_v42( + repo_root=repo_root, + models_dir=models_dir, + images_dir=images_dir, + summary_path=summary_path, + output_root=prep_output_root, + seed=seed, + benchmark_limit=benchmark_limit, + teacher_route_limit=teacher_route_limit, + verifier_limit=verifier_limit, + budget_limit=budget_limit, + diversity_limit=diversity_limit, + base_communication_limit=base_communication_limit, + base_disagreement_limit=base_disagreement_limit, + base_distill_limit=base_distill_limit, + base_teacher_model_limit=base_teacher_model_limit, + ) + if smoke_train: + stage1_rows, full_rows, smoke_summary = _smoke_training_rows_v42( + stage1_rows=stage1_rows, + full_rows=full_rows, + seed=seed, + ) + dataset_summary = dict(dataset_summary) + dataset_summary["smoke_subset"] = smoke_summary + dataset_summary["stage1_rows"] = len(stage1_rows) + dataset_summary["stage2_rows"] = len(full_rows) + + print(json.dumps({"event": "dataset_built", "summary": dataset_summary}, ensure_ascii=True), flush=True) + _write_json_atomic_v41(stage_resume_dir / "dataset_summary.json", {"dataset_summary": dataset_summary}) + _write_run_state_v42( + run_state_path, + { + "status": "dataset_built", + "stage": "dataset", + "resume_dir": str(stage_resume_dir), + "dataset_summary": dataset_summary, + "smoke_train": bool(smoke_train), + }, + ) + + train_rows, val_rows = split_rows(full_rows, seed=seed + 17) + train_stage1 = [row for row in train_rows if row.domain not in {"vision", "spatial_3d", "video"}] + vocab = build_char_vocab([row.prompt for row in train_rows], min_frequency=1) + response_bank = sorted({row.response_text for row in full_rows if row.response_text}) + print( + json.dumps( + { + "event": "label_space", + "train_rows": len(train_rows), + "val_rows": len(val_rows), + "response_bank": len(response_bank), + "vocab_size": len(vocab), + "smoke_train": bool(smoke_train), + }, + ensure_ascii=True, + ), + flush=True, + ) + _write_run_state_v42( + run_state_path, + { + "status": "label_space_ready", + "stage": "label_space", + "resume_dir": str(stage_resume_dir), + "train_rows": len(train_rows), + "val_rows": len(val_rows), + "response_bank": len(response_bank), + "vocab_size": len(vocab), + "smoke_train": bool(smoke_train), + }, + ) + + max_len = 420 + max_words = 112 + word_buckets = 20480 + runtime = resolve_training_runtime( + repo_root=repo_root, + requested_device=requested_device, + amp_mode=amp_mode, + amp_dtype=amp_dtype, + compile_requested=compile_model, + compile_mode=compile_mode, + grad_accum_steps=grad_accum_steps, + ema_decay=ema_decay, + warmup_steps=warmup_steps, + warmup_ratio=warmup_ratio, + min_lr_scale=min_lr_scale, + batch_size=batch_size, + ) + print(json.dumps({"event": "runtime_config", "runtime": runtime.to_payload()}, ensure_ascii=True), flush=True) + model = OmniCollectiveNetV42( + vocab_size=max(len(vocab), 2), + num_intents=len(OMNI_INTENTS_V2), + num_responses=max(len(response_bank), 1), + num_vision_classes=len(SCIENCE_IMAGE_CLASSES), + num_domains=len(OMNI_DOMAIN_LABELS_V2), + base_embed_dim=144, + text_hidden=int(V42_MODEL_CONFIG["text_hidden"]), + image_channels=64, + word_buckets=word_buckets, + word_embed_dim=128, + deep_text_channels=448, + deep_image_channels=144, + fusion_hidden=int(V42_MODEL_CONFIG["fusion_hidden"]), + memory_slots=36, + depth_steps=14, + expert_count=int(V42_MODEL_CONFIG["expert_count"]), + expert_hidden=int(V42_MODEL_CONFIG["expert_hidden"]), + context_top_k=4, + expert_top_k=int(V42_MODEL_CONFIG["expert_top_k"]), + ).to(runtime.device) + warm_start = _load_expanded_state_from_zip(model, base_zip) + forward_model, runtime = maybe_compile_model(model, runtime) + print(json.dumps({"event": "warm_start", "info": warm_start}, ensure_ascii=True), flush=True) + _write_run_state_v42( + run_state_path, + { + "status": "warm_start_complete", + "stage": "warm_start", + "resume_dir": str(stage_resume_dir), + "warm_start": warm_start, + "runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + }, + ) + + progress_every = 1 if smoke_train else None + checkpoint_every = None if smoke_train else None + if smoke_train: + _cleanup_smoke_checkpoint_temps(stage_resume_dir) + stage1 = _train_stage_resumable_v8( + model=model, + forward_model=forward_model, + train_rows=train_stage1, + val_rows=val_rows, + vocab=vocab, + response_bank=response_bank, + image_size=image_size, + max_len=max_len, + max_words=max_words, + word_buckets=word_buckets, + batch_size=batch_size, + learning_rate=stage1_lr, + epochs=stage1_epochs, + seed=seed + 101, + runtime=runtime, + loss_weights={"intent": 0.56, "response": 1.12, "domain": 0.78, "vision": 0.60}, + balance_weight=0.043, + stage_name="stage1", + stage_dir=stage_resume_dir, + run_state_path=run_state_path, + progress_every_batches=progress_every, + checkpoint_every_batches=checkpoint_every, + grad_accum_steps=grad_accum_steps, + ) + if smoke_train: + _cleanup_completed_smoke_stage(stage_resume_dir, "stage1") + stage2 = _train_stage_resumable_v8( + model=model, + forward_model=forward_model, + train_rows=train_rows, + val_rows=val_rows, + vocab=vocab, + response_bank=response_bank, + image_size=image_size, + max_len=max_len, + max_words=max_words, + word_buckets=word_buckets, + batch_size=max(4, int(batch_size) // 2), + learning_rate=stage2_lr, + epochs=stage2_epochs, + seed=seed + 151, + runtime=runtime, + loss_weights={"intent": 0.50, "response": 1.15, "domain": 0.84, "vision": 1.10}, + balance_weight=0.060, + stage_name="stage2", + stage_dir=stage_resume_dir, + run_state_path=run_state_path, + progress_every_batches=progress_every, + checkpoint_every_batches=checkpoint_every, + grad_accum_steps=grad_accum_steps, + ) + if smoke_train: + _cleanup_completed_smoke_stage(stage_resume_dir, "stage2") + + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + mode_tag = "smoke" if smoke_train else "frontier" + artifact_dir = output_dir / f"supermix_omni_collective_v42_{mode_tag}_{stamp}" + artifact_dir.mkdir(parents=True, exist_ok=True) + weights_path = artifact_dir / f"omni_collective_v42_{mode_tag}.pth" + meta_path = artifact_dir / f"omni_collective_v42_{mode_tag}_meta.json" + summary_out_path = artifact_dir / f"omni_collective_v42_{mode_tag}_summary.json" + zip_path = output_dir / f"supermix_omni_collective_v42_{mode_tag}_{stamp}.zip" + desktop_zip_path = models_dir / zip_path.name + + torch.save(model.state_dict(), weights_path) + parameter_count = int(sum(parameter.numel() for parameter in model.parameters())) + meta = { + "architecture_version": 42, + "family": "omni_collective_v42", + "smoke_train": bool(smoke_train), + "vocab": vocab, + "response_bank": response_bank, + "class_info": {}, + "intent_labels": list(OMNI_INTENTS_V2), + "domain_labels": list(OMNI_DOMAIN_LABELS_V2), + "max_len": max_len, + "image_size": int(image_size), + "embed_dim": 144, + "text_hidden": int(V42_MODEL_CONFIG["text_hidden"]), + "image_channels": 64, + "word_buckets": word_buckets, + "max_words": max_words, + "word_embed_dim": 128, + "deep_text_channels": 448, + "deep_image_channels": 144, + "fusion_hidden": int(V42_MODEL_CONFIG["fusion_hidden"]), + "memory_slots": 36, + "depth_steps": 14, + "expert_count": int(V42_MODEL_CONFIG["expert_count"]), + "expert_hidden": int(V42_MODEL_CONFIG["expert_hidden"]), + "context_top_k": 4, + "expert_top_k": int(V42_MODEL_CONFIG["expert_top_k"]), + "parameter_count": parameter_count, + "warm_start": warm_start, + "stage1": stage1, + "stage2": stage2, + "seed": int(seed), + "training_runtime": runtime.to_payload(), + "deliberation_passes": int(V42_MODEL_CONFIG["deliberation_passes"]), + "minimum_passes": int(V42_MODEL_CONFIG["minimum_passes"]), + "grounding_threshold": 0.60, + "prompt_understanding_mode": "budgeted_route_first_verifier_teacher_specialized_v42", + "notes": [ + "v42 extends the v41 scaffold with benchmark-bridge replay, teacher-role specialization, verifier-repair supervision, and TurboQuant-inspired budget control.", + "The v42 data program keeps v41 carryover rows, then adds Gemma 4, Qwen 3.5, Qwen3-Coder-Next, and Qwen3-Omni inspired route and repair slices.", + ], + } + _write_json_atomic_v41(meta_path, meta) + sample_outputs: List[Dict[str, str]] = [] + if not smoke_train: + engine = OmniCollectiveEngineV42(weights_path=weights_path, meta_path=meta_path, device=runtime.device) + sample_prompts = [ + "Which local model should handle a benchmark-style reasoning prompt if exact score matters most?", + "Translate TurboQuant into a training implication for v42 in one sentence.", + "A failing test appeared after a refactor. What should happen next?", + "Which teacher is best for grounded multimodal explanation slices?", + "Summarize a huge build log into the one failing step and the next action.", + ] + sample_outputs = [{"prompt": prompt, "answer": engine.answer(prompt)} for prompt in sample_prompts] + summary = { + "artifact": zip_path.name, + "family": "omni_collective_v42", + "smoke_train": bool(smoke_train), + "parameter_count": parameter_count, + "dataset_summary": dataset_summary, + "warm_start": warm_start, + "stage1": stage1, + "stage2": stage2, + "training_runtime": runtime.to_payload(), + "sample_outputs": sample_outputs, + "notes": [ + "v42 is designed to recover benchmark strength from v40 while keeping v41's route-first and communication improvements.", + "This scaffold is the first v42 smoke/frontier training entry point on top of the latest v41 artifact.", + "Smoke packaging skips post-train sample inference so quick validation runs stay manageable on CPU-only environments.", + ], + } + _write_json_atomic_v41(summary_out_path, summary) + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as archive: + archive.write(weights_path, arcname=weights_path.name) + archive.write(meta_path, arcname=meta_path.name) + archive.write(summary_out_path, arcname=summary_out_path.name) + if not smoke_train: + models_dir.mkdir(parents=True, exist_ok=True) + shutil.copy2(zip_path, desktop_zip_path) + _write_run_state_v42( + run_state_path, + { + "status": "complete", + "stage": "done", + "resume_dir": str(stage_resume_dir), + "zip_path": str(zip_path), + "desktop_zip_path": str(desktop_zip_path) if not smoke_train else None, + "artifact_dir": str(artifact_dir), + "parameter_count": parameter_count, + "stage1_best_score": float(stage1["best_score"]), + "stage2_best_score": float(stage2["best_score"]), + "runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + }, + ) + return { + "zip_path": str(zip_path), + "desktop_zip_path": str(desktop_zip_path) if not smoke_train else None, + "artifact_dir": str(artifact_dir), + "parameter_count": parameter_count, + "stage1_val": stage1["val_metrics"], + "stage2_val": stage2["val_metrics"], + "warm_start": warm_start, + "dataset_summary": dataset_summary, + "training_runtime": runtime.to_payload(), + "smoke_train": bool(smoke_train), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Prepare and smoke-train the omni_collective_v42 continuation scaffold." + ) + parser.add_argument("--summary", default="", help="Optional v41 frontier summary JSON. Defaults to the latest local v41 summary.") + parser.add_argument("--output_root", default=str(DEFAULT_OUTPUT_ROOT), help="Directory for the generated v42 prep pack.") + parser.add_argument("--output_dir", default=str(DEFAULT_TRAIN_OUTPUT_DIR), help="Directory for v42 train artifacts.") + parser.add_argument("--models_dir", default=str(DEFAULT_MODELS_DIR), help="Models directory for dry-run dataset assembly and optional publish copy.") + parser.add_argument("--base_zip", default="", help="Base v41 artifact zip used for warm start. Defaults to the latest local v41 zip.") + parser.add_argument("--images_dir", default="output/science_vision_dataset/images", help="Image directory for dataset assembly.") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--benchmark_limit", type=int, default=220) + parser.add_argument("--teacher_route_limit", type=int, default=260) + parser.add_argument("--verifier_limit", type=int, default=220) + parser.add_argument("--budget_limit", type=int, default=180) + parser.add_argument("--diversity_limit", type=int, default=260) + parser.add_argument("--base_communication_limit", type=int, default=220) + parser.add_argument("--base_disagreement_limit", type=int, default=140) + parser.add_argument("--base_distill_limit", type=int, default=0) + parser.add_argument("--base_teacher_model_limit", type=int, default=0) + parser.add_argument("--image_size", type=int, default=96) + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--stage1_epochs", type=int, default=1) + parser.add_argument("--stage2_epochs", type=int, default=1) + parser.add_argument("--stage1_lr", type=float, default=0.00020) + parser.add_argument("--stage2_lr", type=float, default=0.00009) + parser.add_argument("--device", default="auto") + parser.add_argument("--amp", default="off") + parser.add_argument("--amp_dtype", default="auto") + parser.add_argument("--compile_model", action="store_true") + parser.add_argument("--compile_mode", default="reduce-overhead") + parser.add_argument("--grad_accum_steps", type=int, default=1) + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--warmup_steps", type=int, default=0) + parser.add_argument("--warmup_ratio", type=float, default=0.05) + parser.add_argument("--min_lr_scale", type=float, default=0.05) + parser.add_argument("--dry_run_dataset", action="store_true", help="Build and summarize the merged v42 dry-run dataset on top of the frozen v41 base.") + parser.add_argument("--train_smoke", action="store_true", help="Run a tiny resumable v42 smoke train on top of the latest v41 base artifact.") + parser.add_argument("--train_frontier", action="store_true", help="Run the full resumable v42 frontier training job.") + parser.add_argument("--stdout", action="store_true", help="Print the prep or smoke summary JSON.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + summary_path = Path(args.summary).resolve() if str(args.summary).strip() else latest_v41_summary_path(REPO_ROOT) + base_zip = Path(args.base_zip).resolve() if str(args.base_zip).strip() else latest_v41_zip_path(REPO_ROOT) + payload = build_v42_prep_pack( + summary_path=summary_path, + output_root=Path(args.output_root), + seed=int(args.seed), + benchmark_limit=int(args.benchmark_limit), + teacher_route_limit=int(args.teacher_route_limit), + verifier_limit=int(args.verifier_limit), + budget_limit=int(args.budget_limit), + diversity_limit=int(args.diversity_limit), + ) + if args.dry_run_dataset: + dry_run_summary = build_training_rows_v42_dry_run( + summary_path=summary_path, + output_root=Path(args.output_root), + seed=int(args.seed), + benchmark_limit=int(args.benchmark_limit), + teacher_route_limit=int(args.teacher_route_limit), + verifier_limit=int(args.verifier_limit), + budget_limit=int(args.budget_limit), + diversity_limit=int(args.diversity_limit), + base_communication_limit=int(args.base_communication_limit), + base_disagreement_limit=int(args.base_disagreement_limit), + ) + payload = payload | {"dry_run_dataset": dry_run_summary} + if args.train_smoke: + smoke_result = train_model_v42( + repo_root=REPO_ROOT, + output_dir=Path(args.output_dir).resolve(), + models_dir=Path(args.models_dir).resolve(), + base_zip=base_zip, + images_dir=Path(args.images_dir).resolve(), + summary_path=summary_path, + prep_output_root=Path(args.output_root).resolve(), + image_size=int(args.image_size), + batch_size=int(args.batch_size), + stage1_epochs=int(args.stage1_epochs), + stage2_epochs=int(args.stage2_epochs), + stage1_lr=float(args.stage1_lr), + stage2_lr=float(args.stage2_lr), + seed=int(args.seed), + benchmark_limit=int(args.benchmark_limit), + teacher_route_limit=int(args.teacher_route_limit), + verifier_limit=int(args.verifier_limit), + budget_limit=int(args.budget_limit), + diversity_limit=int(args.diversity_limit), + base_communication_limit=int(args.base_communication_limit), + base_disagreement_limit=int(args.base_disagreement_limit), + base_distill_limit=int(args.base_distill_limit), + base_teacher_model_limit=int(args.base_teacher_model_limit), + requested_device=str(args.device), + amp_mode=str(args.amp), + amp_dtype=str(args.amp_dtype), + compile_model=bool(args.compile_model), + compile_mode=str(args.compile_mode), + grad_accum_steps=int(args.grad_accum_steps), + ema_decay=float(args.ema_decay), + warmup_steps=int(args.warmup_steps), + warmup_ratio=float(args.warmup_ratio), + min_lr_scale=float(args.min_lr_scale), + smoke_train=True, + ) + payload = payload | {"smoke_train": smoke_result} + if args.train_frontier: + frontier_result = train_model_v42( + repo_root=REPO_ROOT, + output_dir=Path(args.output_dir).resolve(), + models_dir=Path(args.models_dir).resolve(), + base_zip=base_zip, + images_dir=Path(args.images_dir).resolve(), + summary_path=summary_path, + prep_output_root=Path(args.output_root).resolve(), + image_size=int(args.image_size), + batch_size=int(args.batch_size), + stage1_epochs=int(args.stage1_epochs), + stage2_epochs=int(args.stage2_epochs), + stage1_lr=float(args.stage1_lr), + stage2_lr=float(args.stage2_lr), + seed=int(args.seed), + benchmark_limit=int(args.benchmark_limit), + teacher_route_limit=int(args.teacher_route_limit), + verifier_limit=int(args.verifier_limit), + budget_limit=int(args.budget_limit), + diversity_limit=int(args.diversity_limit), + base_communication_limit=int(args.base_communication_limit), + base_disagreement_limit=int(args.base_disagreement_limit), + base_distill_limit=int(args.base_distill_limit), + base_teacher_model_limit=int(args.base_teacher_model_limit), + requested_device=str(args.device), + amp_mode=str(args.amp), + amp_dtype=str(args.amp_dtype), + compile_model=bool(args.compile_model), + compile_mode=str(args.compile_mode), + grad_accum_steps=int(args.grad_accum_steps), + ema_decay=float(args.ema_decay), + warmup_steps=int(args.warmup_steps), + warmup_ratio=float(args.warmup_ratio), + min_lr_scale=float(args.min_lr_scale), + smoke_train=False, + ) + payload = payload | {"train_frontier": frontier_result} + if args.stdout: + print(json.dumps(payload, indent=2, ensure_ascii=True)) + + +if __name__ == "__main__": + main() diff --git a/source/train_omni_collective_v8.py b/source/train_omni_collective_v8.py index 0adc03e09..92ad2b46f 100644 --- a/source/train_omni_collective_v8.py +++ b/source/train_omni_collective_v8.py @@ -9,6 +9,7 @@ import subprocess import sys import time +import warnings import zipfile from collections import defaultdict from datetime import datetime @@ -273,8 +274,21 @@ def _restore_scheduler_state_v8( ), ) if resumed_steps > 0: - for _ in range(resumed_steps): - scheduler.step() + # Cursor restoration is not a training step. Jump directly so a large + # legacy checkpoint does not replay every historical scheduler step. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"Detected call of `lr_scheduler\.step\(\)` before `optimizer\.step\(\)`.*", + category=UserWarning, + ) + warnings.filterwarnings( + "ignore", + message=r"The epoch parameter in `scheduler\.step\(\)` was not necessary.*", + category=UserWarning, + ) + scheduler.step(resumed_steps) + scheduler._step_count = resumed_steps + 1 print( json.dumps( { diff --git a/test_omni_collective_v41.py b/test_omni_collective_v41.py new file mode 100644 index 000000000..6b3bfe0cf --- /dev/null +++ b/test_omni_collective_v41.py @@ -0,0 +1,15 @@ +from source.omni_collective_v41_model import _prompt_variants_v41 + + +def test_prompt_variants_v41_adds_hidden_plan_and_reflection(): + variants = _prompt_variants_v41("Explain why regression tests matter after a refactor.") + + assert any("hidden plan" in item.lower() for item in variants) + assert any("reflection pass" in item.lower() for item in variants) + assert any("human reader" in item.lower() for item in variants) + + +def test_prompt_variants_v41_adds_code_repair_prompt_for_debugging_requests(): + variants = _prompt_variants_v41("The stack trace says Access is denied during checkpoint save. What should I do?") + + assert any("code-review and repair pass" in item.lower() for item in variants) diff --git a/test_omni_collective_v42.py b/test_omni_collective_v42.py new file mode 100644 index 000000000..503565d5f --- /dev/null +++ b/test_omni_collective_v42.py @@ -0,0 +1,15 @@ +from source.omni_collective_v42_model import _budget_hint_v42, _prompt_variants_v42 + + +def test_budget_hint_v42_routes_hard_and_simple_prompts(): + assert _budget_hint_v42("Debug this failing pytest traceback.") == "deep" + assert _budget_hint_v42("Summarize this paragraph.", response_confidence=0.8, domain_confidence=0.8) == "medium" + assert _budget_hint_v42("Say hello.", response_confidence=0.8, domain_confidence=0.8) == "short" + + +def test_prompt_variants_v42_adds_budget_verifier_and_agentic_prompts(): + variants = _prompt_variants_v42("Patch this Python test failure and explain the fix.") + + assert any("reasoning budget" in item.lower() for item in variants) + assert any("verifier pass" in item.lower() for item in variants) + assert any("agentic coding request" in item.lower() for item in variants) diff --git a/test_omni_collective_v8.py b/test_omni_collective_v8.py index 8d0c75a75..137b37f63 100644 --- a/test_omni_collective_v8.py +++ b/test_omni_collective_v8.py @@ -3,6 +3,7 @@ import os import sys import time +import warnings from pathlib import Path import pytest @@ -232,14 +233,18 @@ def test_restore_scheduler_state_v8_falls_back_for_legacy_state(): "_get_lr_called_within_step": False, "_last_lr": [0.6], } - resumed_steps = train_v8._restore_scheduler_state_v8( - scheduler, - legacy_state, - optimizer_steps_done=0, - stage_name="stage1", - checkpoint_source=Path("legacy.pt"), - ) + with warnings.catch_warnings(): + warnings.simplefilter("error") + resumed_steps = train_v8._restore_scheduler_state_v8( + scheduler, + legacy_state, + optimizer_steps_done=0, + stage_name="stage1", + checkpoint_source=Path("legacy.pt"), + ) assert resumed_steps == 5 + assert scheduler.last_epoch == 5 + assert scheduler._step_count == 6 assert float(optimizer.param_groups[0]["lr"]) < 1.0 diff --git a/test_prepare_omni_collective_v41.py b/test_prepare_omni_collective_v41.py new file mode 100644 index 000000000..0b6f919ba --- /dev/null +++ b/test_prepare_omni_collective_v41.py @@ -0,0 +1,91 @@ +import json +import tempfile +from pathlib import Path + +from source.prepare_omni_collective_v41 import build_v41_blueprint, latest_v8_summary_path + + +def _sample_summary(): + return { + "artifact": "supermix_omni_collective_v8_frontier_20260408.zip", + "parameter_count": 112025979, + "dataset_summary": { + "stage1_rows": 34124, + "stage2_rows": 34334, + "teacher_league": { + "teacher_keys": [ + "v40_benchmax", + "omni_collective_v7", + "qwen_v28", + "qwen_v30", + "omni_collective_v6", + ] + }, + }, + "stage1": { + "best_score": 0.3722, + "val_metrics": { + "intent_accuracy": 0.6962, + "response_accuracy": 0.0605, + "vision_accuracy": 0.5384, + "domain_accuracy": 0.625, + }, + }, + "stage2": { + "best_score": 0.4183, + "val_metrics": { + "intent_accuracy": 0.6860, + "response_accuracy": 0.0935, + "vision_accuracy": 0.6923, + "domain_accuracy": 0.6761, + }, + }, + } + + +def test_build_v41_blueprint_raises_targets_from_v8(): + blueprint = build_v41_blueprint(_sample_summary()) + + assert blueprint["family"] == "omni_collective_v41" + assert blueprint["architecture"]["parameter_count_target"] > 112025979 + assert blueprint["architecture"]["parameter_count_target"] <= 132000000 + assert blueprint["architecture"]["expert_count"]["to"] > blueprint["architecture"]["expert_count"]["from"] + assert blueprint["architecture"]["deliberation_passes"]["to"] > blueprint["architecture"]["deliberation_passes"]["from"] + assert blueprint["success_gates"]["response_accuracy_min"] > 0.0935 + assert blueprint["teacher_strategy"]["primary_teachers"][0] == "omni_collective_v8" + + +def test_build_v41_blueprint_keeps_recent_method_references(): + blueprint = build_v41_blueprint(_sample_summary()) + names = {entry["name"] for entry in blueprint["recent_method_references"]} + + assert "DeepSeek-V3 Technical Report" in names + assert "Better & Faster Large Language Models via Multi-token Prediction" in names + assert "s1: Simple test-time scaling" in names + assert "LIMO: Less Is More for Reasoning" in names + assert "RefineCoder" in names + assert "Self-Play Fine-Tuning (SPIN)" in names + assert "Large Language Models have Intrinsic Self-Correction Ability" in names + + +def test_build_v41_blueprint_adds_promotion_protocol(): + blueprint = build_v41_blueprint(_sample_summary()) + + assert "promotion_protocol" in blueprint + assert blueprint["promotion_protocol"]["required_preview_benchmarks"] >= 3 + assert "must_not_regress" in blueprint["promotion_protocol"] + + +def test_latest_v8_summary_path_prefers_newest_summary(): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + older = root / "output" / "supermix_omni_collective_v8_frontier_20260407" + newer = root / "output" / "supermix_omni_collective_v8_frontier_20260408" + older.mkdir(parents=True) + newer.mkdir(parents=True) + older_path = older / "omni_collective_v8_frontier_summary.json" + newer_path = newer / "omni_collective_v8_frontier_summary.json" + older_path.write_text(json.dumps({"artifact": "older"}), encoding="utf-8") + newer_path.write_text(json.dumps({"artifact": "newer"}), encoding="utf-8") + + assert latest_v8_summary_path(root) == newer_path.resolve() diff --git a/test_prepare_omni_collective_v42.py b/test_prepare_omni_collective_v42.py new file mode 100644 index 000000000..39300d884 --- /dev/null +++ b/test_prepare_omni_collective_v42.py @@ -0,0 +1,82 @@ +import json +import tempfile +from pathlib import Path + +from source.prepare_omni_collective_v42 import build_v42_blueprint, latest_v41_summary_path + + +def _sample_summary(): + return { + "artifact": "supermix_omni_collective_v41_frontier_20260410.zip", + "parameter_count": 136200770, + "dataset_summary": { + "stage1_rows": 34086, + "stage2_rows": 34303, + }, + "stage1": { + "best_score": 0.3501, + "val_metrics": { + "intent_accuracy": 0.7400, + "response_accuracy": 0.0890, + "vision_accuracy": 0.1538, + "domain_accuracy": 0.6946, + }, + }, + "stage2": { + "best_score": 0.3575, + "val_metrics": { + "intent_accuracy": 0.5782, + "response_accuracy": 0.1038, + "vision_accuracy": 0.3846, + "domain_accuracy": 0.6874, + }, + }, + } + + +def test_build_v42_blueprint_raises_targets_from_v41(): + blueprint = build_v42_blueprint(_sample_summary()) + + assert blueprint["family"] == "omni_collective_v42" + assert blueprint["architecture"]["parameter_count_target"] > 136200770 + assert blueprint["architecture"]["expert_count"]["to"] > blueprint["architecture"]["expert_count"]["from"] + assert blueprint["architecture"]["deliberation_passes"]["to"] > blueprint["architecture"]["deliberation_passes"]["from"] + assert blueprint["success_gates"]["response_accuracy_min"] > 0.1038 + + +def test_build_v42_blueprint_keeps_new_research_and_teacher_refs(): + blueprint = build_v42_blueprint(_sample_summary()) + names = {entry["name"] for entry in blueprint["recent_method_references"]} + teachers = {entry["model"] for entry in blueprint["teacher_strategy"]["teacher_roles"]} + + assert "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate" in names + assert "Small Language Models Need Strong Verifiers to Self-Correct Reasoning" in names + assert "Learning from Partial Chain-of-Thought via Truncated-Reasoning Self-Distillation" in names + assert "google/gemma-4-31B-it" in names + assert "Qwen/Qwen3.5-397B-A17B" in names + assert "Qwen/Qwen3-Coder-Next" in teachers + assert "Qwen/Qwen3-Omni-30B-A3B-Instruct" in teachers + + +def test_build_v42_blueprint_adds_v40_bridge_context(): + blueprint = build_v42_blueprint(_sample_summary()) + + assert "baseline_findings" in blueprint + assert "v40_strengths_to_import" in blueprint["baseline_findings"] + assert "v40_manifest_snapshot" in blueprint + assert blueprint["teacher_strategy"]["turboquant_runtime_note"]["enabled"] is True + + +def test_latest_v41_summary_path_prefers_newest_summary(): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + older = root / "output" / "supermix_omni_collective_v41_frontier_20260409_010000" + newer = root / "output" / "supermix_omni_collective_v41_frontier_20260410_020000" + older.mkdir(parents=True) + newer.mkdir(parents=True) + older_path = older / "omni_collective_v41_frontier_summary.json" + newer_path = newer / "omni_collective_v41_frontier_summary.json" + older_path.write_text(json.dumps({"artifact": "older"}), encoding="utf-8") + newer_path.write_text(json.dumps({"artifact": "newer"}), encoding="utf-8") + + assert latest_v41_summary_path(root) == newer_path.resolve() diff --git a/test_train_omni_collective_v41.py b/test_train_omni_collective_v41.py new file mode 100644 index 000000000..463c28471 --- /dev/null +++ b/test_train_omni_collective_v41.py @@ -0,0 +1,262 @@ +import json +import json +import tempfile +from pathlib import Path + +from source.train_omni_collective_v41 import ( + _communication_polish_rows_v41, + _code_critique_repair_rows_v41, + _latent_plan_rows_v41, + _reasoning_budget_rows_v41, + _stage_resume_dir_v41, + _smoke_training_rows_v41, + _teacher_disagreement_rows_v41, + _row, + assemble_v41_training_rows, + build_training_rows_v41, + build_training_rows_v41_dry_run, + build_v41_prep_pack, +) + + +def _sample_summary(): + return { + "artifact": "supermix_omni_collective_v8_frontier_20260408.zip", + "parameter_count": 112025979, + "dataset_summary": { + "stage1_rows": 34124, + "stage2_rows": 34334, + "teacher_league": { + "teacher_keys": [ + "v40_benchmax", + "omni_collective_v7", + "qwen_v28", + ] + }, + }, + "stage1": { + "best_score": 0.3722, + "val_metrics": { + "intent_accuracy": 0.6962, + "response_accuracy": 0.0605, + "vision_accuracy": 0.5384, + "domain_accuracy": 0.625, + }, + }, + "stage2": { + "best_score": 0.4183, + "val_metrics": { + "intent_accuracy": 0.6860, + "response_accuracy": 0.0935, + "vision_accuracy": 0.6923, + "domain_accuracy": 0.6761, + }, + }, + "sample_outputs": [ + { + "prompt": "Which local model is best for benchmark-focused reasoning prompts?", + "answer": "Upload an image first, then I can identify the science concept and explain the visual clues.", + } + ], + } + + +def test_communication_polish_rows_v41_covers_multiple_sources(): + rows, counts = _communication_polish_rows_v41(seed=41, limit=40) + + assert rows + assert any(row.source.endswith("clarity_rewrite") for row in rows) + assert any(row.source.endswith("grounded_uncertainty") for row in rows) + assert any(row.source.endswith("coding_explain") for row in rows) + assert any(row.source.endswith("creative_grounded") for row in rows) + assert counts + + +def test_teacher_disagreement_rows_v41_uses_summary_teachers_and_repair_sources(): + rows, summary = _teacher_disagreement_rows_v41(summary=_sample_summary(), seed=41, limit=24) + + assert rows + assert summary["teacher_pair"][0] == "v40_benchmax" + assert any("Two strong local teachers disagree" in row.prompt for row in rows) + assert any(row.source.endswith("route_then_answer") for row in rows) + assert any("v40_benchmax" in row.prompt for row in rows) + + +def test_new_v41_training_slices_cover_plans_repairs_and_budgets(): + latent_rows, latent_summary = _latent_plan_rows_v41(seed=41, limit=16) + repair_rows, repair_summary = _code_critique_repair_rows_v41(seed=41, limit=16) + budget_rows, budget_summary = _reasoning_budget_rows_v41(seed=41, limit=24) + + assert latent_rows + assert any("Hidden plan target" in row.prompt for row in latent_rows) + assert latent_summary + + assert repair_rows + assert any("Access is denied" in row.prompt or "No space left on device" in row.prompt for row in repair_rows) + assert repair_summary + + assert budget_rows + assert any("Reasoning budget: short" in row.prompt for row in budget_rows) + assert any("Reasoning budget: deep" in row.prompt for row in budget_rows) + assert budget_summary + + +def test_build_v41_prep_pack_writes_outputs(): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + summary_path = root / "v8_summary.json" + summary_path.write_text(json.dumps(_sample_summary()), encoding="utf-8") + + payload = build_v41_prep_pack( + summary_path=summary_path, + output_root=root / "v41_prep", + seed=41, + communication_limit=16, + disagreement_limit=12, + ) + + assert payload["communication_polish"]["row_count"] > 0 + assert payload["latent_plan"]["row_count"] > 0 + assert payload["code_critique_repair"]["row_count"] > 0 + assert payload["reasoning_budget"]["row_count"] > 0 + assert payload["teacher_disagreement"]["row_count"] > 0 + assert payload["promotion_eval_pack"]["item_count"] > 0 + assert Path(payload["blueprint_path"]).exists() + assert Path(payload["communication_polish"]["path"]).exists() + assert Path(payload["latent_plan"]["path"]).exists() + assert Path(payload["code_critique_repair"]["path"]).exists() + assert Path(payload["reasoning_budget"]["path"]).exists() + assert Path(payload["teacher_disagreement"]["path"]).exists() + assert Path(payload["promotion_eval_pack"]["path"]).exists() + + +def test_assemble_v41_training_rows_merges_base_and_v41_sources(): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + summary_path = root / "v8_summary.json" + summary_path.write_text(json.dumps(_sample_summary()), encoding="utf-8") + prep_payload = build_v41_prep_pack( + summary_path=summary_path, + output_root=root / "v41_prep", + seed=41, + communication_limit=16, + disagreement_limit=12, + ) + + base_stage1 = [_row("Base stage1 prompt", "Base stage1 answer", source="base_stage1", intent="general", domain="general")] + base_full = list(base_stage1) + [ + _row("Base vision prompt", "Base vision answer", source="base_stage2", intent="vision", domain="vision") + ] + stage1_rows, full_rows, merged = assemble_v41_training_rows( + base_stage1_rows=base_stage1, + base_full_rows=base_full, + base_summary={"source_counts": {"base_stage1": 1, "base_stage2": 1}}, + prep_root=Path(prep_payload["output_root"]), + ) + + assert len(full_rows) > len(base_full) + assert len(stage1_rows) >= len(base_stage1) + assert merged["v41_added_rows"] > 0 + assert "communication_polish" in merged["v41_row_groups"] + assert merged["source_counts"]["base_stage1"] == 1 + + +def test_build_training_rows_v41_uses_base_builder_and_writes_dry_run_summary(monkeypatch): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + summary_path = root / "v8_summary.json" + summary_path.write_text(json.dumps(_sample_summary()), encoding="utf-8") + + from source import train_omni_collective_v41 as module + + def _fake_base_builder(**kwargs): + base_stage1 = [_row("Base prompt", "Base answer", source="base_stage1", intent="general", domain="general")] + base_full = list(base_stage1) + return base_stage1, base_full, {"source_counts": {"base_stage1": 1}, "teacher_league": {"teacher_keys": []}} + + monkeypatch.setattr(module, "build_training_rows_v8", _fake_base_builder) + + stage1_rows, full_rows, merged = build_training_rows_v41( + repo_root=root, + models_dir=root, + images_dir=root, + summary_path=summary_path, + output_root=root / "v41_out", + seed=41, + communication_limit=12, + disagreement_limit=12, + ) + + assert len(stage1_rows) > 1 + assert len(full_rows) > 1 + assert merged["prep_payload"]["family"] == "omni_collective_v41" + assert merged["v41_row_groups"]["teacher_disagreement"]["row_count"] > 0 + + +def test_build_training_rows_v41_dry_run_uses_frozen_counts_and_writes_summary(): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + summary_path = root / "v8_summary.json" + summary = _sample_summary() + summary["dataset_summary"]["source_counts"] = {"base_a": 10, "base_b": 5} + summary_path.write_text(json.dumps(summary), encoding="utf-8") + + payload = build_training_rows_v41_dry_run( + summary_path=summary_path, + output_root=root / "v41_out", + seed=41, + communication_limit=12, + disagreement_limit=12, + ) + + assert payload["dry_run"] is True + assert payload["base_mode"] == "frozen_counts_plus_v41_rows" + assert payload["estimated_stage1_rows"] > summary["dataset_summary"]["stage1_rows"] + assert payload["estimated_stage2_rows"] > summary["dataset_summary"]["stage2_rows"] + assert Path(payload["summary_path"]).exists() + + +def test_smoke_training_rows_v41_keeps_v41_mix_and_multimodal_rows(): + stage1_rows = [ + _row("Explain regression tests.", "They catch unintended changes.", source="communication_polish_v41::clarity_rewrite", intent="general", domain="general"), + _row("Fix this bug.", "Split the logic and add a regression test.", source="teacher_disagreement_v41::repair", intent="coding", domain="coding"), + _row("Choose a model.", "Use v40_benchmax for benchmark reasoning.", source="base_stage1", intent="model_selection", domain="model_selection"), + _row("Plan the work.", "List the steps in order.", source="base_stage1", intent="planning", domain="planning"), + ] + full_rows = list(stage1_rows) + [ + _row("Analyze this image.", "Looks like lab equipment.", source="base_stage2", intent="vision", domain="vision"), + _row("Make a 3D part.", "Use OpenSCAD for a cylinder.", source="base_stage2", intent="general", domain="spatial_3d"), + ] + + smoke_stage1, smoke_full, summary = _smoke_training_rows_v41( + stage1_rows=stage1_rows, + full_rows=full_rows, + seed=41, + ) + + assert smoke_stage1 + assert smoke_full + assert any("_v41" in row.source for row in smoke_full) + assert any(row.domain in {"vision", "spatial_3d"} for row in smoke_full) + assert summary["v41_priority_rows"] >= 1 + + +def test_stage_resume_dir_v41_separates_smoke_and_frontier(tmp_path): + smoke_dir = _stage_resume_dir_v41( + tmp_path, + seed=41, + distill_limit=0, + teacher_model_limit=0, + smoke_train=True, + ) + frontier_dir = _stage_resume_dir_v41( + tmp_path, + seed=41, + distill_limit=0, + teacher_model_limit=0, + smoke_train=False, + ) + + assert smoke_dir != frontier_dir + assert "smoke_" in smoke_dir.name + assert "frontier_" in frontier_dir.name From f595e9c9cd5385bdfcc2ea55f962d3c47d711d20 Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:38:33 +0100 Subject: [PATCH 4/6] feat: add verifier-aware qwen data selection --- source/qwen_supermix_pipeline.py | 370 +++++++++++++++++++++++++--- test_training_improvements.py | 403 +++++++++++++++++++++++++++++-- 2 files changed, 731 insertions(+), 42 deletions(-) diff --git a/source/qwen_supermix_pipeline.py b/source/qwen_supermix_pipeline.py index 550dcc59b..f8ade66c6 100644 --- a/source/qwen_supermix_pipeline.py +++ b/source/qwen_supermix_pipeline.py @@ -1,5 +1,6 @@ import argparse import gc +import heapq import hashlib import json import math @@ -75,6 +76,9 @@ class PreferencePair: reasoning_score: float = 0.0 creativity_score: float = 0.0 knowledge_density_score: float = 0.0 + verifier_score: float = 0.0 + rule_reward: float = 0.0 + verifier_difficulty: float = 0.0 is_followup: bool = False @@ -126,9 +130,12 @@ def _normalize_chat_pair_metadata(raw_metadata: Optional[Dict[str, object]]) -> def _extract_chat_pair_metadata(record: Dict[str, object]) -> Dict[str, object]: raw_meta: Dict[str, object] = {} + nested_metadata = record.get("metadata") + if isinstance(nested_metadata, dict): + raw_meta.update(nested_metadata) for raw_key, raw_value in record.items(): key = _coerce_text(raw_key) - if not key or key in {"user", "assistant", "messages"}: + if not key or key in {"user", "assistant", "messages", "metadata"}: continue if key == "source": key = "record_source" @@ -825,6 +832,81 @@ def _looks_like_reasoning_prompt(text: str) -> bool: return hits >= 1 +def _metadata_number(metadata: Dict[str, object], *keys: str) -> Optional[float]: + if not isinstance(metadata, dict): + return None + lowered = {str(k).strip().lower(): v for k, v in metadata.items()} + for key in keys: + raw = lowered.get(str(key).strip().lower()) + if raw is None: + continue + if isinstance(raw, bool): + return 1.0 if raw else 0.0 + try: + return float(raw) + except (TypeError, ValueError): + continue + return None + + +def _normalize_metric_0_1(value: Optional[float]) -> float: + if value is None: + return 0.0 + parsed = float(value) + if parsed > 10.0: + parsed = parsed / 100.0 + elif parsed > 1.0: + parsed = parsed / 10.0 + return float(max(0.0, min(1.0, parsed))) + + +def _normalize_rule_reward(value: Optional[float]) -> float: + if value is None: + return 0.0 + parsed = float(value) + if parsed > 10.0 or parsed < -10.0: + parsed = parsed / 100.0 + elif parsed > 1.0 or parsed < -1.0: + parsed = parsed / 10.0 + return float(max(-1.0, min(1.0, parsed))) + + +def _chat_pair_verifier_metrics(pair: ChatPair) -> Dict[str, float]: + metadata = pair.metadata if isinstance(pair.metadata, dict) else {} + tests_passed = _metadata_number(metadata, "tests_passed", "passed_tests") + tests_total = _metadata_number(metadata, "tests_total", "total_tests") + pass_rate: Optional[float] = None + if tests_passed is not None and tests_total is not None and tests_total > 0: + pass_rate = float(tests_passed) / float(max(1.0, tests_total)) + + verifier_score = max( + _normalize_metric_0_1(_metadata_number(metadata, "verifier_score", "test_pass_rate", "pass_rate")), + _normalize_metric_0_1(pass_rate), + _normalize_metric_0_1(_metadata_number(metadata, "verified_correct", "verified")), + ) + rule_reward = _normalize_rule_reward( + _metadata_number(metadata, "rule_reward", "verifiable_reward", "reward", "verifier_reward") + ) + difficulty = max( + _normalize_metric_0_1( + _metadata_number(metadata, "test_difficulty", "verifier_difficulty", "difficulty", "problem_difficulty") + ), + 0.0, + ) + verifier_bonus = ( + 0.16 * verifier_score + + 0.08 * max(0.0, rule_reward) + + 0.08 * verifier_score * difficulty + - 0.12 * max(0.0, -rule_reward) + ) + return { + "verifier_score": float(verifier_score), + "rule_reward": float(rule_reward), + "verifier_difficulty": float(difficulty), + "verifier_bonus": float(max(-0.12, min(0.34, verifier_bonus))), + } + + def _iter_clean_pairs_from_jsonl( path: Path, min_chars: int, @@ -2377,6 +2459,9 @@ def build_preference_rows( "reasoning_score": float(max(0.0, pair.reasoning_score)), "creativity_score": float(max(0.0, pair.creativity_score)), "knowledge_density_score": float(max(0.0, pair.knowledge_density_score)), + "verifier_score": float(max(0.0, min(1.0, pair.verifier_score))), + "rule_reward": float(max(-1.0, min(1.0, pair.rule_reward))), + "verifier_difficulty": float(max(0.0, min(1.0, pair.verifier_difficulty))), "is_followup": bool(pair.is_followup), "pair_tokens": int(len(chosen["input_ids"]) + len(rejected["input_ids"])), } @@ -2399,6 +2484,9 @@ def _preference_row_rescore_weight( reasoning_score = max(0.0, float(row.get("reasoning_score", 0.0))) creativity_score = max(0.0, float(row.get("creativity_score", 0.0))) knowledge_density_score = max(0.0, float(row.get("knowledge_density_score", 0.0))) + verifier_score = max(0.0, min(1.0, float(row.get("verifier_score", 0.0)))) + rule_reward = max(-1.0, min(1.0, float(row.get("rule_reward", 0.0)))) + verifier_difficulty = max(0.0, min(1.0, float(row.get("verifier_difficulty", 0.0)))) followup_bonus = 0.04 if bool(row.get("is_followup", False)) else 0.0 easy_focus = 0.84 + 0.46 * min(1.0, quality_gap / 0.24) @@ -2417,7 +2505,14 @@ def _preference_row_rescore_weight( ) prompt_focus = 0.96 + 0.04 * min(2.5, prompt_complexity) selection_focus = 1.0 + 0.03 * min(4.0, selection_score) - raw_weight = base_weight * difficulty_focus * novelty_focus * prompt_focus * selection_focus + verifier_focus = ( + 1.0 + + 0.10 * verifier_score + + 0.06 * verifier_score * verifier_difficulty + + 0.05 * max(0.0, rule_reward) + - 0.08 * max(0.0, -rule_reward) + ) + raw_weight = base_weight * difficulty_focus * novelty_focus * prompt_focus * selection_focus * verifier_focus return float(max(0.1, raw_weight)) @@ -2889,6 +2984,11 @@ def _sft_pair_selection_score( conversation_signal = max(0.0, float(alignment.conversation)) creativity_signal = max(0.0, float(alignment.creativity)) prompt_complexity = min(1.0, float(_prompt_complexity_score(pair.user)) / 1.35) + verifier_metrics = _chat_pair_verifier_metrics(pair) + verifier_score = float(verifier_metrics.get("verifier_score", 0.0)) + rule_reward = float(verifier_metrics.get("rule_reward", 0.0)) + verifier_difficulty = float(verifier_metrics.get("verifier_difficulty", 0.0)) + verifier_bonus = float(verifier_metrics.get("verifier_bonus", 0.0)) word_count = max(1, _word_token_count(assistant)) if word_count <= 18: @@ -2912,6 +3012,7 @@ def _sft_pair_selection_score( utility += 0.05 if "quality_anchor" in str(pair.source or "").lower(): utility += 0.03 + utility += verifier_bonus difficulty = min( 1.0, @@ -2920,6 +3021,8 @@ def _sft_pair_selection_score( + 0.18 * min(1.0, float(word_count) / 180.0) + 0.18 * quality_signal, ) + if verifier_score > 0.0 or rule_reward != 0.0: + difficulty = max(difficulty, 0.35 * prompt_complexity + 0.65 * verifier_difficulty) score = utility mode_norm = str(mode).strip().lower() @@ -2939,6 +3042,10 @@ def _sft_pair_selection_score( "prompt_complexity": float(prompt_complexity), "difficulty": float(difficulty), "compactness": float(compactness), + "verifier_score": float(verifier_score), + "rule_reward": float(rule_reward), + "verifier_difficulty": float(verifier_difficulty), + "verifier_bonus": float(verifier_bonus), "words": float(word_count), "utility": float(utility), "score": float(score), @@ -3046,6 +3153,97 @@ def _apply_sft_coverage_topk_scores( return adjusted +def _select_sft_coverage_greedy( + scored: Sequence[Tuple[float, float, ChatPair, Dict[str, float]]], + keep_n: int, + budget_mode: str, + budget_power: float, + token_budget: Optional[float] = None, + min_keep: int = 0, + max_keep: int = 0, +) -> List[Tuple[float, float, ChatPair, Dict[str, float]]]: + """Select a quality-anchored coverage set with exact lazy greedy updates. + + Coverage gains only decrease as buckets become represented, so cached + marginal scores are valid upper bounds. A heap therefore preserves the + exact greedy choice without rescoring every remaining row on every pick. + """ + if not scored: + return [] + keep_n = max(1, min(len(scored), int(keep_n))) + entries = list(scored) + selected: List[Tuple[float, float, ChatPair, Dict[str, float]]] = [] + seen_sources: set = set() + seen_groups: set = set() + seen_styles: set = set() + seen_lengths: set = set() + token_mode = str(budget_mode).strip().lower() == "tokens" + budget_power = max(0.0, min(1.0, float(budget_power))) + token_budget_limit = None if token_budget is None else max(0.0, float(token_budget)) + min_keep = max(0, min(keep_n, int(min_keep))) + max_keep = max(0, min(keep_n, int(max_keep))) + if max_keep > 0: + max_keep = max(max_keep, min_keep) + selected_tokens = 0.0 + + def _marginal_score(entry: Tuple[float, float, ChatPair, Dict[str, float]]) -> float: + _budget_value, score, _pair, metrics = entry + source_key = str(metrics.get("source_key", "")) + group_key = str(metrics.get("group_key", "")) + style_key = str(metrics.get("style_key", "")) + length_key = str(metrics.get("length_key", "")) + quality_signal = max(0.0, min(1.0, float(metrics.get("quality_signal", 0.0)))) + coverage_gain = 0.0 + if source_key and source_key not in seen_sources: + coverage_gain += 0.18 + if group_key and group_key not in seen_groups: + coverage_gain += 0.30 + if style_key and style_key not in seen_styles: + coverage_gain += 0.20 + if length_key and length_key not in seen_lengths: + coverage_gain += 0.06 + marginal = float(score) + coverage_gain * (0.70 + 0.30 * quality_signal) + if token_mode: + tokens = max(8.0, float(metrics.get("estimated_tokens", 0.0))) + marginal /= math.pow(tokens, budget_power) + return float(marginal) + + heap = [ + (-_marginal_score(entry), -float(entry[1]), idx) + for idx, entry in enumerate(entries) + ] + heapq.heapify(heap) + while heap and len(selected) < keep_n: + if max_keep > 0 and len(selected) >= max_keep: + break + _neg_cached, _neg_base, idx = heapq.heappop(heap) + entry = entries[idx] + current = _marginal_score(entry) + entry_tokens = max(0.0, float(entry[3].get("estimated_tokens", 0.0))) + if token_budget_limit is not None: + must_keep = not selected or len(selected) < min_keep + under_budget = (selected_tokens + entry_tokens) <= token_budget_limit + if not must_keep and not under_budget: + # Remaining budget only shrinks, so this row can never become + # feasible later and must not update the represented buckets. + continue + current_key = (current, float(entry[1]), -idx) + if heap: + next_idx = int(heap[0][2]) + next_upper_key = (-float(heap[0][0]), -float(heap[0][1]), -next_idx) + if current_key < next_upper_key: + heapq.heappush(heap, (-current, -float(entry[1]), idx)) + continue + selected.append(entry) + metrics = entry[3] + seen_sources.add(str(metrics.get("source_key", ""))) + seen_groups.add(str(metrics.get("group_key", ""))) + seen_styles.add(str(metrics.get("style_key", ""))) + seen_lengths.add(str(metrics.get("length_key", ""))) + selected_tokens += entry_tokens + return selected + + def _is_scoped_sft_selection_pair( pair: ChatPair, scope: str, @@ -3151,28 +3349,41 @@ def _select_sft_training_pairs( if budget_mode == "tokens" and keep_ratio < 0.999: token_budget = max(0.0, total_estimated_tokens * keep_ratio) min_keep_target = min(min_keep, total) if min_keep > 0 else 0 - for entry in scored: - if max_keep > 0 and len(selected) >= max_keep: - break - entry_tokens = float(entry[3].get("estimated_tokens", 0.0)) - must_keep = len(selected) < max(1, min_keep_target) - under_budget = (selected_estimated_tokens + entry_tokens) <= token_budget - if not selected or must_keep or under_budget: - selected.append(entry) - selected_estimated_tokens += entry_tokens - if min_keep_target > 0 and len(selected) < min_keep_target: - selected_ids = {id(entry[2]) for entry in selected} + if mode == "coverage_topk": + selected = _select_sft_coverage_greedy( + scored, + keep_n=total, + budget_mode=budget_mode, + budget_power=budget_power, + token_budget=token_budget, + min_keep=min_keep_target, + max_keep=max_keep, + ) + else: for entry in scored: - if id(entry[2]) in selected_ids: - continue - selected.append(entry) - selected_estimated_tokens += float(entry[3].get("estimated_tokens", 0.0)) - selected_ids.add(id(entry[2])) - if len(selected) >= min_keep_target: + if max_keep > 0 and len(selected) >= max_keep: break - if not selected: - selected = [scored[0]] - selected_estimated_tokens = float(selected[0][3].get("estimated_tokens", 0.0)) + entry_tokens = float(entry[3].get("estimated_tokens", 0.0)) + must_keep = len(selected) < max(1, min_keep_target) + under_budget = (selected_estimated_tokens + entry_tokens) <= token_budget + if not selected or must_keep or under_budget: + selected.append(entry) + selected_estimated_tokens += entry_tokens + if min_keep_target > 0 and len(selected) < min_keep_target: + selected_ids = {id(entry[2]) for entry in selected} + for entry in scored: + if id(entry[2]) in selected_ids: + continue + selected.append(entry) + selected_estimated_tokens += float(entry[3].get("estimated_tokens", 0.0)) + selected_ids.add(id(entry[2])) + if len(selected) >= min_keep_target: + break + if not selected: + selected = [scored[0]] + selected_estimated_tokens = float( + sum(float(entry[3].get("estimated_tokens", 0.0)) for entry in selected) + ) else: keep_n = total if keep_ratio < 0.999: @@ -3182,7 +3393,15 @@ def _select_sft_training_pairs( if min_keep > 0: keep_n = max(keep_n, min(min_keep, total)) keep_n = max(1, min(total, keep_n)) - selected = scored[:keep_n] + if mode == "coverage_topk": + selected = _select_sft_coverage_greedy( + scored, + keep_n=keep_n, + budget_mode=budget_mode, + budget_power=budget_power, + ) + else: + selected = scored[:keep_n] selected_estimated_tokens = float( sum(float(entry[3].get("estimated_tokens", 0.0)) for entry in selected) ) @@ -3200,6 +3419,8 @@ def _mean(metric_name: str, rows: Sequence[Tuple[float, float, ChatPair, Dict[st f"quality={_mean('quality', scored):.3f}->{_mean('quality', selected):.3f} " f"density={_mean('density', scored):.3f}->{_mean('density', selected):.3f} " f"reason={_mean('reasoning', scored):.3f}->{_mean('reasoning', selected):.3f} " + f"verifier={_mean('verifier_score', scored):.3f}->{_mean('verifier_score', selected):.3f} " + f"rule_reward={_mean('rule_reward', scored):.3f}->{_mean('rule_reward', selected):.3f} " f"difficulty={_mean('difficulty', scored):.3f}->{_mean('difficulty', selected):.3f} " f"compactness={_mean('compactness', scored):.3f}->{_mean('compactness', selected):.3f} " f"words={_mean('words', scored):.1f}->{_mean('words', selected):.1f} " @@ -3752,17 +3973,26 @@ def _preference_selection_score( reasoning_score = max(0.0, float(pair.reasoning_score)) creativity_score = max(0.0, float(pair.creativity_score)) knowledge_density_score = max(0.0, float(pair.knowledge_density_score)) + verifier_score = max(0.0, min(1.0, float(pair.verifier_score))) + rule_reward = max(-1.0, min(1.0, float(pair.rule_reward))) + verifier_difficulty = max(0.0, min(1.0, float(pair.verifier_difficulty))) + verifier_bonus = ( + 0.16 * verifier_score + + 0.08 * max(0.0, rule_reward) + + 0.08 * verifier_score * verifier_difficulty + - 0.12 * max(0.0, -rule_reward) + ) followup_bonus = 0.10 if bool(pair.is_followup) else 0.0 if mode == "margin_topk": - return float(base_weight * (quality_gap + 0.35 * rejected_similarity + 0.08 * prompt_complexity)) + return float(base_weight * (quality_gap + 0.35 * rejected_similarity + 0.08 * prompt_complexity + verifier_bonus)) if mode == "capacity_aware": target = max(0.0, min(1.0, float(hardness_target))) bw = max(0.05, float(hardness_bandwidth)) z = (rejected_similarity - target) / bw hardness_window = math.exp(-0.5 * z * z) - return float(base_weight * hardness_window * (quality_gap + 0.12 * prompt_complexity + 0.05)) + return float(base_weight * hardness_window * (quality_gap + 0.12 * prompt_complexity + verifier_bonus + 0.05)) if mode == "innovation_mix": target = max(0.0, min(1.0, float(hardness_target))) @@ -3779,7 +4009,7 @@ def _preference_selection_score( return float( base_weight * (0.35 + 0.65 * hardness_window) - * (quality_gap + 0.12 * prompt_complexity + novelty_bonus + 0.05) + * (quality_gap + 0.12 * prompt_complexity + novelty_bonus + verifier_bonus + 0.05) ) if mode == "coverage_margin": @@ -3797,7 +4027,7 @@ def _preference_selection_score( return float( base_weight * (0.25 + 0.75 * hardness_window) - * (quality_gap + 0.14 * prompt_complexity + novelty_bonus + 0.05) + * (quality_gap + 0.14 * prompt_complexity + novelty_bonus + verifier_bonus + 0.05) ) return float(max(0.25, pair.weight)) @@ -3873,6 +4103,72 @@ def _apply_preference_coverage_scores( return adjusted +def _select_preference_coverage_greedy( + scored: Sequence[Tuple[float, PreferencePair]], + keep_n: int, +) -> List[PreferencePair]: + """Select preference pairs with exact lazy-greedy coverage updates.""" + if not scored: + return [] + keep_n = max(1, min(len(scored), int(keep_n))) + entries = list(scored) + selected: List[PreferencePair] = [] + seen_signatures: set = set() + seen_styles: set = set() + seen_hardness: set = set() + seen_lengths: set = set() + + def _length_key(pair: PreferencePair) -> str: + return _word_count_bucket( + max( + _word_token_count(_fast_cleanup_response_text(pair.chosen)), + _word_token_count(_fast_cleanup_response_text(pair.rejected)), + ) + ) + + def _marginal_score(entry: Tuple[float, PreferencePair]) -> float: + score, pair = entry + signature_key = _prompt_signature(pair.user)[:96] or "" + style_key = _preference_pair_style_bucket(pair) + hardness_key = _preference_pair_hardness_bucket(pair) + length_key = _length_key(pair) + quality_anchor = 0.60 + 0.40 * min(1.0, max(0.0, float(pair.quality_gap)) / 0.45) + coverage_gain = 0.0 + if signature_key not in seen_signatures: + coverage_gain += 0.06 + if style_key not in seen_styles: + coverage_gain += 0.34 + if hardness_key not in seen_hardness: + coverage_gain += 0.08 + if length_key not in seen_lengths: + coverage_gain += 0.04 + return float(score + coverage_gain * quality_anchor) + + heap = [ + (-_marginal_score(entry), -float(entry[0]), idx) + for idx, entry in enumerate(entries) + ] + heapq.heapify(heap) + while heap and len(selected) < keep_n: + _neg_cached, _neg_base, idx = heapq.heappop(heap) + entry = entries[idx] + current = _marginal_score(entry) + current_key = (current, float(entry[0]), -idx) + if heap: + next_idx = int(heap[0][2]) + next_upper_key = (-float(heap[0][0]), -float(heap[0][1]), -next_idx) + if current_key < next_upper_key: + heapq.heappush(heap, (-current, -float(entry[0]), idx)) + continue + _score, pair = entry + selected.append(pair) + seen_signatures.add(_prompt_signature(pair.user)[:96] or "") + seen_styles.add(_preference_pair_style_bucket(pair)) + seen_hardness.add(_preference_pair_hardness_bucket(pair)) + seen_lengths.add(_length_key(pair)) + return selected + + def _select_preference_pairs( pairs: Sequence[PreferencePair], strategy: str, @@ -3922,7 +4218,10 @@ def _select_preference_pairs( if min_keep > 0: keep_n = max(keep_n, min(min_keep, total)) keep_n = max(1, min(total, keep_n)) - selected = [p for _s, p in scored[:keep_n]] + if mode == "coverage_margin": + selected = _select_preference_coverage_greedy(scored, keep_n=keep_n) + else: + selected = [p for _s, p in scored[:keep_n]] def _mean(vals: Sequence[float]) -> float: if not vals: @@ -4563,6 +4862,18 @@ def build_preference_pairs_with_generation( pair_weight *= float(max(0.7, coding_focus_boost)) if prompt_is_reasoning: pair_weight *= float(max(0.7, reasoning_focus_boost)) + verifier_metrics = _chat_pair_verifier_metrics(pair) + verifier_score = float(verifier_metrics.get("verifier_score", 0.0)) + rule_reward = float(verifier_metrics.get("rule_reward", 0.0)) + verifier_difficulty = float(verifier_metrics.get("verifier_difficulty", 0.0)) + if verifier_score > 0.0 or rule_reward != 0.0: + pair_weight *= ( + 1.0 + + 0.08 * verifier_score + + 0.05 * verifier_score * verifier_difficulty + + 0.04 * max(0.0, rule_reward) + - 0.06 * max(0.0, -rule_reward) + ) pair_weight = float(max(0.75, min(3.25, pair_weight))) out.append( PreferencePair( @@ -4577,6 +4888,9 @@ def build_preference_pairs_with_generation( reasoning_score=float(max(0.0, float(chosen_alignment.reasoning))), creativity_score=float(max(0.0, float(chosen_alignment.creativity))), knowledge_density_score=float(max(0.0, float(chosen_density))), + verifier_score=float(verifier_score), + rule_reward=float(rule_reward), + verifier_difficulty=float(verifier_difficulty), is_followup=bool(chosen_alignment.is_followup), ) ) diff --git a/test_training_improvements.py b/test_training_improvements.py index c958dbed8..6abe14774 100644 --- a/test_training_improvements.py +++ b/test_training_improvements.py @@ -72,6 +72,17 @@ def _import_preference_selection_helpers(): return PreferencePair, _select_preference_pairs +def _import_verifier_helpers(): + from source.qwen_supermix_pipeline import ( + ChatPair, + PreferencePair, + _chat_pair_verifier_metrics, + build_preference_rows, + ) + + return ChatPair, PreferencePair, _chat_pair_verifier_metrics, build_preference_rows + + def _import_split_and_io_helpers(): from source.qwen_supermix_pipeline import ChatPair, load_saved_chat_pairs, save_jsonl, split_train_eval @@ -673,24 +684,33 @@ def test_sft_coverage_selection_keeps_rare_high_quality_group(): ) = _import_pipeline() pairs = [ ChatPair( - user="Explain binary search step by step and give the time complexity.", + user="Explain binary search and its time complexity.", assistant=( - "Check the midpoint, discard the half that cannot contain the target, and repeat. " - "Because each comparison halves the remaining range, binary search runs in O(log n) time." + "Binary search works on a sorted array by comparing the middle element " + "with the target. If they match, the search ends. Otherwise, the half where " + "the target cannot lie is discarded. Each step halves the search space, " + "so the algorithm runs in O(log n) time." ), source="conversation_data.coding_knowledge_2026_02_19.jsonl", ), ChatPair( - user="Explain merge sort step by step and give the time complexity.", + user="Explain merge sort and its time complexity.", assistant=( - "Split the array recursively, sort each half, and merge the sorted halves. " - "The merge work is linear per level, so merge sort runs in O(n log n) time." + "Merge sort splits an array in half recursively until each sub-array has " + "one element. It then merges adjacent sub-arrays in sorted order. Since " + "each level of merging processes all n elements and there are log n levels, " + "the total time complexity is O(n log n)." ), source="conversation_data.coding_knowledge_2026_02_19.jsonl", ), ChatPair( - user="What happened on July 20, 1969?", - assistant="Apollo 11 landed on the Moon, and Neil Armstrong became the first person to walk on it.", + user="Describe the Apollo 11 Moon landing.", + assistant=( + "Apollo 11 was the first crewed mission to land on the Moon. On July 20, " + "1969, Neil Armstrong and Buzz Aldrin descended to the lunar surface while " + "Michael Collins orbited overhead. Armstrong became the first human to walk " + "on the Moon, calling it a giant leap for mankind." + ), source="conversation_data.world_events_2026_02_19.jsonl", metadata={"event_id": "apollo_11_moon_landing", "topic": "space_history"}, ), @@ -744,10 +764,10 @@ def test_preference_coverage_margin_preserves_style_diversity(): quality_gap=0.34, rejected_similarity=0.44, prompt_complexity=0.66, - conversation_score=0.18, - reasoning_score=0.79, - creativity_score=0.06, - knowledge_density_score=0.86, + conversation_score=0.10, + reasoning_score=0.42, + creativity_score=0.04, + knowledge_density_score=0.44, ), PreferencePair( user="Write a creative but concise reflection on learning algorithms.", @@ -757,9 +777,9 @@ def test_preference_coverage_margin_preserves_style_diversity(): quality_gap=0.30, rejected_similarity=0.40, prompt_complexity=0.60, - conversation_score=0.58, + conversation_score=0.92, reasoning_score=0.22, - creativity_score=0.94, + creativity_score=0.96, knowledge_density_score=0.28, ), ] @@ -787,6 +807,355 @@ def test_preference_coverage_margin_preserves_style_diversity(): print(" [ok] coverage_margin keeps style-diverse preference pairs") +def test_sft_lazy_greedy_coverage_avoids_redundant_clusters(): + from source.qwen_supermix_pipeline import ChatPair, _select_sft_coverage_greedy + + first = ChatPair(user="Explain search.", assistant="Use a midpoint and halve the interval.", source="coding") + duplicate = ChatPair(user="Explain sorting.", assistant="Split, order, and merge the intervals.", source="coding") + diverse = ChatPair(user="Write a moon vignette.", assistant="Silver dust rose beneath a quiet lander.", source="creative") + common = { + "source_key": "coding", + "group_key": "algorithms", + "style_key": "reasoning", + "length_key": "short", + "quality_signal": 1.0, + "estimated_tokens": 24.0, + } + scored = [ + (1.00, 1.00, first, dict(common)), + (0.99, 0.99, duplicate, dict(common)), + ( + 0.80, + 0.80, + diverse, + { + "source_key": "creative", + "group_key": "lunar_fiction", + "style_key": "creative", + "length_key": "medium", + "quality_signal": 1.0, + "estimated_tokens": 28.0, + }, + ), + ] + + selected = _select_sft_coverage_greedy(scored, keep_n=2, budget_mode="pairs", budget_power=0.5) + selected_pairs = [entry[2] for entry in selected] + + assert selected_pairs == [first, diverse] + print(" [ok] lazy-greedy SFT coverage avoids a redundant second cluster row") + + +def test_sft_token_budget_coverage_uses_lazy_greedy(monkeypatch): + import source.qwen_supermix_pipeline as pipeline + + pairs = [ + pipeline.ChatPair( + user="Explain binary search briefly.", + assistant="Check the midpoint and discard half of the sorted interval each step.", + source="coding", + ), + pipeline.ChatPair( + user="Explain merge sort briefly.", + assistant="Split recursively, then merge sorted halves in linear work per level.", + source="coding", + ), + pipeline.ChatPair( + user="Write a concise lunar vignette.", + assistant="Silver dust rose beneath the lander as Earth watched from a black sky.", + source="creative", + ), + ] + original = pipeline._select_sft_coverage_greedy + calls = [] + + def recording_selector( + scored, + keep_n, + budget_mode, + budget_power, + token_budget=None, + min_keep=0, + max_keep=0, + ): + calls.append((keep_n, budget_mode, budget_power, token_budget)) + return original( + scored, + keep_n, + budget_mode, + budget_power, + token_budget=token_budget, + min_keep=min_keep, + max_keep=max_keep, + ) + + monkeypatch.setattr(pipeline, "_select_sft_coverage_greedy", recording_selector) + selected = pipeline._select_sft_training_pairs( + pairs, + strategy="coverage_topk", + keep_ratio=0.60, + min_keep=0, + max_keep=0, + hardness_target=0.56, + hardness_bandwidth=0.30, + budget_mode="tokens", + budget_power=0.5, + ) + + assert selected + assert len(calls) == 1 + assert calls[0][:3] == (len(pairs), "tokens", 0.5) + assert calls[0][3] is not None + print(" [ok] token-budget coverage is ordered by lazy-greedy marginals") + + +def test_sft_token_budget_lazy_greedy_does_not_cover_skipped_rows(): + from source.qwen_supermix_pipeline import ChatPair, _select_sft_coverage_greedy + + first = ChatPair(user="Anchor", assistant="Anchor answer", source="anchor") + too_large = ChatPair(user="Large rare", assistant="Large rare answer", source="rare") + feasible_rare = ChatPair(user="Small rare", assistant="Small rare answer", source="rare") + feasible_duplicate = ChatPair(user="Duplicate", assistant="Duplicate answer", source="anchor") + common = { + "quality_signal": 1.0, + "source_key": "anchor", + "group_key": "common", + "style_key": "reasoning", + "length_key": "short", + "estimated_tokens": 5.0, + } + rare = { + "quality_signal": 1.0, + "source_key": "rare", + "group_key": "rare-event", + "style_key": "creative", + "length_key": "medium", + } + scored = [ + (2.0, 2.0, first, dict(common)), + (1.1, 1.1, too_large, {**rare, "estimated_tokens": 6.0}), + (0.2, 0.2, feasible_rare, {**rare, "estimated_tokens": 5.0}), + (0.8, 0.8, feasible_duplicate, dict(common)), + ] + + selected = _select_sft_coverage_greedy( + scored, + keep_n=len(scored), + budget_mode="tokens", + budget_power=0.0, + token_budget=10.0, + ) + + assert [entry[2] for entry in selected] == [first, feasible_rare] + print(" [ok] infeasible rows do not consume lazy-greedy coverage credit") + + +def test_preference_lazy_greedy_coverage_avoids_redundant_styles(): + from source.qwen_supermix_pipeline import PreferencePair, _select_preference_coverage_greedy + + first = PreferencePair( + user="Explain the algorithm.", + chosen="It halves the search interval after each midpoint comparison.", + rejected="It searches efficiently.", + quality_gap=0.45, + rejected_similarity=0.40, + reasoning_score=0.90, + knowledge_density_score=0.70, + ) + duplicate = PreferencePair( + user="Explain the algorithm.", + chosen="It repeatedly eliminates half of the sorted search range.", + rejected="It is a good algorithm.", + quality_gap=0.45, + rejected_similarity=0.40, + reasoning_score=0.88, + knowledge_density_score=0.68, + ) + diverse = PreferencePair( + user="Write a tiny lunar image.", + chosen="Moonlight pooled like silver water around the lander.", + rejected="The moon looked nice.", + quality_gap=0.45, + rejected_similarity=0.40, + creativity_score=0.96, + conversation_score=0.60, + ) + selected = _select_preference_coverage_greedy( + [(1.00, first), (0.99, duplicate), (0.78, diverse)], + keep_n=2, + ) + + assert selected == [first, diverse] + print(" [ok] lazy-greedy preference coverage avoids a redundant second style") + + +def test_sft_selection_uses_verifier_reward_for_hard_math_code_rows(): + ( + _build_lr_lambda, + _LengthBucketBatchSampler, + _pair_knowledge_density_score, + _select_sft_training_pairs, + _counterfactual_reject_variants, + ChatPair, + ) = _import_pipeline() + del _build_lr_lambda, _LengthBucketBatchSampler, _pair_knowledge_density_score, _counterfactual_reject_variants + + verified = ChatPair( + user="Solve 17 * 19 and explain the modular arithmetic check.", + assistant="17 * 19 = 323. A quick check is 323 mod 17 = 0 and 323 mod 19 = 0, so the product is consistent.", + source="math_verifier.jsonl", + metadata={ + "verifier_score": 1.0, + "rule_reward": 1.0, + "test_difficulty": 0.95, + "tests_passed": 8, + "tests_total": 8, + }, + ) + unverified = ChatPair( + user="Explain binary search and its time complexity.", + assistant=( + "Binary search compares the midpoint of a sorted array with the target and discards the impossible half. " + "Repeating this halves the search space each step, so the runtime is O(log n)." + ), + source="coding_knowledge.jsonl", + ) + + selected = _select_sft_training_pairs( + [unverified, verified], + strategy="utility_topk", + keep_ratio=0.5, + min_keep=1, + max_keep=1, + hardness_target=0.45, + hardness_bandwidth=0.25, + ) + + assert selected == [verified], "Verifier-backed hard math row should win a tight SFT selection budget" + print(" [ok] verifier reward boosts hard SFT rows") + + +def test_preference_selection_uses_verifier_reward_margin(): + PreferencePair, _select_preference_pairs = _import_preference_selection_helpers() + verified = PreferencePair( + user="Write a tested Python function for modular exponentiation.", + chosen="Use pow(base, exp, mod), validate mod is positive, and return the computed residue.", + rejected="Loop over multiplication without checking the modulus.", + weight=1.0, + quality_gap=0.22, + rejected_similarity=0.30, + prompt_complexity=1.0, + reasoning_score=0.45, + knowledge_density_score=0.45, + verifier_score=1.0, + rule_reward=1.0, + verifier_difficulty=0.90, + ) + unverified = PreferencePair( + user="Explain merge sort briefly.", + chosen="Merge sort recursively splits an array, sorts each half, and merges sorted halves in O(n log n).", + rejected="Merge sort repeatedly scans the array until it looks sorted.", + weight=1.0, + quality_gap=0.34, + rejected_similarity=0.30, + prompt_complexity=1.0, + reasoning_score=0.30, + knowledge_density_score=0.30, + ) + + selected = _select_preference_pairs( + [unverified, verified], + strategy="margin_topk", + keep_ratio=0.5, + min_keep=1, + max_keep=1, + hardness_target=0.45, + hardness_bandwidth=0.25, + ) + + assert selected == [verified], "Verifier-backed preference pair should survive tight margin selection" + print(" [ok] verifier reward boosts preference-pair selection") + + +def test_verifier_metadata_normalization_and_nested_jsonl_loading(): + ChatPair, _PreferencePair, _chat_pair_verifier_metrics, _build_preference_rows = _import_verifier_helpers() + del _PreferencePair, _build_preference_rows + _SplitPair, _split_train_eval, save_jsonl, load_saved_chat_pairs = _import_split_and_io_helpers() + del _SplitPair, _split_train_eval + from source.qwen_supermix_pipeline import load_jsonl_pairs + + pair = ChatPair( + user="Solve 2x + 3 = 7.", + assistant="Subtract 3 to get 2x = 4, so x = 2.", + source="math_verifier.jsonl", + metadata={"tests_passed": 8, "tests_total": 10, "rule_reward": 75, "verifier_difficulty": 0.4}, + ) + metrics = _chat_pair_verifier_metrics(pair) + assert abs(metrics["verifier_score"] - 0.8) < 1e-9 + assert abs(metrics["rule_reward"] - 0.75) < 1e-9 + assert abs(metrics["verifier_difficulty"] - 0.4) < 1e-9 + + negative = ChatPair( + user="Return the exact result of a unit test.", + assistant="The test failed.", + metadata={"rule_reward": -50}, + ) + neg_metrics = _chat_pair_verifier_metrics(negative) + assert abs(neg_metrics["rule_reward"] + 0.5) < 1e-9 + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "verifier_pairs.jsonl" + path.write_text( + ( + '{"user":"Solve 3+4.","assistant":"7","source":"math.jsonl",' + '"metadata":{"tests_passed":1,"tests_total":1,"rule_reward":1.0,' + '"verifier_difficulty":0.4}}\n' + ), + encoding="utf-8", + ) + loaded = load_saved_chat_pairs(path) + + assert loaded[0].metadata["tests_passed"] == 1 + assert loaded[0].metadata["rule_reward"] == 1.0 + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "raw_verifier_pairs.jsonl" + path.write_text( + ( + '{"user":"Solve 8+5.","assistant":"8 + 5 = 13.","source":"math_raw.jsonl",' + '"metadata":{"tests_passed":1,"tests_total":1,"rule_reward":1.0,' + '"verifier_difficulty":0.5}}\n' + ), + encoding="utf-8", + ) + raw_loaded = load_jsonl_pairs([str(path)], max_records=10, min_chars=1) + + assert raw_loaded[0].metadata["tests_passed"] == 1 + assert raw_loaded[0].metadata["record_source"] == "math_raw.jsonl" + print(" [ok] verifier metadata normalizes and nested JSONL metadata loads") + + +def test_preference_rows_preserve_verifier_fields(): + _ChatPair, PreferencePair, _chat_pair_verifier_metrics, build_preference_rows = _import_verifier_helpers() + del _ChatPair, _chat_pair_verifier_metrics + pair = PreferencePair( + user="Solve the checked example.", + chosen="The checked answer is correct.", + rejected="The answer is probably correct.", + verifier_score=0.9, + rule_reward=0.8, + verifier_difficulty=0.6, + ) + + rows = build_preference_rows(_ToyTokenizer(), [pair], max_length=256) + + assert len(rows) == 1 + assert rows[0]["verifier_score"] == 0.9 + assert rows[0]["rule_reward"] == 0.8 + assert rows[0]["verifier_difficulty"] == 0.6 + print(" [ok] preference rows preserve verifier fields") + + def test_benchmark_sample_comparison_ranks_worst_regressions_first(): build_benchmark_sample_comparison = _import_benchmark_helpers() base_rows = [ @@ -946,6 +1315,12 @@ def run_all(): ("Scoped token-budget SFT selection", test_sft_scoped_selection_only_trims_teacher_subset), ("Coverage-aware SFT selection", test_sft_coverage_selection_keeps_rare_high_quality_group), ("Coverage-aware preference selection", test_preference_coverage_margin_preserves_style_diversity), + ("Lazy-greedy SFT coverage", test_sft_lazy_greedy_coverage_avoids_redundant_clusters), + ("Lazy-greedy preference coverage", test_preference_lazy_greedy_coverage_avoids_redundant_styles), + ("Verifier-aware SFT selection", test_sft_selection_uses_verifier_reward_for_hard_math_code_rows), + ("Verifier-aware preference selection", test_preference_selection_uses_verifier_reward_margin), + ("Verifier metadata normalization", test_verifier_metadata_normalization_and_nested_jsonl_loading), + ("Verifier preference row fields", test_preference_rows_preserve_verifier_fields), ("Benchmark sample comparison", test_benchmark_sample_comparison_ranks_worst_regressions_first), ("Auto eval split grouping", test_split_train_eval_auto_groups_metadata_events), ("Chat-pair metadata roundtrip", test_save_jsonl_round_trips_pair_metadata), From c3d8153569583caf9f3d571b9d865850fb71471b Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:38:51 +0100 Subject: [PATCH 5/6] feat: unify adaptive routing across runtime surfaces --- chat_web_app.py | 549 +-- runtime_python/chat_app.py | 423 +- runtime_python/chat_web_app.py | 464 +- .../benchmark_auto_collective_loop_common.py | 295 ++ source/chat_app.py | 423 +- source/chat_image_variant_app.py | 15 +- source/chat_web_app.py | 321 +- source/multimodel_memory.py | 862 +++- source/multimodel_runtime.py | 3188 ++++++++++++- source/supermix_multimodel_web_app.py | 1862 +++++++- test_multimodel_memory_tools.py | 369 ++ test_multimodel_runtime.py | 4155 ++++++++++++++++- test_route_policy_shadow_web_status.py | 241 + test_runtime_compute_controls.py | 738 +++ test_supermix_multimodel_web_app.py | 1069 +++++ 15 files changed, 14042 insertions(+), 932 deletions(-) create mode 100644 source/benchmark_auto_collective_loop_common.py create mode 100644 test_route_policy_shadow_web_status.py create mode 100644 test_runtime_compute_controls.py diff --git a/chat_web_app.py b/chat_web_app.py index 699239154..fda53e06b 100644 --- a/chat_web_app.py +++ b/chat_web_app.py @@ -1,535 +1,48 @@ -import argparse -import json -import threading -import time -import uuid -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -from flask import Flask, jsonify, request -import torch - -import chat_app -from device_utils import configure_torch_runtime, resolve_device - - -HTML = """ -Champion AI • Premium Chat - - - - -
-
-
-
-

Champion AI

-
-

Premium retrieval interface with Gated Experts.

-
-
-
-
-
- - - -
-
Ready...
-
-
-
-

Chat Session

Initializing...
-
-
-
-
- - -
-
-
-""" - - - -class Engine: - def __init__(self, device: Any, device_info: Dict[str, Any], defaults: Dict[str, Any]): - self.device = device - self.device_info = dict(device_info or {}) - self.defaults = dict(defaults) - self.lock = threading.RLock() - self.model = None - self.weights_path: Optional[str] = None - self.meta_path: Optional[str] = None - self.feature_mode = "legacy" - self.model_size = "base" - self.buckets: Dict[int, List[Dict[str, Any]]] = {} - self.available_labels: List[int] = list(range(chat_app.MODEL_CLASSES)) - self.sessions: Dict[str, List[Tuple[str, str]]] = {} - self.recent: Dict[str, List[str]] = {} - - def status(self) -> Dict[str, Any]: - with self.lock: - return { - "loaded": self.model is not None, - "weights": self.weights_path, - "meta": self.meta_path, - "feature_mode": self.feature_mode, - "model_size": self.model_size, - "available_labels": len(self.available_labels), - "device": self.device_info.get("resolved", str(self.device)), - "sessions": len(self.sessions), - } - - def _parse_buckets(self, meta: Dict[str, Any]) -> None: - buckets: Dict[int, List[Dict[str, Any]]] = {} - raw = meta.get("buckets", {}) - if isinstance(raw, dict): - for k, v in raw.items(): - try: - label = int(k) - except Exception: - continue - if isinstance(v, list) and v: - buckets[label] = v - self.buckets = buckets - self.available_labels = sorted(buckets.keys()) or list(range(chat_app.MODEL_CLASSES)) - - def load(self, weights: str, meta_path: str) -> Dict[str, Any]: - t0 = time.perf_counter() - weights = str(Path(weights)) - meta_path = str(Path(meta_path)) - if not Path(weights).exists(): - raise FileNotFoundError(f"Weights not found: {weights}") - if not Path(meta_path).exists(): - raise FileNotFoundError(f"Metadata not found: {meta_path}") - - meta = chat_app.load_metadata(meta_path) - feature_mode = str(meta.get("feature_mode", "legacy")).strip().lower() - if feature_mode not in {"legacy", "context_v2", "context_v3", "context_v4", "context_mix_v1", "context_mix_v2_mm"}: - feature_mode = "legacy" - - sd = chat_app.safe_load_state_dict(weights) - inferred = chat_app.detect_model_size_from_state_dict(sd) - resolved_model_size = self.defaults.get("model_size", "auto") - if resolved_model_size == "auto": - resolved_model_size = inferred +"""Compatibility entrypoint for the canonical source chat web app. - expansion_dim = chat_app._resolve_expansion_dim(None, meta, "expansion_dim", 1024 if resolved_model_size in {"xxlarge","xxxlarge","ultralarge","megalarge"} else (768 if resolved_model_size == "xlarge" else 512), inferred, {"large","xlarge","xxlarge","xxxlarge","ultralarge","megalarge"}, chat_app.detect_large_head_expansion_dim, sd) - extra_expansion_dim = chat_app._resolve_expansion_dim(None, meta, "extra_expansion_dim", max(2048 if resolved_model_size in {"xxlarge","xxxlarge","ultralarge","megalarge"} else 1024, expansion_dim * 2), inferred, {"xlarge","xxlarge","xxxlarge","ultralarge","megalarge"}, chat_app.detect_xlarge_aux_expansion_dim, sd) - third_expansion_dim = chat_app._resolve_expansion_dim(None, meta, "third_expansion_dim", max(3072, extra_expansion_dim + expansion_dim), inferred, {"xxlarge","xxxlarge","ultralarge","megalarge"}, chat_app.detect_xxlarge_third_expansion_dim, sd) - fourth_expansion_dim = chat_app._resolve_expansion_dim(None, meta, "fourth_expansion_dim", max(4096, third_expansion_dim + expansion_dim), inferred, {"xxxlarge","ultralarge","megalarge"}, chat_app.detect_xxxlarge_fourth_expansion_dim, sd) - fifth_expansion_dim = chat_app._resolve_expansion_dim(None, meta, "fifth_expansion_dim", max(6144, fourth_expansion_dim + expansion_dim), inferred, {"ultralarge","megalarge"}, chat_app.detect_ultralarge_fifth_expansion_dim, sd) - sixth_expansion_dim = chat_app._resolve_expansion_dim(None, meta, "sixth_expansion_dim", max(8192, fifth_expansion_dim + expansion_dim), inferred, {"megalarge"}, chat_app.detect_megalarge_sixth_expansion_dim, sd) - adapter_dropout = float(meta.get("adapter_dropout", 0.1)) +The implementation lives in ``source/chat_web_app.py``. Keeping this thin +wrapper at the repository root prevents imports from binding to an older copy +when tests or scripts import ``chat_web_app`` before adding ``source`` to +``sys.path``. +""" - model = chat_app.build_model( - model_size=resolved_model_size, - expansion_dim=expansion_dim, - dropout=adapter_dropout, - extra_expansion_dim=extra_expansion_dim, - third_expansion_dim=third_expansion_dim, - fourth_expansion_dim=fourth_expansion_dim, - fifth_expansion_dim=fifth_expansion_dim, - sixth_expansion_dim=sixth_expansion_dim, - ).to(self.device).eval() - missing, unexpected = chat_app.load_weights_for_model(model, sd, model_size=resolved_model_size) - if missing or unexpected: - raise RuntimeError(f"State dict mismatch. Missing={missing}, Unexpected={unexpected}") +from __future__ import annotations - with self.lock: - self.model = model - self.weights_path = weights - self.meta_path = meta_path - self.feature_mode = feature_mode - self.model_size = resolved_model_size - self._parse_buckets(meta) - self.sessions.clear() - self.recent.clear() - - return {"ok": True, "load_ms": round((time.perf_counter()-t0)*1000,1), **self.status()} - - def clear(self, session_id: str) -> None: - with self.lock: - self.sessions.pop(session_id, None) - self.recent.pop(session_id, None) - - def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None, response_temperature: Optional[float] = None, show_top_responses: int = 0) -> Dict[str, Any]: - if not user_text.strip(): - raise ValueError("Empty message") - with self.lock: - if self.model is None: - raise RuntimeError("No model loaded") - model = self.model - feature_mode = self.feature_mode - buckets = self.buckets - labels = list(self.available_labels) - history = list(self.sessions.get(session_id, [])) - recent_msgs = list(self.recent.get(session_id, [])) - t0 = time.perf_counter() - t_infer = 0.0 - t_rank = 0.0 - - context = chat_app.build_context(history, user_text=user_text, max_turns=int(self.defaults.get("max_turns", 2))) - tt = time.perf_counter() - x = chat_app.text_to_model_input(context, feature_mode=feature_mode).to(self.device) - with torch.no_grad(): - logits = model(x)[0, 0] - t_infer += time.perf_counter() - tt - - idx = torch.tensor(labels, dtype=torch.long, device=logits.device) - avail_logits = logits.index_select(0, idx) - probs = torch.softmax(avail_logits, dim=0) - pool_mode = str(self.defaults.get("pool_mode", "all")) - if pool_mode == "all": - top_pos = list(range(len(labels))) - else: - k = max(1, min(int(self.defaults.get("top_labels", 3)), len(labels))) - top_pos = torch.topk(avail_logits, k=k).indices.tolist() - - pooled: List[Dict[str, Any]] = [] - for pos in top_pos: - label = labels[int(pos)] - bucket_score = float(probs[int(pos)].item()) - for row in buckets.get(label, []): - m = dict(row) - m["bucket_score"] = bucket_score - m["_source"] = "model" - pooled.append(m) - if (not pooled) and buckets: - label = chat_app.choose_bucket_from_logits(logits, labels, temperature=float(self.defaults.get("temperature", 0.0))) - pooled = list(buckets.get(label, [])) - - dedup: Dict[str, Dict[str, Any]] = {} - for row in pooled: - text = str(row.get("text", "")).strip() - if not text: - continue - prev = dedup.get(text) - if prev is None: - d = dict(row) - d["_sources_set"] = {row.get("_source", "unknown")} - dedup[text] = d - continue - src = row.get("_source", "unknown") - base = max(float(prev.get("bucket_score", 0.0)), float(row.get("bucket_score", 0.0))) - if src not in prev["_sources_set"]: - base *= 1.10 - prev["_sources_set"].add(src) - prev["bucket_score"] = base - prev["count"] = int(prev.get("count", 1)) + int(row.get("count", 1)) - for k in list(dedup): - dedup[k].pop("_sources_set", None) - dedup[k].pop("_source", None) - pooled = list(dedup.values()) - - resolved_style = chat_app.infer_style_mode(user_text, requested_mode=style_mode or str(self.defaults.get("style_mode", "auto"))) - top_candidates: List[Dict[str, Any]] = [] - show_n = max(0, int(show_top_responses)) - if show_n > 0 and pooled: - tt = time.perf_counter() - ranked, scores = chat_app.rank_response_candidates(pooled, query_text=user_text, recent_assistant_messages=recent_msgs, style_mode=resolved_style) - t_rank += time.perf_counter() - tt - shown = 0 - for ridx in ranked: - txt = str(pooled[ridx].get("text", "")).strip() - if not txt: - continue - top_candidates.append({"score": float(scores[ridx].item()), "text": txt}) - shown += 1 - if shown >= show_n: - break - - tt = time.perf_counter() - resp = chat_app.pick_response( - pooled, - query_text=user_text, - recent_assistant_messages=recent_msgs, - response_temperature=float(self.defaults.get("response_temperature", 0.08) if response_temperature is None else response_temperature), - style_mode=resolved_style, - creativity=max(0.0, min(1.0, float(self.defaults.get("creativity", 0.2)))), - ) - t_rank += time.perf_counter() - tt - resp = chat_app.cleanup_response_text(resp) or "I do not have a trained response for that yet." - - with self.lock: - hist = self.sessions.setdefault(session_id, []) - hist.append((user_text, resp)) - if len(hist) > 40: - del hist[:-40] - recent = self.recent.setdefault(session_id, []) - recent.append(resp) - if len(recent) > 24: - del recent[:-24] - - return { - "ok": True, - "session_id": session_id, - "response": resp, - "style_mode": resolved_style, - "timing_ms": { - "infer": round(t_infer * 1000, 1), - "rank_pick": round(t_rank * 1000, 1), - "total": round((time.perf_counter() - t0) * 1000, 1), - }, - "top_candidates": top_candidates, - } +import importlib.util +import sys +from pathlib import Path -def build_app(engine: Engine, default_weights: str, default_meta: str): - app = Flask(__name__) +_ROOT = Path(__file__).resolve().parent +_SOURCE_DIR = _ROOT / "source" +_IMPL_PATH = _SOURCE_DIR / "chat_web_app.py" - @app.get('/') - def index(): - html = HTML.replace("", f"") - html = html.replace("", f"") - return html +if str(_SOURCE_DIR) not in sys.path: + sys.path.insert(0, str(_SOURCE_DIR)) - @app.get('/api/status') - def api_status(): - return jsonify({"ok": True, "status": engine.status()}) +_SPEC = importlib.util.spec_from_file_location("_supermix_source_chat_web_app", _IMPL_PATH) +if _SPEC is None or _SPEC.loader is None: + raise ImportError(f"Cannot load chat web app implementation from {_IMPL_PATH}") - @app.post('/api/load') - def api_load(): - p = request.get_json(force=True, silent=True) or {} - try: - return jsonify(engine.load(str(p.get('weights') or '').strip(), str(p.get('meta') or '').strip())) - except FileNotFoundError as e: - return jsonify({"ok": False, "error": str(e)}), 404 - except Exception as e: - return jsonify({"ok": False, "error": str(e)}), 400 +_IMPL = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _IMPL +_SPEC.loader.exec_module(_IMPL) - @app.post('/api/chat') - def api_chat(): - p = request.get_json(force=True, silent=True) or {} - sid = str(p.get('session_id') or '').strip() or str(uuid.uuid4()) - msg = str(p.get('message') or '').strip() - try: - return jsonify(engine.chat( - session_id=sid, - user_text=msg, - style_mode=p.get('style_mode'), - response_temperature=p.get('response_temperature'), - show_top_responses=int(p.get('show_top_responses') or 0), - )) - except Exception as e: - return jsonify({"ok": False, "error": str(e)}), 400 +__all__ = [name for name in vars(_IMPL) if not name.startswith("_")] - @app.post('/api/clear') - def api_clear(): - p = request.get_json(force=True, silent=True) or {} - sid = str(p.get('session_id') or '').strip() - if not sid: - return jsonify({"ok": False, "error": "session_id required"}), 400 - engine.clear(sid) - return jsonify({"ok": True, "cleared": True, "session_id": sid}) +for _name in __all__: + globals()[_name] = getattr(_IMPL, _name) - return app +# Keep the root compatibility surface behaviorally identical for callers that +# validate CLI default inheritance after another test imported this wrapper +# first. The canonical helper remains private in the source implementation. +for _name in ("_runtime_compute_cli_overrides",): + globals()[_name] = getattr(_IMPL, _name) def main() -> None: - ap = argparse.ArgumentParser(description='Web interface for Champion chat model (loads local weights + metadata).') - ap.add_argument('--weights', default='champion_model_chat_supermix_v27_500k_ft.pth') - ap.add_argument('--meta', default='chat_model_meta_supermix_v27_500k.json') - ap.add_argument('--autoload', action='store_true') - ap.add_argument('--host', default='127.0.0.1') - ap.add_argument('--port', type=int, default=8000) - ap.add_argument('--device', default='auto') - ap.add_argument('--device_preference', default='cuda,npu,xpu,dml,mps,cpu') - ap.add_argument('--torch_num_threads', type=int, default=0) - ap.add_argument('--torch_interop_threads', type=int, default=0) - ap.add_argument('--matmul_precision', choices=['highest', 'high', 'medium'], default='high') - ap.add_argument('--disable_tf32', action='store_true') - ap.add_argument('--model_size', choices=['auto','base','large','xlarge','xxlarge','xxxlarge','ultralarge','megalarge'], default='auto') - ap.add_argument('--max_turns', type=int, default=2) - ap.add_argument('--top_labels', type=int, default=3) - ap.add_argument('--pool_mode', choices=['all','topk'], default='all') - ap.add_argument('--response_temperature', type=float, default=0.08) - ap.add_argument('--temperature', type=float, default=0.0) - ap.add_argument('--style_mode', choices=['auto','balanced','creative','concise','analyst'], default='auto') - ap.add_argument('--creativity', type=float, default=0.2) - args = ap.parse_args() - - configure_torch_runtime( - torch_num_threads=int(args.torch_num_threads), - torch_interop_threads=int(args.torch_interop_threads), - allow_tf32=not bool(args.disable_tf32), - matmul_precision=str(args.matmul_precision), - ) - device, device_info = resolve_device(args.device, preference=args.device_preference) - engine = Engine(device, device_info, { - 'model_size': args.model_size, - 'max_turns': int(args.max_turns), - 'top_labels': int(args.top_labels), - 'pool_mode': str(args.pool_mode), - 'response_temperature': float(args.response_temperature), - 'temperature': float(args.temperature), - 'style_mode': str(args.style_mode), - 'creativity': float(args.creativity), - }) - if args.autoload: - try: - print(engine.load(args.weights, args.meta)) - except Exception as e: - print(f'Autoload failed: {e}') - app = build_app(engine, str(args.weights), str(args.meta)) - print(f'Web UI: http://{args.host}:{args.port}') - app.run(host=args.host, port=int(args.port), threaded=True) + _IMPL.main() -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/runtime_python/chat_app.py b/runtime_python/chat_app.py index c42331cba..58d4d3027 100644 --- a/runtime_python/chat_app.py +++ b/runtime_python/chat_app.py @@ -1,9 +1,10 @@ import argparse +import inspect import json import re import time from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from concurrent.futures import ThreadPoolExecutor import torch @@ -71,6 +72,11 @@ def _int_or_default(value, default: int) -> int: _FOLLOWUP_QUERY_RE = re.compile(r"\b(it|that|this|they|them|same|previous|above|more|continue|deeper|expand)\b", re.I) VALID_RUNTIME_MODEL_SIZES = SUPPORTED_MODEL_SIZES +MAX_RUNTIME_REASONING_CYCLES = 64 +DEFAULT_ADAPTIVE_EXIT_ENTROPY = 0.2 +DEFAULT_PREDICTION_STABILITY_PATIENCE = 2 +DEFAULT_PREDICTION_STABILITY_TOL = 5e-3 +AUTO_REASONING_CYCLE_BUCKETS = (1, 3, 8, 16) def _build_db_query(user: str, history: List[Tuple[str, str]], memory_rows: List[Dict], max_turns: int = 2) -> str: @@ -145,6 +151,257 @@ def resolve_runtime_model_size(requested_model_size: str, meta_model_size: str, return resolved_model_size, warning +def _coerce_optional_positive_int(value: Any, max_value: Optional[int] = None) -> Optional[int]: + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + if stripped in {"", "default", "none", "off", "auto"}: + return None + value = stripped + try: + parsed = int(value) + except (TypeError, ValueError): + return None + parsed = max(1, parsed) + if max_value is not None: + parsed = min(parsed, max(1, int(max_value))) + return parsed + + +def _is_auto_reasoning_cycles(value: Any) -> bool: + return isinstance(value, str) and value.strip().lower() in {"auto", "adaptive", "smart"} + + +def _format_reasoning_cycles_setting(value: Any) -> Any: + if _is_auto_reasoning_cycles(value): + return "auto" + return _coerce_optional_positive_int(value, MAX_RUNTIME_REASONING_CYCLES) or "default" + + +def estimate_auto_reasoning_cycles(context: str, max_value: Optional[int] = None) -> Dict[str, Any]: + text = str(context or "") + lowered = text.lower() + reasons: List[str] = [] + score = 0 + + char_count = len(text) + if char_count >= 1600: + score += 3 + reasons.append("long_context") + elif char_count >= 700: + score += 2 + reasons.append("medium_context") + elif char_count >= 280: + score += 1 + reasons.append("short_context") + + if re.search(r"```|traceback|exception|error|bug|fix|debug|refactor|class\s+|def\s+|function|import\s+|select\s+", lowered): + score += 2 + reasons.append("code_or_debug") + if re.search(r"\b(prove|derive|equation|calculate|solve|theorem|complexity|probability|optimi[sz]e|benchmarks?)\b|[=+\-*/^]{2,}", lowered): + score += 2 + reasons.append("math_or_verification") + if re.search(r"\b(compare|tradeoff|plan|architect|design|multi-step|step by step|analy[sz]e|research|evidence)\b", lowered): + score += 1 + reasons.append("deliberation") + if re.search(r"\b(agent|tool|browser|github|file|tests?|verify|integration|runtime)\b", lowered): + score += 1 + reasons.append("tool_or_runtime") + if re.search(r"\b(quick|brief|concise|short answer|one line)\b", lowered): + score -= 1 + reasons.append("brevity_requested") + + if score <= 0: + cycles = AUTO_REASONING_CYCLE_BUCKETS[0] + elif score <= 2: + cycles = AUTO_REASONING_CYCLE_BUCKETS[1] + elif score <= 4: + cycles = AUTO_REASONING_CYCLE_BUCKETS[2] + else: + cycles = AUTO_REASONING_CYCLE_BUCKETS[3] + cycles = _coerce_optional_positive_int(cycles, max_value) or AUTO_REASONING_CYCLE_BUCKETS[0] + + return { + "mode": "auto", + "score": int(score), + "reasons": reasons or ["simple_prompt"], + "cycles": int(cycles), + } + + +def _coerce_nonnegative_float(value: Any, default: float, max_value: Optional[float] = None) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + parsed = float(default) + parsed = max(0.0, parsed) + if max_value is not None: + parsed = min(parsed, float(max_value)) + return parsed + + +def _coerce_nonnegative_int(value: Any, default: int, max_value: Optional[int] = None) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + parsed = int(default) + parsed = max(0, parsed) + if max_value is not None: + parsed = min(parsed, max(0, int(max_value))) + return parsed + + +def _coerce_bool(value: Any) -> bool: + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on", "y"} + return bool(value) + + +def _model_forward_accepts_kwarg(model: Any, kwarg: str) -> bool: + forward = getattr(model, "forward", None) + if forward is None: + return False + try: + params = inspect.signature(forward).parameters.values() + except (TypeError, ValueError): + return False + for param in params: + if param.kind == inspect.Parameter.VAR_KEYWORD or param.name == kwarg: + return True + return False + + +def model_supports_runtime_compute(model: Any) -> bool: + return ( + _model_forward_accepts_kwarg(model, "reasoning_cycles") + or _model_forward_accepts_kwarg(model, "adaptive_compute") + ) + + +def _to_optional_scalar(value: Any) -> Optional[float]: + if value is None: + return None + try: + if hasattr(value, "detach"): + return float(value.detach().cpu().item()) + return float(value) + except (TypeError, ValueError, RuntimeError): + return None + + +def collect_runtime_compute_metrics(model: Any) -> Dict[str, Any]: + metrics: Dict[str, Any] = { + "cycles_used": None, + "ponder_cost": None, + "consistency_loss": None, + "gating_entropy": None, + "prediction_streak": None, + "prediction_confidence_delta": None, + "exit_reason": None, + } + attr_map = { + "last_cycles_used": "cycles_used", + "last_ponder_cost": "ponder_cost", + "last_consistency_loss": "consistency_loss", + "last_gating_entropy": "gating_entropy", + "last_prediction_streak": "prediction_streak", + "last_prediction_confidence_delta": "prediction_confidence_delta", + } + modules = list(model.modules()) if hasattr(model, "modules") else [model] + for module in modules: + for attr, key in attr_map.items(): + if metrics[key] is not None: + continue + scalar = _to_optional_scalar(getattr(module, attr, None)) + if scalar is not None: + metrics[key] = scalar + if metrics["exit_reason"] is None: + reason = getattr(module, "last_exit_reason", None) + if isinstance(reason, str) and reason.strip(): + metrics["exit_reason"] = reason.strip() + return metrics + + +def get_last_cycles_used(model: Any) -> Optional[float]: + return collect_runtime_compute_metrics(model).get("cycles_used") + + +def forward_with_runtime_compute( + model: Any, + x: torch.Tensor, + reasoning_cycles: Any = None, + adaptive_compute: Any = False, + exit_tol: Any = 1e-3, + exit_entropy_threshold: Any = DEFAULT_ADAPTIVE_EXIT_ENTROPY, + prediction_stability_patience: Any = DEFAULT_PREDICTION_STABILITY_PATIENCE, + prediction_stability_tol: Any = DEFAULT_PREDICTION_STABILITY_TOL, + auto_reasoning_context: str = "", +) -> Tuple[torch.Tensor, Dict[str, Any]]: + """Forward a model while applying optional v50 runtime-compute controls. + + Legacy model variants ignore these controls because their forward signatures + do not accept them. This keeps the runtime knobs safe for mixed checkpoints. + """ + auto_policy: Optional[Dict[str, Any]] = None + if _is_auto_reasoning_cycles(reasoning_cycles): + auto_policy = estimate_auto_reasoning_cycles( + auto_reasoning_context, + MAX_RUNTIME_REASONING_CYCLES, + ) + cycles = int(auto_policy["cycles"]) + else: + cycles = _coerce_optional_positive_int(reasoning_cycles, MAX_RUNTIME_REASONING_CYCLES) + adaptive = _coerce_bool(adaptive_compute) + tol = _coerce_nonnegative_float(exit_tol, 1e-3) + entropy_threshold = _coerce_nonnegative_float( + exit_entropy_threshold, + DEFAULT_ADAPTIVE_EXIT_ENTROPY, + ) + stability_patience = _coerce_nonnegative_int( + prediction_stability_patience, + DEFAULT_PREDICTION_STABILITY_PATIENCE, + MAX_RUNTIME_REASONING_CYCLES, + ) + stability_tol = _coerce_nonnegative_float( + prediction_stability_tol, + DEFAULT_PREDICTION_STABILITY_TOL, + ) + kwargs: Dict[str, Any] = {} + + if cycles is not None and _model_forward_accepts_kwarg(model, "reasoning_cycles"): + kwargs["reasoning_cycles"] = cycles + if adaptive and _model_forward_accepts_kwarg(model, "adaptive_compute"): + kwargs["adaptive_compute"] = True + if _model_forward_accepts_kwarg(model, "exit_tol"): + kwargs["exit_tol"] = tol + if _model_forward_accepts_kwarg(model, "exit_entropy_threshold"): + kwargs["exit_entropy_threshold"] = entropy_threshold + if _model_forward_accepts_kwarg(model, "prediction_stability_patience"): + kwargs["prediction_stability_patience"] = stability_patience + if _model_forward_accepts_kwarg(model, "prediction_stability_tol"): + kwargs["prediction_stability_tol"] = stability_tol + + output = model(x, **kwargs) if kwargs else model(x) + metrics = collect_runtime_compute_metrics(model) + diagnostics = { + "requested_reasoning_cycles": cycles, + "selected_reasoning_cycles": cycles, + "reasoning_budget_mode": "auto" if auto_policy is not None else ("manual" if cycles is not None else "default"), + "auto_reasoning_policy": auto_policy, + "max_reasoning_cycles": MAX_RUNTIME_REASONING_CYCLES, + "adaptive_compute": adaptive, + "exit_tol": tol, + "exit_entropy_threshold": entropy_threshold, + "prediction_stability_patience": stability_patience, + "prediction_stability_tol": stability_tol, + "applied": bool(kwargs), + "applied_kwargs": dict(kwargs), + **metrics, + } + return output, diagnostics + + def _default_expansion_dim_for_model_size(model_size: str) -> int: if model_size in THIRD_EXPANSION_DIM_MODEL_SIZES: return 1024 @@ -180,6 +437,12 @@ def _print_chat_help() -> None: print(" /clear Clear in-memory conversation history") print(" /style Set style: auto|balanced|creative|concise|analyst") print(" /creativity <0-1> Set creative rewrite strength") + print(" /cycles Set runtime reasoning cycles for future turns") + print(" /adaptive on|off Toggle v50 convergence early-exit") + print(" /exit_tol Set v50 adaptive-compute convergence tolerance") + print(" /exit_entropy Set v50 adaptive-compute entropy exit threshold") + print(" /stability Stable prediction cycles required (0 disables)") + print(" /stability_tol Maximum confidence drift across the stable streak") print(" /top Show top reranked candidates each turn (0 disables)") print(" /timing on|off Toggle per-turn timing output") print(" /memory on|off Toggle memory retrieval/writes for this session") @@ -204,6 +467,41 @@ def main(): ap.add_argument("--fifth_expansion_dim", type=int, default=None) ap.add_argument("--sixth_expansion_dim", type=int, default=None) ap.add_argument("--adapter_dropout", type=float, default=None) + ap.add_argument( + "--reasoning_cycles", + type=str, + default=None, + help="Optional runtime reasoning cycles, or 'auto' to choose from prompt difficulty; omit for checkpoint default.", + ) + ap.add_argument( + "--adaptive_compute", + action="store_true", + help="Enable v50 convergence early-exit during inference when the model supports it.", + ) + ap.add_argument( + "--adaptive_exit_tol", + type=float, + default=1e-3, + help="Convergence tolerance for v50 adaptive compute.", + ) + ap.add_argument( + "--adaptive_exit_entropy", + type=float, + default=DEFAULT_ADAPTIVE_EXIT_ENTROPY, + help="Prediction entropy threshold for v50 adaptive compute.", + ) + ap.add_argument( + "--prediction_stability_patience", + type=int, + default=DEFAULT_PREDICTION_STABILITY_PATIENCE, + help="Stable full-output predictions required before adaptive exit (0 disables).", + ) + ap.add_argument( + "--prediction_stability_tol", + type=float, + default=DEFAULT_PREDICTION_STABILITY_TOL, + help="Maximum confidence drift allowed across a stable prediction streak.", + ) ap.add_argument("--device", default="auto") ap.add_argument( "--device_preference", @@ -412,11 +710,22 @@ def main(): session_started_at = time.time() turn_count = 0 last_turn_timing: Dict[str, float] = {} + last_compute_diag: Dict[str, Any] = {} print(f"\n{TerminalColors.SYSTEM}--- Session Info ---") print(f"Loaded: {Path(args.weights).name} [{resolved_model_size}] | Available labels: {len(available_labels)}") print(f"Device: {device_info.get('resolved', args.device)} | Threads intra={torch.get_num_threads()} interop={torch.get_num_interop_threads()}") print(f"Feature mode: {feature_mode}{feature_mode_note} | Style mode: {args.style_mode} (creativity={max(0.0, min(1.0, float(args.creativity))):.2f})") + compute_support = "yes" if model_supports_runtime_compute(model) else "no" + print( + "Test-time compute: " + f"supported={compute_support} cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) if llm_db: print(f"LLM DB: {args.llm_db} (top_k={args.db_top_k})") @@ -462,12 +771,31 @@ def main(): f" style={args.style_mode} creativity={float(args.creativity):.2f} " f"top_debug={int(args.show_top_responses)} timing={'on' if show_timing else 'off'}" ) + print( + f" compute_supported={'yes' if model_supports_runtime_compute(model) else 'no'} " + f"cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) print( f" memory={'on' if session_memory_enabled else 'off'} " f"(db={'ready' if memory_db is not None else 'missing'}) | " f"llm_db={'on' if session_db_enabled else 'off'} " f"(db={'ready' if llm_db is not None else 'missing'})" ) + if last_compute_diag: + print( + " last_compute=" + f"applied={last_compute_diag.get('applied')} " + f"cycles_used={last_compute_diag.get('cycles_used')} " + f"ponder={last_compute_diag.get('ponder_cost')} " + f"gate_entropy={last_compute_diag.get('gating_entropy')} " + f"exit_reason={last_compute_diag.get('exit_reason')} " + f"prediction_streak={last_compute_diag.get('prediction_streak')}" + ) if last_turn_timing: print( " last_timing=" @@ -483,6 +811,15 @@ def main(): f" pool_mode={args.pool_mode} top_labels={int(args.top_labels)} " f"show_top_responses={int(args.show_top_responses)}" ) + print( + f" compute_supported={'yes' if model_supports_runtime_compute(model) else 'no'} " + f"cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) print( f" memory={'on' if session_memory_enabled else 'off'} top_k={int(args.memory_top_k)} " f"pool={int(args.memory_pool_size)} score_scale={float(args.memory_score_scale):.3f}" @@ -503,6 +840,68 @@ def main(): print(f"{TerminalColors.SYSTEM}Creativity set to {float(args.creativity):.2f}.{TerminalColors.RESET}") except Exception: print(f"{TerminalColors.SYSTEM}Usage: /creativity 0.0-1.0{TerminalColors.RESET}") + elif cmd == "cycles": + raw_cycles = arg.strip().lower() + if raw_cycles in {"auto", "adaptive", "smart"}: + args.reasoning_cycles = "auto" + print(f"{TerminalColors.SYSTEM}Reasoning cycles set to auto.{TerminalColors.RESET}") + else: + cycles = _coerce_optional_positive_int(arg) + if cycles is None and raw_cycles not in {"", "default", "off", "none"}: + print(f"{TerminalColors.SYSTEM}Usage: /cycles {TerminalColors.RESET}") + continue + args.reasoning_cycles = cycles + label = cycles if cycles is not None else "default" + print(f"{TerminalColors.SYSTEM}Reasoning cycles set to {label}.{TerminalColors.RESET}") + elif cmd == "adaptive": + if arg.lower() in {"on", "1", "true"}: + args.adaptive_compute = True + elif arg.lower() in {"off", "0", "false"}: + args.adaptive_compute = False + else: + print(f"{TerminalColors.SYSTEM}Usage: /adaptive on|off{TerminalColors.RESET}") + continue + print(f"{TerminalColors.SYSTEM}Adaptive compute {'enabled' if args.adaptive_compute else 'disabled'}.{TerminalColors.RESET}") + elif cmd == "exit_tol": + try: + args.adaptive_exit_tol = _coerce_nonnegative_float(arg, args.adaptive_exit_tol) + print(f"{TerminalColors.SYSTEM}Adaptive exit tolerance set to {float(args.adaptive_exit_tol):.4g}.{TerminalColors.RESET}") + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /exit_tol {TerminalColors.RESET}") + elif cmd == "exit_entropy": + try: + args.adaptive_exit_entropy = _coerce_nonnegative_float( + arg, + args.adaptive_exit_entropy, + ) + print(f"{TerminalColors.SYSTEM}Adaptive exit entropy set to {float(args.adaptive_exit_entropy):.4g}.{TerminalColors.RESET}") + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /exit_entropy {TerminalColors.RESET}") + elif cmd == "stability": + try: + args.prediction_stability_patience = _coerce_nonnegative_int( + arg, + args.prediction_stability_patience, + MAX_RUNTIME_REASONING_CYCLES, + ) + print( + f"{TerminalColors.SYSTEM}Prediction stability patience set to " + f"{int(args.prediction_stability_patience)}.{TerminalColors.RESET}" + ) + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /stability {TerminalColors.RESET}") + elif cmd == "stability_tol": + try: + args.prediction_stability_tol = _coerce_nonnegative_float( + arg, + args.prediction_stability_tol, + ) + print( + f"{TerminalColors.SYSTEM}Prediction stability tolerance set to " + f"{float(args.prediction_stability_tol):.4g}.{TerminalColors.RESET}" + ) + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /stability_tol {TerminalColors.RESET}") elif cmd == "top": try: args.show_top_responses = max(0, int(arg)) @@ -585,7 +984,19 @@ def main(): x = text_to_model_input(context, feature_mode=feature_mode).to(device) with torch.no_grad(): - logits = model(x)[0, 0] # (10,) + model_out, compute_diag = forward_with_runtime_compute( + model, + x, + reasoning_cycles=args.reasoning_cycles, + adaptive_compute=args.adaptive_compute, + exit_tol=args.adaptive_exit_tol, + exit_entropy_threshold=args.adaptive_exit_entropy, + prediction_stability_patience=args.prediction_stability_patience, + prediction_stability_tol=args.prediction_stability_tol, + auto_reasoning_context=context, + ) + logits = model_out[0, 0] # (10,) + last_compute_diag = compute_diag t_infer += max(0.0, time.perf_counter() - _t) idx = torch.tensor(available_labels, dtype=torch.long, device=logits.device) @@ -732,6 +1143,14 @@ def main(): f"{TerminalColors.SYSTEM}Timing:{TerminalColors.RESET} " + ", ".join(f"{k}={_format_ms(v)}" for k, v in last_turn_timing.items()) ) + if last_compute_diag: + print( + f"{TerminalColors.SYSTEM}Compute:{TerminalColors.RESET} " + f"applied={last_compute_diag.get('applied')} " + f"cycles_used={last_compute_diag.get('cycles_used')} " + f"ponder={last_compute_diag.get('ponder_cost')} " + f"gate_entropy={last_compute_diag.get('gating_entropy')}" + ) finally: executor.shutdown(wait=False) diff --git a/runtime_python/chat_web_app.py b/runtime_python/chat_web_app.py index 1849dabd2..b73932839 100644 --- a/runtime_python/chat_web_app.py +++ b/runtime_python/chat_web_app.py @@ -3,6 +3,7 @@ import threading import time import uuid +from html import escape from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -12,6 +13,14 @@ import chat_app from device_utils import configure_torch_runtime, resolve_device +try: + from route_policy_shadow_registry import RouteShadowAssignmentRegistry +except ImportError: # repository compatibility when only runtime_python is first on sys.path + try: + from source.route_policy_shadow_registry import RouteShadowAssignmentRegistry + except ImportError: # standalone legacy bundles may omit the shadow console modules + RouteShadowAssignmentRegistry = None # type: ignore[assignment,misc] + HTML = """ Supermix v27 Intelligence @@ -179,6 +188,10 @@ } .btns { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; margin-top: 10px; } .btn-full { grid-column: span 2; } +.shadow-registry { margin-top:12px; padding:12px; border:1px solid rgba(45,212,191,.2); border-radius:12px; background:rgba(0,0,0,.16); } +.shadow-registry-head { display:flex; align-items:center; justify-content:space-between; gap:8px; color:var(--text); font-size:.72rem; font-weight:700; text-transform:uppercase; letter-spacing:.05em; } +.shadow-registry-head button { padding:7px 10px; font-size:.7rem; } +.shadow-registry-status { margin-top:8px; color:var(--text-dim); font:11px/1.45 'Courier New',monospace; white-space:pre-wrap; overflow-wrap:anywhere; } @media (max-width: 900px) { .wrap { grid-template-columns: 1fr; height: auto; overflow: visible; } @@ -203,6 +216,18 @@
+
+
+
+
+
+
+
+
+
+
+
+
@@ -211,6 +236,10 @@
System idle.
+
+
Shadow registry - read only
+
Not loaded. Browser mutation, execution, activation, and promotion are unavailable.
+
@@ -222,6 +251,7 @@
+
@@ -236,22 +266,67 @@ promptEl.style.height = (promptEl.scrollHeight) + 'px'; }); -function add(kind,text,timing,top){ - const d=document.createElement('div'); - d.className='msg '+kind; - let html = `
${kind==='user'?'Human':'Supermix'}
${text}
`; - +function fmtNum(value,digits=3){const n=Number(value);return Number.isFinite(n)?n.toFixed(digits):null;} +function reasoningCyclesValue(){const raw=el('reasoningCycles').value.trim();if(!raw)return null;const low=raw.toLowerCase();if(['auto','adaptive','smart'].includes(low))return 'auto';const n=Number(raw);return Number.isFinite(n)?n:raw;} +function add(kind,text,timing,top,compute){ + const d=document.createElement('div'); + d.className='msg '+kind; + const who=document.createElement('div'); + who.className='who'; + who.textContent=kind==='user'?'Human':'Supermix'; + const body=document.createElement('div'); + body.style.whiteSpace='pre-wrap'; + body.textContent=text; + d.appendChild(who); + d.appendChild(body); + if(timing){ - html += `
Engine Latency: ${timing.total}ms | Infer: ${timing.infer}ms
`; + const cycles = timing.cycles_used !== undefined && timing.cycles_used !== null ? ` | Cycles: ${timing.cycles_used}` : ''; + const t=document.createElement('div'); + t.className='tim'; + t.textContent=`Engine Latency: ${timing.total}ms | Infer: ${timing.infer}ms${cycles}`; + d.appendChild(t); + } + if(compute&&compute.applied){ + const parts=[]; + if(compute.reasoning_budget_mode==='auto') parts.push('mode auto'); + if(compute.requested_reasoning_cycles!==undefined&&compute.requested_reasoning_cycles!==null) parts.push(`requested ${compute.requested_reasoning_cycles}`); + if(compute.cycles_used!==undefined&&compute.cycles_used!==null) parts.push(`used ${compute.cycles_used}`); + if(compute.exit_reason) parts.push(`exit ${compute.exit_reason}`); + const streak=fmtNum(compute.prediction_streak); if(streak) parts.push(`stable ${streak}`); + const drift=fmtNum(compute.prediction_confidence_delta); if(drift) parts.push(`drift ${drift}`); + const ponder=fmtNum(compute.ponder_cost); if(ponder) parts.push(`ponder ${ponder}`); + const consistency=fmtNum(compute.consistency_loss); if(consistency) parts.push(`consistency ${consistency}`); + const entropy=fmtNum(compute.gating_entropy); if(entropy) parts.push(`gate entropy ${entropy}`); + const exitEntropy=fmtNum(compute.exit_entropy_threshold); if(exitEntropy) parts.push(`exit entropy ${exitEntropy}`); + if(compute.auto_reasoning_policy&&Array.isArray(compute.auto_reasoning_policy.reasons)) parts.push(`policy ${compute.auto_reasoning_policy.reasons.slice(0,3).join(',')}`); + if(parts.length){ + const c=document.createElement('div'); + c.className='tim'; + c.textContent='Compute: '+parts.join(' | '); + d.appendChild(c); + } } if(top&&top.length){ - html += `
- Neural Probabilities:
${top.map((c,i)=>`${i+1}. ${c.text.slice(0,100)}... (${(c.score*100).toFixed(1)}%)`).join('
')} -
`; + const x=document.createElement('div'); + x.className='tim'; + x.style.borderTop='1px solid rgba(255,255,255,0.05)'; + x.style.paddingTop='8px'; + x.style.marginTop='8px'; + const title=document.createElement('b'); + title.textContent='Neural Probabilities:'; + x.appendChild(title); + top.forEach((c,i)=>{ + const row=document.createElement('div'); + const score=Number(c.score); + const scoreText=Number.isFinite(score)?(score*100).toFixed(1):'n/a'; + row.textContent=`${i+1}. ${String(c.text||'').slice(0,100)}... (${scoreText}%)`; + x.appendChild(row); + }); + d.appendChild(x); } - d.innerHTML = html; - msgs.appendChild(d); - msgs.scrollTop=msgs.scrollHeight; + msgs.appendChild(d); + msgs.scrollTop=msgs.scrollHeight; } async function jget(path){ @@ -273,8 +348,37 @@ el('metaLine').textContent = d.status.loaded ? `${d.status.model_size.toUpperCase()} CORE | ${d.status.device.toUpperCase()}` : 'ENGINE OFFLINE'; if(!el('weights').value&&d.status.weights) el('weights').value=d.status.weights; if(!el('meta').value&&d.status.meta) el('meta').value=d.status.meta; + if(!el('reasoningCycles').value&&d.status.reasoning_cycles) el('reasoningCycles').value=d.status.reasoning_cycles; + el('adaptiveCompute').value = d.status.adaptive_compute ? 'on' : 'off'; + if(d.status.adaptive_exit_tol !== undefined) el('exitTol').value = d.status.adaptive_exit_tol; + if(d.status.adaptive_exit_entropy !== undefined) el('exitEntropy').value = d.status.adaptive_exit_entropy; + if(d.status.prediction_stability_patience !== undefined) el('stabilityPatience').value = d.status.prediction_stability_patience; + if(d.status.prediction_stability_tol !== undefined) el('stabilityTol').value = d.status.prediction_stability_tol; }catch(e){ el('statusBox').textContent='TELEMETRY ERROR: '+e.message; } } +function renderRouteShadowRegistry(snapshot){ + const campaigns=Array.isArray(snapshot&&snapshot.campaigns)?snapshot.campaigns:[]; + const committed=campaigns.reduce((n,row)=>n+(Number(row.commitment_count)||0),0); + const matched=campaigns.reduce((n,row)=>n+(Number(row.matched_assignment_count)||0),0); + const processed=campaigns.reduce((n,row)=>n+(Number(row.processed_reveal_count)||0),0); + const mismatched=campaigns.reduce((n,row)=>n+(Number(row.mismatched_assignment_count)||0),0); + const chain=snapshot&&snapshot.event_chain; + if(!snapshot||snapshot.available!==true){ + el('routeShadowStatus').textContent=`Not initialized at ${(snapshot&&snapshot.registry_location)||'the canonical memory path'}.\nRead only - execution, activation, and promotion unavailable.`; + return; + } + const states=[...new Set(campaigns.map(row=>String(row.state||'unknown').replaceAll('_',' ')))]; + el('routeShadowStatus').textContent= + `${snapshot.ok?'VERIFIED':'VERIFICATION FAILED'} | campaigns ${campaigns.length} | assignments ${matched}/${committed} matched | reveals ${processed} processed | mismatches ${mismatched}\n`+ + `chain ${chain&&chain.ok?'verified':'failed'} (${Number(chain&&chain.verified_events)||0} events) | states ${states.join(', ')||'none'}\n`+ + 'Local chain only; no external anchor. Read only - execution, activation, and promotion unavailable.'; +} +async function refreshRouteShadowRegistry(){ + const button=el('routeShadowRefresh');button.disabled=true;el('routeShadowStatus').textContent='Reading local registry...'; + try{const d=await jget('/api/route_shadow_registry/status');renderRouteShadowRegistry(d.route_shadow_registry||{});} + catch(e){el('routeShadowStatus').textContent='REGISTRY STATUS ERROR: '+e.message;} + finally{button.disabled=false;} +} async function loadModel(){ el('statusBox').textContent='SYNCING NEURAL WEIGHTS...'; try{ @@ -289,10 +393,25 @@ promptEl.value=''; promptEl.style.height = 'auto'; try{ - const d=await jpost('/api/chat',{session_id:sid,message:text,style_mode:el('style').value,response_temperature:Number(el('rt').value),show_top_responses:Number(el('showTop').value)}); - add('bot',d.response,d.timing_ms,d.top_candidates); + const cycles = reasoningCyclesValue(); + const d=await jpost('/api/chat',{session_id:sid,message:text,style_mode:el('style').value,response_temperature:Number(el('rt').value),show_top_responses:Number(el('showTop').value),reasoning_cycles:cycles,adaptive_compute:el('adaptiveCompute').value==='on',adaptive_exit_tol:Number(el('exitTol').value),adaptive_exit_entropy:Number(el('exitEntropy').value),prediction_stability_patience:Number(el('stabilityPatience').value),prediction_stability_tol:Number(el('stabilityTol').value)}); + add('bot',d.response,d.timing_ms,d.top_candidates,d.compute); }catch(e){ add('bot','CORE ERROR: '+e.message); } } +async function sweepCompute(){ + const text=promptEl.value.trim(); if(!text) return; + try{ + const d=await jpost('/api/compute_sweep',{session_id:sid,message:text,cycles:[1,3,8],adaptive_compute:el('adaptiveCompute').value==='on',adaptive_exit_tol:Number(el('exitTol').value),adaptive_exit_entropy:Number(el('exitEntropy').value),prediction_stability_patience:Number(el('stabilityPatience').value),prediction_stability_tol:Number(el('stabilityTol').value)}); + const lines=['Compute sweep for draft prompt:']; + d.rows.forEach((row)=>{ + const conf=fmtNum(row.confidence); + const entropy=fmtNum(row.entropy); + const reason=row.compute&&row.compute.exit_reason?` | exit ${row.compute.exit_reason}`:''; + lines.push(`cycles ${row.requested_cycles}: ${row.latency_ms}ms | used ${row.cycles_used} | label ${row.predicted_label} | conf ${conf||'n/a'} | entropy ${entropy||'n/a'}${reason}`); + }); + add('bot',lines.join('\\n'),null,null,d.rows[d.rows.length-1]?.compute||null); + }catch(e){ add('bot','SWEEP ERROR: '+e.message); } +} async function clearSess(){ try{ await jpost('/api/clear',{session_id:sid}); @@ -300,17 +419,85 @@ add('bot','Neural cache purged. Ready for fresh session.'); }catch(e){ add('bot','PURGE ERROR: '+e.message); } } -el('loadBtn').onclick=loadModel; el('statusBtn').onclick=refresh; el('clearBtn').onclick=clearSess; el('sendBtn').onclick=send; +el('loadBtn').onclick=loadModel; el('statusBtn').onclick=refresh; el('clearBtn').onclick=clearSess; el('sendBtn').onclick=send; el('sweepBtn').onclick=sweepCompute; el('routeShadowRefresh').onclick=refreshRouteShadowRegistry; promptEl.addEventListener('keydown',e=>{ if(e.key==='Enter'&&!e.shiftKey){ e.preventDefault();send(); } }); -refresh(); +refresh(); refreshRouteShadowRegistry(); """ +_RUNTIME_COMPUTE_DEFAULT_KEYS = ( + "reasoning_cycles", + "adaptive_compute", + "adaptive_exit_tol", + "adaptive_exit_entropy", + "prediction_stability_patience", + "prediction_stability_tol", +) + + +def _library_runtime_compute_defaults() -> Dict[str, Any]: + return { + "reasoning_cycles": None, + "adaptive_compute": False, + "adaptive_exit_tol": 1e-3, + "adaptive_exit_entropy": chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY, + "prediction_stability_patience": chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, + "prediction_stability_tol": chat_app.DEFAULT_PREDICTION_STABILITY_TOL, + } + + +def _normalize_runtime_compute_defaults(values: Dict[str, Any]) -> Dict[str, Any]: + raw_cycles = values.get("reasoning_cycles") + reasoning_cycles: Any + if chat_app._is_auto_reasoning_cycles(raw_cycles): + reasoning_cycles = "auto" + else: + reasoning_cycles = chat_app._coerce_optional_positive_int( + raw_cycles, + chat_app.MAX_RUNTIME_REASONING_CYCLES, + ) + return { + "reasoning_cycles": reasoning_cycles, + "adaptive_compute": chat_app._coerce_bool(values.get("adaptive_compute")), + "adaptive_exit_tol": chat_app._coerce_nonnegative_float( + values.get("adaptive_exit_tol"), + 1e-3, + ), + "adaptive_exit_entropy": chat_app._coerce_nonnegative_float( + values.get("adaptive_exit_entropy"), + chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY, + ), + "prediction_stability_patience": chat_app._coerce_nonnegative_int( + values.get("prediction_stability_patience"), + chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, + chat_app.MAX_RUNTIME_REASONING_CYCLES, + ), + "prediction_stability_tol": chat_app._coerce_nonnegative_float( + values.get("prediction_stability_tol"), + chat_app.DEFAULT_PREDICTION_STABILITY_TOL, + ), + } + + +def _runtime_compute_cli_overrides(args: argparse.Namespace) -> Dict[str, Any]: + """Return only compute options the CLI user actually supplied.""" + values = { + "reasoning_cycles": getattr(args, "reasoning_cycles", None), + "adaptive_compute": getattr(args, "adaptive_compute", None), + "adaptive_exit_tol": getattr(args, "adaptive_exit_tol", None), + "adaptive_exit_entropy": getattr(args, "adaptive_exit_entropy", None), + "prediction_stability_patience": getattr(args, "prediction_stability_patience", None), + "prediction_stability_tol": getattr(args, "prediction_stability_tol", None), + } + return {key: value for key, value in values.items() if value is not None} + + class Engine: def __init__(self, device: Any, device_info: Dict[str, Any], defaults: Dict[str, Any]): self.device = device self.device_info = dict(device_info or {}) - self.defaults = dict(defaults) + self._constructor_defaults = dict(defaults or {}) + self.defaults = self._build_effective_defaults({}) self.lock = threading.RLock() self.model = None self.weights_path: Optional[str] = None @@ -321,6 +508,32 @@ def __init__(self, device: Any, device_info: Dict[str, Any], defaults: Dict[str, self.available_labels: List[int] = list(range(chat_app.MODEL_CLASSES)) self.sessions: Dict[str, List[Tuple[str, str]]] = {} self.recent: Dict[str, List[str]] = {} + registry_path = self._constructor_defaults.get("route_shadow_registry_path") + self.route_shadow_registry_path = Path( + registry_path or Path("tmp") / "memory" / "route-policy-shadow-registry.sqlite3" + ).expanduser().resolve() + + def _build_effective_defaults(self, meta: Dict[str, Any]) -> Dict[str, Any]: + runtime_defaults = _library_runtime_compute_defaults() + metadata_defaults = meta.get("runtime_defaults") + if isinstance(metadata_defaults, dict): + runtime_defaults.update( + { + key: metadata_defaults[key] + for key in _RUNTIME_COMPUTE_DEFAULT_KEYS + if key in metadata_defaults + } + ) + runtime_defaults.update( + { + key: self._constructor_defaults[key] + for key in _RUNTIME_COMPUTE_DEFAULT_KEYS + if key in self._constructor_defaults + } + ) + effective = dict(self._constructor_defaults) + effective.update(_normalize_runtime_compute_defaults(runtime_defaults)) + return effective def status(self) -> Dict[str, Any]: with self.lock: @@ -333,8 +546,56 @@ def status(self) -> Dict[str, Any]: "available_labels": len(self.available_labels), "device": self.device_info.get("resolved", str(self.device)), "sessions": len(self.sessions), + "runtime_compute_supported": chat_app.model_supports_runtime_compute(self.model) if self.model is not None else False, + "reasoning_cycles": chat_app._format_reasoning_cycles_setting(self.defaults.get("reasoning_cycles")), + "adaptive_compute": bool(self.defaults.get("adaptive_compute", False)), + "adaptive_exit_tol": chat_app._coerce_nonnegative_float(self.defaults.get("adaptive_exit_tol", 1e-3), 1e-3), + "adaptive_exit_entropy": chat_app._coerce_nonnegative_float(self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + "prediction_stability_patience": chat_app._coerce_nonnegative_int(self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, chat_app.MAX_RUNTIME_REASONING_CYCLES), + "prediction_stability_tol": chat_app._coerce_nonnegative_float(self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), chat_app.DEFAULT_PREDICTION_STABILITY_TOL), } + def route_shadow_registry_snapshot(self) -> Dict[str, Any]: + """Return compatible shadow-registry status without exposing mutations.""" + + registry_path = self.route_shadow_registry_path + if not registry_path.is_file(): + return { + "ok": True, + "available": False, + "status": "not_initialized", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + if RouteShadowAssignmentRegistry is None: + return { + "ok": False, + "available": True, + "status": "reader_unavailable", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + snapshot = RouteShadowAssignmentRegistry(registry_path, read_only=True).snapshot() + return { + **snapshot, + "available": True, + "status": "verified" if snapshot.get("ok") else "verification_failed", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + } + def _parse_buckets(self, meta: Dict[str, Any]) -> None: buckets: Dict[int, List[Dict[str, Any]]] = {} raw = meta.get("buckets", {}) @@ -359,13 +620,14 @@ def load(self, weights: str, meta_path: str) -> Dict[str, Any]: raise FileNotFoundError(f"Metadata not found: {meta_path}") meta = chat_app.load_metadata(meta_path) + effective_defaults = self._build_effective_defaults(meta) raw_feature_mode = str(meta.get("feature_mode", "legacy")).strip().lower() feature_mode = chat_app.resolve_feature_mode(raw_feature_mode, smarter_auto=True) sd = chat_app.safe_load_state_dict(weights) inferred = chat_app.detect_model_size_from_state_dict(sd) resolved_model_size, _ = chat_app.resolve_runtime_model_size( - str(self.defaults.get("model_size", "auto")), + str(effective_defaults.get("model_size", "auto")), str(meta.get("model_size", "")), inferred, ) @@ -398,6 +660,7 @@ def load(self, weights: str, meta_path: str) -> Dict[str, Any]: self.meta_path = meta_path self.feature_mode = feature_mode self.model_size = resolved_model_size + self.defaults = effective_defaults self._parse_buckets(meta) self.sessions.clear() self.recent.clear() @@ -409,7 +672,86 @@ def clear(self, session_id: str) -> None: self.sessions.pop(session_id, None) self.recent.pop(session_id, None) - def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None, response_temperature: Optional[float] = None, show_top_responses: int = 0) -> Dict[str, Any]: + def compute_sweep( + self, + session_id: str, + user_text: str, + cycles: Any = None, + adaptive_compute: Any = None, + adaptive_exit_tol: Any = None, + adaptive_exit_entropy: Any = None, + prediction_stability_patience: Any = None, + prediction_stability_tol: Any = None, + ) -> Dict[str, Any]: + if not user_text.strip(): + raise ValueError("Empty message") + requested_cycles: List[int] = [] + raw_cycles = cycles if isinstance(cycles, list) and cycles else [1, 3, 8] + for raw in raw_cycles: + parsed = chat_app._coerce_optional_positive_int(raw, chat_app.MAX_RUNTIME_REASONING_CYCLES) + if parsed is not None and parsed not in requested_cycles: + requested_cycles.append(parsed) + if not requested_cycles: + requested_cycles = [1, 3, 8] + + with self.lock: + if self.model is None: + raise RuntimeError("No model loaded") + model = self.model + feature_mode = self.feature_mode + labels = list(self.available_labels) + history = list(self.sessions.get(session_id, [])) + + context = chat_app.build_context(history, user_text=user_text, max_turns=int(self.defaults.get("max_turns", 2))) + x = chat_app.text_to_model_input(context, feature_mode=feature_mode).to(self.device) + idx = torch.tensor(labels, dtype=torch.long, device=self.device) + rows: List[Dict[str, Any]] = [] + + with torch.no_grad(): + for cycle_count in requested_cycles: + t0 = time.perf_counter() + model_out, compute_diag = chat_app.forward_with_runtime_compute( + model, + x, + reasoning_cycles=cycle_count, + adaptive_compute=adaptive_compute if adaptive_compute is not None else self.defaults.get("adaptive_compute", False), + exit_tol=adaptive_exit_tol if adaptive_exit_tol is not None else self.defaults.get("adaptive_exit_tol", 1e-3), + exit_entropy_threshold=adaptive_exit_entropy if adaptive_exit_entropy is not None else self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + prediction_stability_patience=prediction_stability_patience if prediction_stability_patience is not None else self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), + prediction_stability_tol=prediction_stability_tol if prediction_stability_tol is not None else self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), + ) + latency_ms = round((time.perf_counter() - t0) * 1000, 1) + logits = model_out[0, 0] + avail_logits = logits.index_select(0, idx) + probs = torch.softmax(avail_logits, dim=0) + top_pos = int(torch.argmax(probs).item()) + entropy = -torch.sum(probs * torch.log(probs + 1e-9)) + rows.append({ + "requested_cycles": cycle_count, + "latency_ms": latency_ms, + "cycles_used": compute_diag.get("cycles_used"), + "predicted_label": int(labels[top_pos]), + "confidence": float(probs[top_pos].item()), + "entropy": float(entropy.item()), + "compute": compute_diag, + }) + + return {"ok": True, "session_id": session_id, "rows": rows} + + def chat( + self, + session_id: str, + user_text: str, + style_mode: Optional[str] = None, + response_temperature: Optional[float] = None, + show_top_responses: int = 0, + reasoning_cycles: Any = None, + adaptive_compute: Any = None, + adaptive_exit_tol: Any = None, + adaptive_exit_entropy: Any = None, + prediction_stability_patience: Any = None, + prediction_stability_tol: Any = None, + ) -> Dict[str, Any]: if not user_text.strip(): raise ValueError("Empty message") with self.lock: @@ -429,7 +771,18 @@ def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None tt = time.perf_counter() x = chat_app.text_to_model_input(context, feature_mode=feature_mode).to(self.device) with torch.no_grad(): - logits = model(x)[0, 0] + model_out, compute_diag = chat_app.forward_with_runtime_compute( + model, + x, + reasoning_cycles=reasoning_cycles if reasoning_cycles is not None else self.defaults.get("reasoning_cycles"), + adaptive_compute=adaptive_compute if adaptive_compute is not None else self.defaults.get("adaptive_compute", False), + exit_tol=adaptive_exit_tol if adaptive_exit_tol is not None else self.defaults.get("adaptive_exit_tol", 1e-3), + exit_entropy_threshold=adaptive_exit_entropy if adaptive_exit_entropy is not None else self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + prediction_stability_patience=prediction_stability_patience if prediction_stability_patience is not None else self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), + prediction_stability_tol=prediction_stability_tol if prediction_stability_tol is not None else self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), + auto_reasoning_context=context, + ) + logits = model_out[0, 0] t_infer += time.perf_counter() - tt idx = torch.tensor(labels, dtype=torch.long, device=logits.device) @@ -526,7 +879,9 @@ def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None "infer": round(t_infer * 1000, 1), "rank_pick": round(t_rank * 1000, 1), "total": round((time.perf_counter() - t0) * 1000, 1), + "cycles_used": compute_diag.get("cycles_used"), }, + "compute": compute_diag, "top_candidates": top_candidates, } @@ -536,14 +891,29 @@ def build_app(engine: Engine, default_weights: str, default_meta: str): @app.get('/') def index(): - html = HTML.replace("", f"") - html = html.replace("", f"") + html = HTML.replace( + "", + f"", + ) + html = html.replace( + "", + f"", + ) return html @app.get('/api/status') def api_status(): return jsonify({"ok": True, "status": engine.status()}) + @app.get('/api/route_shadow_registry/status') + def api_route_shadow_registry_status(): + try: + response = jsonify({"ok": True, "route_shadow_registry": engine.route_shadow_registry_snapshot()}) + response.headers["Cache-Control"] = "no-store" + return response + except Exception as e: + return jsonify({"ok": False, "error": str(e)}), 500 + @app.post('/api/load') def api_load(): p = request.get_json(force=True, silent=True) or {} @@ -566,6 +936,31 @@ def api_chat(): style_mode=p.get('style_mode'), response_temperature=p.get('response_temperature'), show_top_responses=int(p.get('show_top_responses') or 0), + reasoning_cycles=p.get('reasoning_cycles'), + adaptive_compute=p.get('adaptive_compute'), + adaptive_exit_tol=p.get('adaptive_exit_tol'), + adaptive_exit_entropy=p.get('adaptive_exit_entropy'), + prediction_stability_patience=p.get('prediction_stability_patience'), + prediction_stability_tol=p.get('prediction_stability_tol'), + )) + except Exception as e: + return jsonify({"ok": False, "error": str(e)}), 400 + + @app.post('/api/compute_sweep') + def api_compute_sweep(): + p = request.get_json(force=True, silent=True) or {} + sid = str(p.get('session_id') or '').strip() or str(uuid.uuid4()) + msg = str(p.get('message') or '').strip() + try: + return jsonify(engine.compute_sweep( + session_id=sid, + user_text=msg, + cycles=p.get('cycles'), + adaptive_compute=p.get('adaptive_compute'), + adaptive_exit_tol=p.get('adaptive_exit_tol'), + adaptive_exit_entropy=p.get('adaptive_exit_entropy'), + prediction_stability_patience=p.get('prediction_stability_patience'), + prediction_stability_tol=p.get('prediction_stability_tol'), )) except Exception as e: return jsonify({"ok": False, "error": str(e)}), 400 @@ -603,6 +998,25 @@ def main() -> None: ap.add_argument('--temperature', type=float, default=0.0) ap.add_argument('--style_mode', choices=['auto','balanced','creative','concise','analyst'], default='auto') ap.add_argument('--creativity', type=float, default=0.2) + ap.add_argument('--reasoning_cycles', type=str, default=None) + adaptive_compute_group = ap.add_mutually_exclusive_group() + adaptive_compute_group.add_argument( + '--adaptive_compute', + dest='adaptive_compute', + action='store_true', + help='enable adaptive compute, overriding checkpoint metadata', + ) + adaptive_compute_group.add_argument( + '--no_adaptive_compute', + dest='adaptive_compute', + action='store_false', + help='disable adaptive compute, overriding checkpoint metadata', + ) + ap.set_defaults(adaptive_compute=None) + ap.add_argument('--adaptive_exit_tol', type=float, default=None) + ap.add_argument('--adaptive_exit_entropy', type=float, default=None) + ap.add_argument('--prediction_stability_patience', type=int, default=None) + ap.add_argument('--prediction_stability_tol', type=float, default=None) args = ap.parse_args() configure_torch_runtime( @@ -612,7 +1026,7 @@ def main() -> None: matmul_precision=str(args.matmul_precision), ) device, device_info = resolve_device(args.device, preference=args.device_preference) - engine = Engine(device, device_info, { + engine_defaults = { 'model_size': args.model_size, 'max_turns': int(args.max_turns), 'top_labels': int(args.top_labels), @@ -621,7 +1035,9 @@ def main() -> None: 'temperature': float(args.temperature), 'style_mode': str(args.style_mode), 'creativity': float(args.creativity), - }) + } + engine_defaults.update(_runtime_compute_cli_overrides(args)) + engine = Engine(device, device_info, engine_defaults) if args.autoload: try: print(engine.load(args.weights, args.meta)) diff --git a/source/benchmark_auto_collective_loop_common.py b/source/benchmark_auto_collective_loop_common.py new file mode 100644 index 000000000..0802159cc --- /dev/null +++ b/source/benchmark_auto_collective_loop_common.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import sys +import time +from collections import Counter +from difflib import SequenceMatcher +from pathlib import Path +from typing import Dict, List, Sequence + +SOURCE_DIR = Path(__file__).resolve().parent +if str(SOURCE_DIR) not in sys.path: + sys.path.append(str(SOURCE_DIR)) + +from benchmark_all_models_common import ( + _extract_answer, + _overall_exact_score, + _stable_hash, + _write_csv, + _write_jsonl, + build_benchmark_items, + draw_graph, +) +from multimodel_catalog import DEFAULT_COMMON_SUMMARY, DEFAULT_MODELS_DIR, discover_model_records +from multimodel_runtime import UnifiedModelManager +from qwen_supermix_pipeline import token_f1 + + +AUTO_COLLECTIVE_LOOP_KEY = "auto_collective_loop" + + +def _normalize_response_text(value: object) -> str: + return " ".join(str(value or "").strip().split()) + + +def _sorted_counts(counter: Counter[str]) -> Dict[str, int]: + return dict(sorted(counter.items(), key=lambda item: (-item[1], item[0].lower()))) + + +def benchmark_auto_collective_loop( + *, + manager: UnifiedModelManager, + items, + log_every: int, + loop_max_steps: int, + consultant_keys: Sequence[str], + consultant_limit: int, + backend_cache_size: int, +) -> tuple[list[dict], list[dict]]: + benchmark_names = sorted({item.benchmark for item in items}) + details: List[dict] = [] + per_benchmark_exact: Dict[str, List[float]] = {name: [] for name in benchmark_names} + token_scores: List[float] = [] + char_scores: List[float] = [] + gen_seconds: List[float] = [] + route_counts: Counter[str] = Counter() + consulted_counts: Counter[str] = Counter() + loop_steps_used: List[int] = [] + start_time = time.time() + + settings = { + "agent_mode": "collective_loop", + "loop_max_steps": int(loop_max_steps), + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + } + if consultant_keys: + settings["collective_consultant_keys"] = list(consultant_keys) + if consultant_limit > 0: + settings["collective_consultant_limit"] = int(consultant_limit) + + print(f"[bench] loading {AUTO_COLLECTIVE_LOOP_KEY} (1/1)") + for item_index, item in enumerate(items, start=1): + item_id = f"{item.benchmark}:{item_index:04d}:{_stable_hash(item.prompt)[:12]}" + t0 = time.time() + payload = manager.handle_prompt( + session_id=f"benchmark::{AUTO_COLLECTIVE_LOOP_KEY}::{item_id}", + prompt=item.prompt, + model_key="auto", + action_mode="auto", + settings=settings, + ) + elapsed = time.time() - t0 + + prediction = _normalize_response_text(payload.get("response")) + extracted = _extract_answer(item, prediction) + exact = 1.0 if _normalize_response_text(extracted).lower() == _normalize_response_text(item.reference_extracted).lower() else 0.0 + token = float(token_f1(item.reference_text, prediction)) + char = float(SequenceMatcher(None, item.reference_text.lower(), prediction.lower()).ratio()) + + agent_trace = payload.get("agent_trace") if isinstance(payload.get("agent_trace"), dict) else {} + consultation_rows = agent_trace.get("consultation_rows") if isinstance(agent_trace.get("consultation_rows"), list) else [] + consultation_keys: List[str] = [] + for row in consultation_rows: + if not isinstance(row, dict): + continue + consultant_key = str(row.get("model_key") or row.get("model_label") or "").strip() + if consultant_key: + consultation_keys.append(consultant_key) + consulted_counts[consultant_key] += 1 + loop_steps = agent_trace.get("loop_steps") if isinstance(agent_trace.get("loop_steps"), list) else [] + loop_step_count = int(payload.get("timing", {}).get("loop_steps") or len(loop_steps) or 0) + loop_steps_used.append(loop_step_count) + + active_model_key = str(payload.get("active_model_key") or payload.get("model_key") or "").strip() + if active_model_key: + route_counts[active_model_key] += 1 + + per_benchmark_exact[item.benchmark].append(exact) + token_scores.append(token) + char_scores.append(char) + gen_seconds.append(elapsed) + + details.append( + { + "model": AUTO_COLLECTIVE_LOOP_KEY, + "family": "router", + "benchmark": item.benchmark, + "item_id": item_id, + "item_index": item_index, + "prompt": item.prompt, + "prompt_hash": _stable_hash(item.prompt), + "reference_text": item.reference_text, + "reference_hash": _stable_hash(item.reference_text), + "reference_extracted": item.reference_extracted, + "prediction": prediction, + "prediction_hash": _stable_hash(prediction), + "prediction_extracted": extracted, + "exact": exact, + "is_exact": bool(exact >= 1.0), + "token_f1": token, + "char_similarity": char, + "gen_seconds": elapsed, + "active_model_key": active_model_key, + "active_model_label": str(payload.get("active_model_label") or ""), + "route_reason": str(payload.get("route_reason") or ""), + "agent_mode": str(agent_trace.get("agent_mode") or ""), + "loop_steps": loop_step_count, + "consultation_model_keys": consultation_keys, + "consultation_count": len(consultation_keys), + } + ) + if log_every > 0 and item_index % log_every == 0: + print(f"[bench] {AUTO_COLLECTIVE_LOOP_KEY} {item_index}/{len(items)} done") + + benchmark_scores = { + name: float(sum(values) / max(1, len(values))) + for name, values in per_benchmark_exact.items() + } + summary_rows = [ + { + "model": AUTO_COLLECTIVE_LOOP_KEY, + "family": "router", + "overall_exact": _overall_exact_score(benchmark_scores), + "avg_token_f1": float(sum(token_scores) / max(1, len(token_scores))), + "avg_char_similarity": float(sum(char_scores) / max(1, len(char_scores))), + "avg_gen_seconds": float(sum(gen_seconds) / max(1, len(gen_seconds))), + "model_seconds": float(time.time() - start_time), + "benchmarks": benchmark_scores, + "runtime_settings": { + "model_key": "auto", + "action_mode": "auto", + "agent_mode": "collective_loop", + "loop_max_steps": int(loop_max_steps), + "collective_consultant_keys": list(consultant_keys), + "collective_consultant_limit": int(consultant_limit), + "backend_cache_size": int(backend_cache_size), + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + "routed_model_counts": _sorted_counts(route_counts), + "consultation_model_counts": _sorted_counts(consulted_counts), + "avg_loop_steps": float(sum(loop_steps_used) / max(1, len(loop_steps_used))), + } + ] + print(f"[bench] finished {AUTO_COLLECTIVE_LOOP_KEY} overall_exact={summary_rows[0]['overall_exact']:.4f}") + return summary_rows, details + + +def main() -> int: + parser = argparse.ArgumentParser(description="Benchmark the multimodel auto router with collective loop mode on sampled common benchmarks.") + parser.add_argument("--output_dir", required=True) + parser.add_argument("--models_dir", default=str(DEFAULT_MODELS_DIR)) + parser.add_argument("--common_summary", default=str(DEFAULT_COMMON_SUMMARY)) + parser.add_argument("--device_preference", default="cpu") + parser.add_argument("--backend_cache_size", type=int, default=3) + parser.add_argument("--sample_per_benchmark", type=int, default=50) + parser.add_argument("--seed", type=int, default=20260327) + parser.add_argument("--log_every", type=int, default=10) + parser.add_argument("--loop_max_steps", type=int, default=3) + parser.add_argument( + "--collective_consultant_keys", + default="omni_collective_v41,v40_benchmax,omni_collective_v8,v33_final,qwen_v28,omni_collective_v7,math_equation_micro_v1", + help="Comma-separated chat model keys to prioritize in the collective worker panel.", + ) + parser.add_argument("--collective_consultant_limit", type=int, default=7) + args = parser.parse_args() + + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + models_dir = Path(args.models_dir).resolve() + common_summary = Path(args.common_summary).resolve() + consultant_keys = [part.strip() for part in str(args.collective_consultant_keys or "").split(",") if part.strip()] + + records = discover_model_records(models_dir=models_dir, common_summary_path=common_summary) + if not records: + raise RuntimeError(f"No local model records discovered under {models_dir}") + + manager = UnifiedModelManager( + records=tuple(records), + extraction_root=output_dir / "runtime_extract", + generated_dir=output_dir / "generated", + device_preference=str(args.device_preference), + models_dir=models_dir, + common_summary_path=common_summary, + backend_cache_size=int(args.backend_cache_size), + ) + + items = build_benchmark_items(sample_per_benchmark=int(args.sample_per_benchmark), seed=int(args.seed)) + benchmark_names = sorted({item.benchmark for item in items}) + + prompts_jsonl = output_dir / "benchmark_items.jsonl" + _write_jsonl( + prompts_jsonl, + ( + { + "item_id": f"{item.benchmark}:{index:04d}:{_stable_hash(item.prompt)[:12]}", + "item_index": index, + "benchmark": item.benchmark, + "prompt": item.prompt, + "prompt_hash": _stable_hash(item.prompt), + "reference_text": item.reference_text, + "reference_hash": _stable_hash(item.reference_text), + "reference_extracted": item.reference_extracted, + "max_new_tokens": item.max_new_tokens, + } + for index, item in enumerate(items, start=1) + ), + ) + + summary_rows, details = benchmark_auto_collective_loop( + manager=manager, + items=items, + log_every=int(args.log_every), + loop_max_steps=int(args.loop_max_steps), + consultant_keys=consultant_keys, + consultant_limit=int(args.collective_consultant_limit), + backend_cache_size=int(args.backend_cache_size), + ) + + graph_path = output_dir / "benchmark_all_models_common_graph.png" + draw_graph(graph_path, summary_rows, benchmark_names) + + summary_path = output_dir / "benchmark_all_models_common_summary.json" + details_path = output_dir / "benchmark_all_models_common_details.jsonl" + csv_path = output_dir / "benchmark_all_models_common_table.csv" + + _write_jsonl(details_path, details) + _write_csv(csv_path, summary_rows, benchmark_names) + summary_path.write_text( + json.dumps( + { + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "models_dir": str(models_dir), + "common_summary": str(common_summary), + "output_dir": str(output_dir), + "device_preference": str(args.device_preference), + "backend_cache_size": int(args.backend_cache_size), + "sample_per_benchmark": int(args.sample_per_benchmark), + "benchmarks": benchmark_names, + "models_benchmarked": [AUTO_COLLECTIVE_LOOP_KEY], + "summary_rows": summary_rows, + "artifacts": { + "prompts_jsonl": str(prompts_jsonl), + "details_jsonl": str(details_path), + "table_csv": str(csv_path), + "graph_png": str(graph_path), + }, + }, + indent=2, + ), + encoding="utf-8", + ) + + print(f"[bench] wrote summary to {summary_path}") + print(f"[bench] wrote graph to {graph_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/source/chat_app.py b/source/chat_app.py index c42331cba..58d4d3027 100644 --- a/source/chat_app.py +++ b/source/chat_app.py @@ -1,9 +1,10 @@ import argparse +import inspect import json import re import time from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from concurrent.futures import ThreadPoolExecutor import torch @@ -71,6 +72,11 @@ def _int_or_default(value, default: int) -> int: _FOLLOWUP_QUERY_RE = re.compile(r"\b(it|that|this|they|them|same|previous|above|more|continue|deeper|expand)\b", re.I) VALID_RUNTIME_MODEL_SIZES = SUPPORTED_MODEL_SIZES +MAX_RUNTIME_REASONING_CYCLES = 64 +DEFAULT_ADAPTIVE_EXIT_ENTROPY = 0.2 +DEFAULT_PREDICTION_STABILITY_PATIENCE = 2 +DEFAULT_PREDICTION_STABILITY_TOL = 5e-3 +AUTO_REASONING_CYCLE_BUCKETS = (1, 3, 8, 16) def _build_db_query(user: str, history: List[Tuple[str, str]], memory_rows: List[Dict], max_turns: int = 2) -> str: @@ -145,6 +151,257 @@ def resolve_runtime_model_size(requested_model_size: str, meta_model_size: str, return resolved_model_size, warning +def _coerce_optional_positive_int(value: Any, max_value: Optional[int] = None) -> Optional[int]: + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + if stripped in {"", "default", "none", "off", "auto"}: + return None + value = stripped + try: + parsed = int(value) + except (TypeError, ValueError): + return None + parsed = max(1, parsed) + if max_value is not None: + parsed = min(parsed, max(1, int(max_value))) + return parsed + + +def _is_auto_reasoning_cycles(value: Any) -> bool: + return isinstance(value, str) and value.strip().lower() in {"auto", "adaptive", "smart"} + + +def _format_reasoning_cycles_setting(value: Any) -> Any: + if _is_auto_reasoning_cycles(value): + return "auto" + return _coerce_optional_positive_int(value, MAX_RUNTIME_REASONING_CYCLES) or "default" + + +def estimate_auto_reasoning_cycles(context: str, max_value: Optional[int] = None) -> Dict[str, Any]: + text = str(context or "") + lowered = text.lower() + reasons: List[str] = [] + score = 0 + + char_count = len(text) + if char_count >= 1600: + score += 3 + reasons.append("long_context") + elif char_count >= 700: + score += 2 + reasons.append("medium_context") + elif char_count >= 280: + score += 1 + reasons.append("short_context") + + if re.search(r"```|traceback|exception|error|bug|fix|debug|refactor|class\s+|def\s+|function|import\s+|select\s+", lowered): + score += 2 + reasons.append("code_or_debug") + if re.search(r"\b(prove|derive|equation|calculate|solve|theorem|complexity|probability|optimi[sz]e|benchmarks?)\b|[=+\-*/^]{2,}", lowered): + score += 2 + reasons.append("math_or_verification") + if re.search(r"\b(compare|tradeoff|plan|architect|design|multi-step|step by step|analy[sz]e|research|evidence)\b", lowered): + score += 1 + reasons.append("deliberation") + if re.search(r"\b(agent|tool|browser|github|file|tests?|verify|integration|runtime)\b", lowered): + score += 1 + reasons.append("tool_or_runtime") + if re.search(r"\b(quick|brief|concise|short answer|one line)\b", lowered): + score -= 1 + reasons.append("brevity_requested") + + if score <= 0: + cycles = AUTO_REASONING_CYCLE_BUCKETS[0] + elif score <= 2: + cycles = AUTO_REASONING_CYCLE_BUCKETS[1] + elif score <= 4: + cycles = AUTO_REASONING_CYCLE_BUCKETS[2] + else: + cycles = AUTO_REASONING_CYCLE_BUCKETS[3] + cycles = _coerce_optional_positive_int(cycles, max_value) or AUTO_REASONING_CYCLE_BUCKETS[0] + + return { + "mode": "auto", + "score": int(score), + "reasons": reasons or ["simple_prompt"], + "cycles": int(cycles), + } + + +def _coerce_nonnegative_float(value: Any, default: float, max_value: Optional[float] = None) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + parsed = float(default) + parsed = max(0.0, parsed) + if max_value is not None: + parsed = min(parsed, float(max_value)) + return parsed + + +def _coerce_nonnegative_int(value: Any, default: int, max_value: Optional[int] = None) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + parsed = int(default) + parsed = max(0, parsed) + if max_value is not None: + parsed = min(parsed, max(0, int(max_value))) + return parsed + + +def _coerce_bool(value: Any) -> bool: + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on", "y"} + return bool(value) + + +def _model_forward_accepts_kwarg(model: Any, kwarg: str) -> bool: + forward = getattr(model, "forward", None) + if forward is None: + return False + try: + params = inspect.signature(forward).parameters.values() + except (TypeError, ValueError): + return False + for param in params: + if param.kind == inspect.Parameter.VAR_KEYWORD or param.name == kwarg: + return True + return False + + +def model_supports_runtime_compute(model: Any) -> bool: + return ( + _model_forward_accepts_kwarg(model, "reasoning_cycles") + or _model_forward_accepts_kwarg(model, "adaptive_compute") + ) + + +def _to_optional_scalar(value: Any) -> Optional[float]: + if value is None: + return None + try: + if hasattr(value, "detach"): + return float(value.detach().cpu().item()) + return float(value) + except (TypeError, ValueError, RuntimeError): + return None + + +def collect_runtime_compute_metrics(model: Any) -> Dict[str, Any]: + metrics: Dict[str, Any] = { + "cycles_used": None, + "ponder_cost": None, + "consistency_loss": None, + "gating_entropy": None, + "prediction_streak": None, + "prediction_confidence_delta": None, + "exit_reason": None, + } + attr_map = { + "last_cycles_used": "cycles_used", + "last_ponder_cost": "ponder_cost", + "last_consistency_loss": "consistency_loss", + "last_gating_entropy": "gating_entropy", + "last_prediction_streak": "prediction_streak", + "last_prediction_confidence_delta": "prediction_confidence_delta", + } + modules = list(model.modules()) if hasattr(model, "modules") else [model] + for module in modules: + for attr, key in attr_map.items(): + if metrics[key] is not None: + continue + scalar = _to_optional_scalar(getattr(module, attr, None)) + if scalar is not None: + metrics[key] = scalar + if metrics["exit_reason"] is None: + reason = getattr(module, "last_exit_reason", None) + if isinstance(reason, str) and reason.strip(): + metrics["exit_reason"] = reason.strip() + return metrics + + +def get_last_cycles_used(model: Any) -> Optional[float]: + return collect_runtime_compute_metrics(model).get("cycles_used") + + +def forward_with_runtime_compute( + model: Any, + x: torch.Tensor, + reasoning_cycles: Any = None, + adaptive_compute: Any = False, + exit_tol: Any = 1e-3, + exit_entropy_threshold: Any = DEFAULT_ADAPTIVE_EXIT_ENTROPY, + prediction_stability_patience: Any = DEFAULT_PREDICTION_STABILITY_PATIENCE, + prediction_stability_tol: Any = DEFAULT_PREDICTION_STABILITY_TOL, + auto_reasoning_context: str = "", +) -> Tuple[torch.Tensor, Dict[str, Any]]: + """Forward a model while applying optional v50 runtime-compute controls. + + Legacy model variants ignore these controls because their forward signatures + do not accept them. This keeps the runtime knobs safe for mixed checkpoints. + """ + auto_policy: Optional[Dict[str, Any]] = None + if _is_auto_reasoning_cycles(reasoning_cycles): + auto_policy = estimate_auto_reasoning_cycles( + auto_reasoning_context, + MAX_RUNTIME_REASONING_CYCLES, + ) + cycles = int(auto_policy["cycles"]) + else: + cycles = _coerce_optional_positive_int(reasoning_cycles, MAX_RUNTIME_REASONING_CYCLES) + adaptive = _coerce_bool(adaptive_compute) + tol = _coerce_nonnegative_float(exit_tol, 1e-3) + entropy_threshold = _coerce_nonnegative_float( + exit_entropy_threshold, + DEFAULT_ADAPTIVE_EXIT_ENTROPY, + ) + stability_patience = _coerce_nonnegative_int( + prediction_stability_patience, + DEFAULT_PREDICTION_STABILITY_PATIENCE, + MAX_RUNTIME_REASONING_CYCLES, + ) + stability_tol = _coerce_nonnegative_float( + prediction_stability_tol, + DEFAULT_PREDICTION_STABILITY_TOL, + ) + kwargs: Dict[str, Any] = {} + + if cycles is not None and _model_forward_accepts_kwarg(model, "reasoning_cycles"): + kwargs["reasoning_cycles"] = cycles + if adaptive and _model_forward_accepts_kwarg(model, "adaptive_compute"): + kwargs["adaptive_compute"] = True + if _model_forward_accepts_kwarg(model, "exit_tol"): + kwargs["exit_tol"] = tol + if _model_forward_accepts_kwarg(model, "exit_entropy_threshold"): + kwargs["exit_entropy_threshold"] = entropy_threshold + if _model_forward_accepts_kwarg(model, "prediction_stability_patience"): + kwargs["prediction_stability_patience"] = stability_patience + if _model_forward_accepts_kwarg(model, "prediction_stability_tol"): + kwargs["prediction_stability_tol"] = stability_tol + + output = model(x, **kwargs) if kwargs else model(x) + metrics = collect_runtime_compute_metrics(model) + diagnostics = { + "requested_reasoning_cycles": cycles, + "selected_reasoning_cycles": cycles, + "reasoning_budget_mode": "auto" if auto_policy is not None else ("manual" if cycles is not None else "default"), + "auto_reasoning_policy": auto_policy, + "max_reasoning_cycles": MAX_RUNTIME_REASONING_CYCLES, + "adaptive_compute": adaptive, + "exit_tol": tol, + "exit_entropy_threshold": entropy_threshold, + "prediction_stability_patience": stability_patience, + "prediction_stability_tol": stability_tol, + "applied": bool(kwargs), + "applied_kwargs": dict(kwargs), + **metrics, + } + return output, diagnostics + + def _default_expansion_dim_for_model_size(model_size: str) -> int: if model_size in THIRD_EXPANSION_DIM_MODEL_SIZES: return 1024 @@ -180,6 +437,12 @@ def _print_chat_help() -> None: print(" /clear Clear in-memory conversation history") print(" /style Set style: auto|balanced|creative|concise|analyst") print(" /creativity <0-1> Set creative rewrite strength") + print(" /cycles Set runtime reasoning cycles for future turns") + print(" /adaptive on|off Toggle v50 convergence early-exit") + print(" /exit_tol Set v50 adaptive-compute convergence tolerance") + print(" /exit_entropy Set v50 adaptive-compute entropy exit threshold") + print(" /stability Stable prediction cycles required (0 disables)") + print(" /stability_tol Maximum confidence drift across the stable streak") print(" /top Show top reranked candidates each turn (0 disables)") print(" /timing on|off Toggle per-turn timing output") print(" /memory on|off Toggle memory retrieval/writes for this session") @@ -204,6 +467,41 @@ def main(): ap.add_argument("--fifth_expansion_dim", type=int, default=None) ap.add_argument("--sixth_expansion_dim", type=int, default=None) ap.add_argument("--adapter_dropout", type=float, default=None) + ap.add_argument( + "--reasoning_cycles", + type=str, + default=None, + help="Optional runtime reasoning cycles, or 'auto' to choose from prompt difficulty; omit for checkpoint default.", + ) + ap.add_argument( + "--adaptive_compute", + action="store_true", + help="Enable v50 convergence early-exit during inference when the model supports it.", + ) + ap.add_argument( + "--adaptive_exit_tol", + type=float, + default=1e-3, + help="Convergence tolerance for v50 adaptive compute.", + ) + ap.add_argument( + "--adaptive_exit_entropy", + type=float, + default=DEFAULT_ADAPTIVE_EXIT_ENTROPY, + help="Prediction entropy threshold for v50 adaptive compute.", + ) + ap.add_argument( + "--prediction_stability_patience", + type=int, + default=DEFAULT_PREDICTION_STABILITY_PATIENCE, + help="Stable full-output predictions required before adaptive exit (0 disables).", + ) + ap.add_argument( + "--prediction_stability_tol", + type=float, + default=DEFAULT_PREDICTION_STABILITY_TOL, + help="Maximum confidence drift allowed across a stable prediction streak.", + ) ap.add_argument("--device", default="auto") ap.add_argument( "--device_preference", @@ -412,11 +710,22 @@ def main(): session_started_at = time.time() turn_count = 0 last_turn_timing: Dict[str, float] = {} + last_compute_diag: Dict[str, Any] = {} print(f"\n{TerminalColors.SYSTEM}--- Session Info ---") print(f"Loaded: {Path(args.weights).name} [{resolved_model_size}] | Available labels: {len(available_labels)}") print(f"Device: {device_info.get('resolved', args.device)} | Threads intra={torch.get_num_threads()} interop={torch.get_num_interop_threads()}") print(f"Feature mode: {feature_mode}{feature_mode_note} | Style mode: {args.style_mode} (creativity={max(0.0, min(1.0, float(args.creativity))):.2f})") + compute_support = "yes" if model_supports_runtime_compute(model) else "no" + print( + "Test-time compute: " + f"supported={compute_support} cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) if llm_db: print(f"LLM DB: {args.llm_db} (top_k={args.db_top_k})") @@ -462,12 +771,31 @@ def main(): f" style={args.style_mode} creativity={float(args.creativity):.2f} " f"top_debug={int(args.show_top_responses)} timing={'on' if show_timing else 'off'}" ) + print( + f" compute_supported={'yes' if model_supports_runtime_compute(model) else 'no'} " + f"cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) print( f" memory={'on' if session_memory_enabled else 'off'} " f"(db={'ready' if memory_db is not None else 'missing'}) | " f"llm_db={'on' if session_db_enabled else 'off'} " f"(db={'ready' if llm_db is not None else 'missing'})" ) + if last_compute_diag: + print( + " last_compute=" + f"applied={last_compute_diag.get('applied')} " + f"cycles_used={last_compute_diag.get('cycles_used')} " + f"ponder={last_compute_diag.get('ponder_cost')} " + f"gate_entropy={last_compute_diag.get('gating_entropy')} " + f"exit_reason={last_compute_diag.get('exit_reason')} " + f"prediction_streak={last_compute_diag.get('prediction_streak')}" + ) if last_turn_timing: print( " last_timing=" @@ -483,6 +811,15 @@ def main(): f" pool_mode={args.pool_mode} top_labels={int(args.top_labels)} " f"show_top_responses={int(args.show_top_responses)}" ) + print( + f" compute_supported={'yes' if model_supports_runtime_compute(model) else 'no'} " + f"cycles={_format_reasoning_cycles_setting(args.reasoning_cycles)} " + f"adaptive={'on' if bool(args.adaptive_compute) else 'off'} " + f"exit_tol={_coerce_nonnegative_float(args.adaptive_exit_tol, 1e-3):.4g} " + f"exit_entropy={_coerce_nonnegative_float(args.adaptive_exit_entropy, DEFAULT_ADAPTIVE_EXIT_ENTROPY):.4g} " + f"stability={_coerce_nonnegative_int(args.prediction_stability_patience, DEFAULT_PREDICTION_STABILITY_PATIENCE)} " + f"stability_tol={_coerce_nonnegative_float(args.prediction_stability_tol, DEFAULT_PREDICTION_STABILITY_TOL):.4g}" + ) print( f" memory={'on' if session_memory_enabled else 'off'} top_k={int(args.memory_top_k)} " f"pool={int(args.memory_pool_size)} score_scale={float(args.memory_score_scale):.3f}" @@ -503,6 +840,68 @@ def main(): print(f"{TerminalColors.SYSTEM}Creativity set to {float(args.creativity):.2f}.{TerminalColors.RESET}") except Exception: print(f"{TerminalColors.SYSTEM}Usage: /creativity 0.0-1.0{TerminalColors.RESET}") + elif cmd == "cycles": + raw_cycles = arg.strip().lower() + if raw_cycles in {"auto", "adaptive", "smart"}: + args.reasoning_cycles = "auto" + print(f"{TerminalColors.SYSTEM}Reasoning cycles set to auto.{TerminalColors.RESET}") + else: + cycles = _coerce_optional_positive_int(arg) + if cycles is None and raw_cycles not in {"", "default", "off", "none"}: + print(f"{TerminalColors.SYSTEM}Usage: /cycles {TerminalColors.RESET}") + continue + args.reasoning_cycles = cycles + label = cycles if cycles is not None else "default" + print(f"{TerminalColors.SYSTEM}Reasoning cycles set to {label}.{TerminalColors.RESET}") + elif cmd == "adaptive": + if arg.lower() in {"on", "1", "true"}: + args.adaptive_compute = True + elif arg.lower() in {"off", "0", "false"}: + args.adaptive_compute = False + else: + print(f"{TerminalColors.SYSTEM}Usage: /adaptive on|off{TerminalColors.RESET}") + continue + print(f"{TerminalColors.SYSTEM}Adaptive compute {'enabled' if args.adaptive_compute else 'disabled'}.{TerminalColors.RESET}") + elif cmd == "exit_tol": + try: + args.adaptive_exit_tol = _coerce_nonnegative_float(arg, args.adaptive_exit_tol) + print(f"{TerminalColors.SYSTEM}Adaptive exit tolerance set to {float(args.adaptive_exit_tol):.4g}.{TerminalColors.RESET}") + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /exit_tol {TerminalColors.RESET}") + elif cmd == "exit_entropy": + try: + args.adaptive_exit_entropy = _coerce_nonnegative_float( + arg, + args.adaptive_exit_entropy, + ) + print(f"{TerminalColors.SYSTEM}Adaptive exit entropy set to {float(args.adaptive_exit_entropy):.4g}.{TerminalColors.RESET}") + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /exit_entropy {TerminalColors.RESET}") + elif cmd == "stability": + try: + args.prediction_stability_patience = _coerce_nonnegative_int( + arg, + args.prediction_stability_patience, + MAX_RUNTIME_REASONING_CYCLES, + ) + print( + f"{TerminalColors.SYSTEM}Prediction stability patience set to " + f"{int(args.prediction_stability_patience)}.{TerminalColors.RESET}" + ) + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /stability {TerminalColors.RESET}") + elif cmd == "stability_tol": + try: + args.prediction_stability_tol = _coerce_nonnegative_float( + arg, + args.prediction_stability_tol, + ) + print( + f"{TerminalColors.SYSTEM}Prediction stability tolerance set to " + f"{float(args.prediction_stability_tol):.4g}.{TerminalColors.RESET}" + ) + except Exception: + print(f"{TerminalColors.SYSTEM}Usage: /stability_tol {TerminalColors.RESET}") elif cmd == "top": try: args.show_top_responses = max(0, int(arg)) @@ -585,7 +984,19 @@ def main(): x = text_to_model_input(context, feature_mode=feature_mode).to(device) with torch.no_grad(): - logits = model(x)[0, 0] # (10,) + model_out, compute_diag = forward_with_runtime_compute( + model, + x, + reasoning_cycles=args.reasoning_cycles, + adaptive_compute=args.adaptive_compute, + exit_tol=args.adaptive_exit_tol, + exit_entropy_threshold=args.adaptive_exit_entropy, + prediction_stability_patience=args.prediction_stability_patience, + prediction_stability_tol=args.prediction_stability_tol, + auto_reasoning_context=context, + ) + logits = model_out[0, 0] # (10,) + last_compute_diag = compute_diag t_infer += max(0.0, time.perf_counter() - _t) idx = torch.tensor(available_labels, dtype=torch.long, device=logits.device) @@ -732,6 +1143,14 @@ def main(): f"{TerminalColors.SYSTEM}Timing:{TerminalColors.RESET} " + ", ".join(f"{k}={_format_ms(v)}" for k, v in last_turn_timing.items()) ) + if last_compute_diag: + print( + f"{TerminalColors.SYSTEM}Compute:{TerminalColors.RESET} " + f"applied={last_compute_diag.get('applied')} " + f"cycles_used={last_compute_diag.get('cycles_used')} " + f"ponder={last_compute_diag.get('ponder_cost')} " + f"gate_entropy={last_compute_diag.get('gating_entropy')}" + ) finally: executor.shutdown(wait=False) diff --git a/source/chat_image_variant_app.py b/source/chat_image_variant_app.py index 8c39e616a..33b11cea6 100644 --- a/source/chat_image_variant_app.py +++ b/source/chat_image_variant_app.py @@ -413,10 +413,12 @@ def _refine_prompt_with_text_model(self, prompt: str) -> str: finally: self.text_engine.clear(session_id) - def _build_final_prompt(self, original_prompt: str, style: str, use_text_refiner: bool) -> Dict[str, str]: + def _build_final_prompt(self, original_prompt: str, style: str, use_text_refiner: bool) -> Dict[str, Any]: base = _cleanup_prompt(original_prompt) refined = "" + refiner_model_calls = 0 if use_text_refiner and self.text_engine.status().get("loaded"): + refiner_model_calls = 1 try: candidate = self._refine_prompt_with_text_model(base) if _prompt_is_related(base, candidate): @@ -429,7 +431,12 @@ def _build_final_prompt(self, original_prompt: str, style: str, use_text_refiner final_prompt = ", ".join(extras) if extras else base if "high detail" not in final_prompt.lower(): final_prompt = f"{final_prompt}, high detail, coherent composition" - return {"base_prompt": base, "refined_prompt": refined, "prompt_used": final_prompt} + return { + "base_prompt": base, + "refined_prompt": refined, + "prompt_used": final_prompt, + "refiner_model_calls": refiner_model_calls, + } def generate_image( self, @@ -499,6 +506,8 @@ def generate_image( "seed": seed_value, "guidance_scale": guidance_value, "timing_ms": total_ms, + "refiner_model_calls": int(prompt_payload.get("refiner_model_calls") or 0), + "model_calls": 1 + int(prompt_payload.get("refiner_model_calls") or 0), } meta_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") @@ -512,6 +521,8 @@ def generate_image( "metadata_path": str(meta_path), "image_url": f"/generated/{image_path.name}", "timing_ms": total_ms, + "refiner_model_calls": int(prompt_payload.get("refiner_model_calls") or 0), + "model_calls": 1 + int(prompt_payload.get("refiner_model_calls") or 0), } with self.lock: self.recent_images = [result, *self.recent_images[:5]] diff --git a/source/chat_web_app.py b/source/chat_web_app.py index 1808d96ae..269ac232e 100644 --- a/source/chat_web_app.py +++ b/source/chat_web_app.py @@ -12,6 +12,7 @@ import chat_app from device_utils import configure_torch_runtime, resolve_device +from route_policy_shadow_registry import RouteShadowAssignmentRegistry HTML = """ @@ -80,12 +81,13 @@ input,select,textarea{width:100%;border:1px solid var(--line);background:#fff;color:var(--ink);border-radius:7px;padding:10px 11px;font:inherit}input:focus,select:focus,textarea:focus{outline:none;border-color:var(--accent);box-shadow:0 0 0 3px rgba(8,127,91,.16)} .split,.btns{display:grid;grid-template-columns:1fr 1fr;gap:8px}button{border:0;border-radius:7px;padding:10px 12px;font:inherit;font-weight:800;cursor:pointer;background:var(--accent);color:#fff}button:hover{filter:brightness(.94)}button:disabled{opacity:.55;cursor:not-allowed} button.alt{background:#e7edf0;color:#1e2a31;border:1px solid var(--line)}.status{white-space:pre-wrap;background:#f6f8f9;border:1px solid var(--line);border-radius:7px;padding:11px;min-height:112px;color:#40505c;font:12px ui-monospace,SFMono-Regular,Consolas,monospace} +.shadow-status{display:grid;gap:8px;padding:10px;border:1px solid var(--line);border-radius:7px;background:#f6f8f9;color:#40505c;font:11px/1.45 ui-monospace,SFMono-Regular,Consolas,monospace;white-space:pre-wrap;overflow-wrap:anywhere}.shadow-head{display:flex;align-items:center;justify-content:space-between;gap:8px;font-family:Aptos,Segoe UI,sans-serif;font-weight:800;text-transform:uppercase;letter-spacing:.05em}.shadow-head button{padding:7px 9px;font-size:.72rem} .head{padding:14px 16px;border-bottom:1px solid var(--line);display:flex;justify-content:space-between;gap:12px;align-items:center;background:#fff}.head-title{font-weight:900}.head small{color:var(--muted)} .pillrow{display:flex;gap:6px;flex-wrap:wrap;justify-content:flex-end}.pill{border:1px solid var(--line);background:#f6f8f9;border-radius:999px;padding:5px 9px;color:#40505c;font-size:.76rem} .msgs{padding:16px;overflow:auto;display:flex;flex-direction:column;gap:12px;background:linear-gradient(180deg,#f8fafb,#eef3f5)}.msg{position:relative;border:1px solid var(--line);border-radius:8px;padding:12px 13px;max-width:min(78%,820px);white-space:pre-wrap;line-height:1.45;background:var(--bot);box-shadow:0 6px 18px rgba(23,32,38,.06)} .msg.user{align-self:flex-end;background:var(--user);border-color:#a8d2f2}.msg.bot{align-self:flex-start}.msg.pending{color:var(--muted);font-style:italic}.who{display:flex;justify-content:space-between;gap:12px;align-items:center;margin-bottom:6px;color:#4b5b65;font-size:.72rem;font-weight:900;text-transform:uppercase;letter-spacing:.07em} .copy{padding:4px 7px;border-radius:6px;background:#eef3f5;color:#31414c;border:1px solid var(--line);font-size:.72rem}.tim{margin-top:8px;color:#63717d;font-size:.76rem}details{margin-top:8px;border-top:1px solid var(--line);padding-top:8px;color:#63717d;font-size:.76rem}summary{cursor:pointer;font-weight:800;color:#40505c} -.comp{padding:13px 16px;border-top:1px solid var(--line);display:grid;grid-template-columns:1fr auto;gap:10px;align-items:end;background:#fff}textarea{min-height:54px;max-height:190px;resize:none;line-height:1.42}.hint{grid-column:1/-1;color:var(--muted);font-size:.76rem} +.comp{padding:13px 16px;border-top:1px solid var(--line);display:grid;grid-template-columns:1fr auto auto;gap:10px;align-items:end;background:#fff}textarea{min-height:54px;max-height:190px;resize:none;line-height:1.42}.hint{grid-column:1/-1;color:var(--muted);font-size:.76rem} .quick{display:flex;gap:6px;flex-wrap:wrap;margin-top:2px}.quick button{background:#f6f8f9;color:#31414c;border:1px solid var(--line);font-weight:700;padding:7px 9px;font-size:.78rem} @media (max-width:900px){.wrap{grid-template-columns:1fr;padding:10px}.chat{height:72vh;min-height:560px}.msg{max-width:94%}.split,.btns{grid-template-columns:1fr}} @@ -96,46 +98,124 @@
+
+
+
Loading status...
+
Shadow registry - read only
Not loaded. Execution, activation, and promotion are unavailable.
Web Chat
No model loaded
idle
-
Drafts and the local transcript are kept in this browser session.
+
Drafts and the local transcript are kept in this browser session.
""" +_RUNTIME_COMPUTE_DEFAULT_KEYS = ( + "reasoning_cycles", + "adaptive_compute", + "adaptive_exit_tol", + "adaptive_exit_entropy", + "prediction_stability_patience", + "prediction_stability_tol", +) + + +def _library_runtime_compute_defaults() -> Dict[str, Any]: + return { + "reasoning_cycles": None, + "adaptive_compute": False, + "adaptive_exit_tol": 1e-3, + "adaptive_exit_entropy": chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY, + "prediction_stability_patience": chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, + "prediction_stability_tol": chat_app.DEFAULT_PREDICTION_STABILITY_TOL, + } + + +def _normalize_runtime_compute_defaults(values: Dict[str, Any]) -> Dict[str, Any]: + raw_cycles = values.get("reasoning_cycles") + reasoning_cycles: Any + if chat_app._is_auto_reasoning_cycles(raw_cycles): + reasoning_cycles = "auto" + else: + reasoning_cycles = chat_app._coerce_optional_positive_int( + raw_cycles, + chat_app.MAX_RUNTIME_REASONING_CYCLES, + ) + return { + "reasoning_cycles": reasoning_cycles, + "adaptive_compute": chat_app._coerce_bool(values.get("adaptive_compute")), + "adaptive_exit_tol": chat_app._coerce_nonnegative_float( + values.get("adaptive_exit_tol"), + 1e-3, + ), + "adaptive_exit_entropy": chat_app._coerce_nonnegative_float( + values.get("adaptive_exit_entropy"), + chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY, + ), + "prediction_stability_patience": chat_app._coerce_nonnegative_int( + values.get("prediction_stability_patience"), + chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, + chat_app.MAX_RUNTIME_REASONING_CYCLES, + ), + "prediction_stability_tol": chat_app._coerce_nonnegative_float( + values.get("prediction_stability_tol"), + chat_app.DEFAULT_PREDICTION_STABILITY_TOL, + ), + } + + +def _runtime_compute_cli_overrides(args: argparse.Namespace) -> Dict[str, Any]: + """Return only compute options the CLI user actually supplied.""" + values = { + "reasoning_cycles": getattr(args, "reasoning_cycles", None), + "adaptive_compute": getattr(args, "adaptive_compute", None), + "adaptive_exit_tol": getattr(args, "adaptive_exit_tol", None), + "adaptive_exit_entropy": getattr(args, "adaptive_exit_entropy", None), + "prediction_stability_patience": getattr(args, "prediction_stability_patience", None), + "prediction_stability_tol": getattr(args, "prediction_stability_tol", None), + } + return {key: value for key, value in values.items() if value is not None} + + class Engine: def __init__(self, device: Any, device_info: Dict[str, Any], defaults: Dict[str, Any]): self.device = device self.device_info = dict(device_info or {}) - self.defaults = dict(defaults) + self._constructor_defaults = dict(defaults or {}) + self.defaults = self._build_effective_defaults({}) self.lock = threading.RLock() self.model = None self.weights_path: Optional[str] = None @@ -146,6 +226,32 @@ def __init__(self, device: Any, device_info: Dict[str, Any], defaults: Dict[str, self.available_labels: List[int] = list(range(chat_app.MODEL_CLASSES)) self.sessions: Dict[str, List[Tuple[str, str]]] = {} self.recent: Dict[str, List[str]] = {} + registry_path = self._constructor_defaults.get("route_shadow_registry_path") + self.route_shadow_registry_path = Path( + registry_path or Path("tmp") / "memory" / "route-policy-shadow-registry.sqlite3" + ).expanduser().resolve() + + def _build_effective_defaults(self, meta: Dict[str, Any]) -> Dict[str, Any]: + runtime_defaults = _library_runtime_compute_defaults() + metadata_defaults = meta.get("runtime_defaults") + if isinstance(metadata_defaults, dict): + runtime_defaults.update( + { + key: metadata_defaults[key] + for key in _RUNTIME_COMPUTE_DEFAULT_KEYS + if key in metadata_defaults + } + ) + runtime_defaults.update( + { + key: self._constructor_defaults[key] + for key in _RUNTIME_COMPUTE_DEFAULT_KEYS + if key in self._constructor_defaults + } + ) + effective = dict(self._constructor_defaults) + effective.update(_normalize_runtime_compute_defaults(runtime_defaults)) + return effective def status(self) -> Dict[str, Any]: with self.lock: @@ -158,8 +264,42 @@ def status(self) -> Dict[str, Any]: "available_labels": len(self.available_labels), "device": self.device_info.get("resolved", str(self.device)), "sessions": len(self.sessions), + "runtime_compute_supported": chat_app.model_supports_runtime_compute(self.model) if self.model is not None else False, + "reasoning_cycles": chat_app._format_reasoning_cycles_setting(self.defaults.get("reasoning_cycles")), + "adaptive_compute": bool(self.defaults.get("adaptive_compute", False)), + "adaptive_exit_tol": chat_app._coerce_nonnegative_float(self.defaults.get("adaptive_exit_tol", 1e-3), 1e-3), + "adaptive_exit_entropy": chat_app._coerce_nonnegative_float(self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + "prediction_stability_patience": chat_app._coerce_nonnegative_int(self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE, chat_app.MAX_RUNTIME_REASONING_CYCLES), + "prediction_stability_tol": chat_app._coerce_nonnegative_float(self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), chat_app.DEFAULT_PREDICTION_STABILITY_TOL), } + def route_shadow_registry_snapshot(self) -> Dict[str, Any]: + """Return compatible shadow-registry status without exposing mutations.""" + + registry_path = self.route_shadow_registry_path + if not registry_path.is_file(): + return { + "ok": True, + "available": False, + "status": "not_initialized", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + snapshot = RouteShadowAssignmentRegistry(registry_path, read_only=True).snapshot() + return { + **snapshot, + "available": True, + "status": "verified" if snapshot.get("ok") else "verification_failed", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + } + def _parse_buckets(self, meta: Dict[str, Any]) -> None: buckets: Dict[int, List[Dict[str, Any]]] = {} raw = meta.get("buckets", {}) @@ -184,13 +324,14 @@ def load(self, weights: str, meta_path: str) -> Dict[str, Any]: raise FileNotFoundError(f"Metadata not found: {meta_path}") meta = chat_app.load_metadata(meta_path) + effective_defaults = self._build_effective_defaults(meta) raw_feature_mode = str(meta.get("feature_mode", "legacy")).strip().lower() feature_mode = chat_app.resolve_feature_mode(raw_feature_mode, smarter_auto=True) sd = chat_app.safe_load_state_dict(weights) inferred = chat_app.detect_model_size_from_state_dict(sd) resolved_model_size, _ = chat_app.resolve_runtime_model_size( - str(self.defaults.get("model_size", "auto")), + str(effective_defaults.get("model_size", "auto")), str(meta.get("model_size", "")), inferred, ) @@ -223,6 +364,7 @@ def load(self, weights: str, meta_path: str) -> Dict[str, Any]: self.meta_path = meta_path self.feature_mode = feature_mode self.model_size = resolved_model_size + self.defaults = effective_defaults self._parse_buckets(meta) self.sessions.clear() self.recent.clear() @@ -234,7 +376,86 @@ def clear(self, session_id: str) -> None: self.sessions.pop(session_id, None) self.recent.pop(session_id, None) - def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None, response_temperature: Optional[float] = None, show_top_responses: int = 0) -> Dict[str, Any]: + def compute_sweep( + self, + session_id: str, + user_text: str, + cycles: Any = None, + adaptive_compute: Any = None, + adaptive_exit_tol: Any = None, + adaptive_exit_entropy: Any = None, + prediction_stability_patience: Any = None, + prediction_stability_tol: Any = None, + ) -> Dict[str, Any]: + if not user_text.strip(): + raise ValueError("Empty message") + requested_cycles: List[int] = [] + raw_cycles = cycles if isinstance(cycles, list) and cycles else [1, 3, 8] + for raw in raw_cycles: + parsed = chat_app._coerce_optional_positive_int(raw, chat_app.MAX_RUNTIME_REASONING_CYCLES) + if parsed is not None and parsed not in requested_cycles: + requested_cycles.append(parsed) + if not requested_cycles: + requested_cycles = [1, 3, 8] + + with self.lock: + if self.model is None: + raise RuntimeError("No model loaded") + model = self.model + feature_mode = self.feature_mode + labels = list(self.available_labels) + history = list(self.sessions.get(session_id, [])) + + context = chat_app.build_context(history, user_text=user_text, max_turns=int(self.defaults.get("max_turns", 2))) + x = chat_app.text_to_model_input(context, feature_mode=feature_mode).to(self.device) + idx = torch.tensor(labels, dtype=torch.long, device=self.device) + rows: List[Dict[str, Any]] = [] + + with torch.no_grad(): + for cycle_count in requested_cycles: + t0 = time.perf_counter() + model_out, compute_diag = chat_app.forward_with_runtime_compute( + model, + x, + reasoning_cycles=cycle_count, + adaptive_compute=adaptive_compute if adaptive_compute is not None else self.defaults.get("adaptive_compute", False), + exit_tol=adaptive_exit_tol if adaptive_exit_tol is not None else self.defaults.get("adaptive_exit_tol", 1e-3), + exit_entropy_threshold=adaptive_exit_entropy if adaptive_exit_entropy is not None else self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + prediction_stability_patience=prediction_stability_patience if prediction_stability_patience is not None else self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), + prediction_stability_tol=prediction_stability_tol if prediction_stability_tol is not None else self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), + ) + latency_ms = round((time.perf_counter() - t0) * 1000, 1) + logits = model_out[0, 0] + avail_logits = logits.index_select(0, idx) + probs = torch.softmax(avail_logits, dim=0) + top_pos = int(torch.argmax(probs).item()) + entropy = -torch.sum(probs * torch.log(probs + 1e-9)) + rows.append({ + "requested_cycles": cycle_count, + "latency_ms": latency_ms, + "cycles_used": compute_diag.get("cycles_used"), + "predicted_label": int(labels[top_pos]), + "confidence": float(probs[top_pos].item()), + "entropy": float(entropy.item()), + "compute": compute_diag, + }) + + return {"ok": True, "session_id": session_id, "rows": rows} + + def chat( + self, + session_id: str, + user_text: str, + style_mode: Optional[str] = None, + response_temperature: Optional[float] = None, + show_top_responses: int = 0, + reasoning_cycles: Any = None, + adaptive_compute: Any = None, + adaptive_exit_tol: Any = None, + adaptive_exit_entropy: Any = None, + prediction_stability_patience: Any = None, + prediction_stability_tol: Any = None, + ) -> Dict[str, Any]: if not user_text.strip(): raise ValueError("Empty message") with self.lock: @@ -254,7 +475,18 @@ def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None tt = time.perf_counter() x = chat_app.text_to_model_input(context, feature_mode=feature_mode).to(self.device) with torch.no_grad(): - logits = model(x)[0, 0] + model_out, compute_diag = chat_app.forward_with_runtime_compute( + model, + x, + reasoning_cycles=reasoning_cycles if reasoning_cycles is not None else self.defaults.get("reasoning_cycles"), + adaptive_compute=adaptive_compute if adaptive_compute is not None else self.defaults.get("adaptive_compute", False), + exit_tol=adaptive_exit_tol if adaptive_exit_tol is not None else self.defaults.get("adaptive_exit_tol", 1e-3), + exit_entropy_threshold=adaptive_exit_entropy if adaptive_exit_entropy is not None else self.defaults.get("adaptive_exit_entropy", chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY), + prediction_stability_patience=prediction_stability_patience if prediction_stability_patience is not None else self.defaults.get("prediction_stability_patience", chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE), + prediction_stability_tol=prediction_stability_tol if prediction_stability_tol is not None else self.defaults.get("prediction_stability_tol", chat_app.DEFAULT_PREDICTION_STABILITY_TOL), + auto_reasoning_context=context, + ) + logits = model_out[0, 0] t_infer += time.perf_counter() - tt idx = torch.tensor(labels, dtype=torch.long, device=logits.device) @@ -351,7 +583,9 @@ def chat(self, session_id: str, user_text: str, style_mode: Optional[str] = None "infer": round(t_infer * 1000, 1), "rank_pick": round(t_rank * 1000, 1), "total": round((time.perf_counter() - t0) * 1000, 1), + "cycles_used": compute_diag.get("cycles_used"), }, + "compute": compute_diag, "top_candidates": top_candidates, } @@ -375,6 +609,15 @@ def index(): def api_status(): return jsonify({"ok": True, "status": engine.status()}) + @app.get('/api/route_shadow_registry/status') + def api_route_shadow_registry_status(): + try: + response = jsonify({"ok": True, "route_shadow_registry": engine.route_shadow_registry_snapshot()}) + response.headers["Cache-Control"] = "no-store" + return response + except Exception as e: + return jsonify({"ok": False, "error": str(e)}), 500 + @app.post('/api/load') def api_load(): p = request.get_json(force=True, silent=True) or {} @@ -397,6 +640,31 @@ def api_chat(): style_mode=p.get('style_mode'), response_temperature=p.get('response_temperature'), show_top_responses=int(p.get('show_top_responses') or 0), + reasoning_cycles=p.get('reasoning_cycles'), + adaptive_compute=p.get('adaptive_compute'), + adaptive_exit_tol=p.get('adaptive_exit_tol'), + adaptive_exit_entropy=p.get('adaptive_exit_entropy'), + prediction_stability_patience=p.get('prediction_stability_patience'), + prediction_stability_tol=p.get('prediction_stability_tol'), + )) + except Exception as e: + return jsonify({"ok": False, "error": str(e)}), 400 + + @app.post('/api/compute_sweep') + def api_compute_sweep(): + p = request.get_json(force=True, silent=True) or {} + sid = str(p.get('session_id') or '').strip() or str(uuid.uuid4()) + msg = str(p.get('message') or '').strip() + try: + return jsonify(engine.compute_sweep( + session_id=sid, + user_text=msg, + cycles=p.get('cycles'), + adaptive_compute=p.get('adaptive_compute'), + adaptive_exit_tol=p.get('adaptive_exit_tol'), + adaptive_exit_entropy=p.get('adaptive_exit_entropy'), + prediction_stability_patience=p.get('prediction_stability_patience'), + prediction_stability_tol=p.get('prediction_stability_tol'), )) except Exception as e: return jsonify({"ok": False, "error": str(e)}), 400 @@ -434,6 +702,25 @@ def main() -> None: ap.add_argument('--temperature', type=float, default=0.0) ap.add_argument('--style_mode', choices=['auto','balanced','creative','concise','analyst'], default='auto') ap.add_argument('--creativity', type=float, default=0.2) + ap.add_argument('--reasoning_cycles', type=str, default=None) + adaptive_compute_group = ap.add_mutually_exclusive_group() + adaptive_compute_group.add_argument( + '--adaptive_compute', + dest='adaptive_compute', + action='store_true', + help='enable adaptive compute, overriding checkpoint metadata', + ) + adaptive_compute_group.add_argument( + '--no_adaptive_compute', + dest='adaptive_compute', + action='store_false', + help='disable adaptive compute, overriding checkpoint metadata', + ) + ap.set_defaults(adaptive_compute=None) + ap.add_argument('--adaptive_exit_tol', type=float, default=None) + ap.add_argument('--adaptive_exit_entropy', type=float, default=None) + ap.add_argument('--prediction_stability_patience', type=int, default=None) + ap.add_argument('--prediction_stability_tol', type=float, default=None) args = ap.parse_args() configure_torch_runtime( @@ -443,7 +730,7 @@ def main() -> None: matmul_precision=str(args.matmul_precision), ) device, device_info = resolve_device(args.device, preference=args.device_preference) - engine = Engine(device, device_info, { + engine_defaults = { 'model_size': args.model_size, 'max_turns': int(args.max_turns), 'top_labels': int(args.top_labels), @@ -452,7 +739,9 @@ def main() -> None: 'temperature': float(args.temperature), 'style_mode': str(args.style_mode), 'creativity': float(args.creativity), - }) + } + engine_defaults.update(_runtime_compute_cli_overrides(args)) + engine = Engine(device, device_info, engine_defaults) if args.autoload: try: print(engine.load(args.weights, args.meta)) diff --git a/source/multimodel_memory.py b/source/multimodel_memory.py index 9eb9ce9df..ae6009b70 100644 --- a/source/multimodel_memory.py +++ b/source/multimodel_memory.py @@ -1,6 +1,8 @@ from __future__ import annotations +import hashlib import json +import math import re import time from pathlib import Path @@ -9,6 +11,70 @@ TOKEN_RE = re.compile(r"[a-z0-9]{3,}", re.IGNORECASE) MULTISPACE_RE = re.compile(r"\s+") +ROUTE_MODES = ("off", "collective", "loop", "collective_loop") +ROUTE_FEEDBACK_DECAY = 0.6 +ROUTE_FEEDBACK_CONFIDENCE_Z = 1.645 +ROUTE_PREFERENCE_SIGNAL_FLOOR = 0.75 +ROUTE_EVIDENCE_INTERVAL_KIND = "recency_weighted_wilson_heuristic" +FEEDBACK_RATINGS = { + "up": ("up", 1), + "good": ("up", 1), + "approve": ("up", 1), + "approved": ("up", 1), + "positive": ("up", 1), + "down": ("down", -1), + "bad": ("down", -1), + "reject": ("down", -1), + "rejected": ("down", -1), + "negative": ("down", -1), +} +FEEDBACK_INTENTS = { + "good": { + "quality": 1, + "depth_preference": 0, + "cost_pressure": 0, + "latency_pressure": 0, + }, + "bad_quality": { + "quality": -1, + "depth_preference": 0, + "cost_pressure": 0, + "latency_pressure": 0, + }, + "needs_deeper": { + "quality": None, + "depth_preference": 1, + "cost_pressure": 0, + "latency_pressure": 0, + }, + "too_costly": { + "quality": None, + "depth_preference": -1, + "cost_pressure": 1, + "latency_pressure": 0, + }, + "too_slow": { + "quality": None, + "depth_preference": -1, + "cost_pressure": 0, + "latency_pressure": 1, + }, +} +FEEDBACK_INTENT_ALIASES = { + "up": "good", + "positive": "good", + "approve": "good", + "approved": "good", + "bad": "bad_quality", + "down": "bad_quality", + "negative": "bad_quality", + "deeper": "needs_deeper", + "more_depth": "needs_deeper", + "cost": "too_costly", + "expensive": "too_costly", + "slow": "too_slow", + "latency": "too_slow", +} MEMORY_PATTERNS = ( ( @@ -73,34 +139,448 @@ def _now_ts() -> float: return float(time.time()) +def _safe_float(value: Any, *, limit: float = 1_000_000.0, digits: int = 3) -> Optional[float]: + try: + cooked = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(cooked): + return None + return round(max(0.0, min(limit, cooked)), digits) + + +def _safe_int(value: Any, *, limit: int = 1_000_000) -> Optional[int]: + cooked = _safe_float(value, limit=float(limit), digits=0) + if cooked is None: + return None + return int(cooked) + + +def _compact_route_economics(raw: Any, *, trace: Dict[str, Any], policy: Dict[str, Any]) -> Dict[str, Any]: + source = raw if isinstance(raw, dict) else {} + if not source and isinstance(trace.get("route_economics"), dict): + source = trace["route_economics"] + + estimate_src = source.get("estimate") if isinstance(source.get("estimate"), dict) else {} + actual_src = source.get("actual") if isinstance(source.get("actual"), dict) else {} + if not estimate_src and isinstance(policy.get("route_economics_estimate"), dict): + estimate_src = policy["route_economics_estimate"] + if not actual_src and isinstance(policy.get("route_economics_actual"), dict): + actual_src = policy["route_economics_actual"] + + estimate: Dict[str, Any] = {} + for key in ("selected_agent_mode", "action_mode", "budget_profile", "cost_preference", "latency_tier"): + text = _norm(estimate_src.get(key) or "", limit=80) + if text: + estimate[key] = text + for key in ("estimated_model_calls", "estimated_tool_calls", "planned_loop_steps", "collective_model_count"): + value = _safe_int(estimate_src.get(key), limit=10_000) + if value is not None: + estimate[key] = value + cost = _safe_float(estimate_src.get("estimated_cost_units"), limit=100_000.0, digits=3) + if cost is not None: + estimate["estimated_cost_units"] = cost + + actual: Dict[str, Any] = {} + for key in ("latency_tier",): + text = _norm(actual_src.get(key) or "", limit=80) + if text: + actual[key] = text + for key in ("model_calls", "tool_calls", "loop_steps", "consultation_count"): + value = _safe_int(actual_src.get(key), limit=10_000) + if value is not None: + actual[key] = value + for key in ("elapsed_ms", "cost_units"): + value = _safe_float(actual_src.get(key), limit=100_000_000.0 if key == "elapsed_ms" else 100_000.0, digits=3) + if value is not None: + actual[key] = value + + compact: Dict[str, Any] = {} + if estimate: + compact["estimate"] = estimate + if actual: + compact["actual"] = actual + return compact + + +def _route_economics_summary(rows: Sequence[Dict[str, Any]]) -> Dict[str, Any]: + actual_rows: List[Dict[str, Any]] = [] + for row in rows: + economics = row.get("route_economics") if isinstance(row.get("route_economics"), dict) else {} + actual = economics.get("actual") if isinstance(economics.get("actual"), dict) else {} + if actual: + actual_rows.append(actual) + + def avg(key: str) -> Optional[float]: + values = [ + float(value) + for value in (actual.get(key) for actual in actual_rows) + if isinstance(value, (int, float)) + ] + if not values: + return None + return round(sum(values) / len(values), 3) + + cost_values = [ + float(value) + for value in (actual.get("cost_units") for actual in actual_rows) + if isinstance(value, (int, float)) + ] + latency_values = [ + float(value) + for value in (actual.get("elapsed_ms") for actual in actual_rows) + if isinstance(value, (int, float)) + ] + summary: Dict[str, Any] = {"sample_count": len(actual_rows)} + if cost_values: + summary["avg_cost_units"] = round(sum(cost_values) / len(cost_values), 3) + summary["total_cost_units"] = round(sum(cost_values), 3) + summary["max_cost_units"] = round(max(cost_values), 3) + if latency_values: + summary["avg_elapsed_ms"] = round(sum(latency_values) / len(latency_values), 3) + summary["max_elapsed_ms"] = round(max(latency_values), 3) + model_calls = avg("model_calls") + if model_calls is not None: + summary["avg_model_calls"] = model_calls + tool_calls = avg("tool_calls") + if tool_calls is not None: + summary["avg_tool_calls"] = tool_calls + return summary + + +def _feedback_axes(row: Dict[str, Any]) -> Dict[str, Any]: + stored = row.get("feedback_axes") if isinstance(row.get("feedback_axes"), dict) else {} + if stored: + quality_raw = stored.get("quality") + quality = None if quality_raw is None else max(-1, min(1, int(quality_raw))) + return { + "quality": quality, + "depth_preference": max(-1, min(1, int(stored.get("depth_preference") or 0))), + "cost_pressure": max(0, min(1, int(stored.get("cost_pressure") or 0))), + "latency_pressure": max(0, min(1, int(stored.get("latency_pressure") or 0))), + } + + # Old route-feedback rows only had a scalar score. Treat them as generic + # quality feedback so existing sessions keep their previous semantics. + return { + "quality": max(-1, min(1, int(row.get("score_delta") or 0))), + "depth_preference": 0, + "cost_pressure": 0, + "latency_pressure": 0, + } + + +def _weighted_wilson_interval( + *, + positive_weight: float, + total_weight: float, + squared_weight: float, + z: float = ROUTE_FEEDBACK_CONFIDENCE_Z, +) -> Dict[str, Optional[float]]: + if total_weight <= 0.0 or squared_weight <= 0.0: + return { + "effective_sample_size": 0.0, + "quality_lower_bound": None, + "quality_upper_bound": None, + "quality_confidence_width": None, + } + effective_n = (total_weight * total_weight) / squared_weight + probability = max(0.0, min(1.0, positive_weight / total_weight)) + z2 = z * z + denominator = 1.0 + (z2 / effective_n) + center = (probability + (z2 / (2.0 * effective_n))) / denominator + radius = ( + z + * math.sqrt( + max(0.0, (probability * (1.0 - probability) / effective_n) + (z2 / (4.0 * effective_n * effective_n))) + ) + / denominator + ) + lower = max(0.0, center - radius) + upper = min(1.0, center + radius) + return { + "effective_sample_size": round(effective_n, 3), + "quality_lower_bound": round(lower, 3), + "quality_upper_bound": round(upper, 3), + "quality_confidence_width": round(upper - lower, 3), + } + + +def _compact_auto_agent_policy(policy: Dict[str, Any], selected_mode: str) -> Dict[str, Any]: + raw_allowed_modes = policy.get("allowed_agent_modes") or policy.get("eligible_agent_modes") or [] + if not isinstance(raw_allowed_modes, (list, tuple)): + raw_allowed_modes = [] + allowed_modes = [ + _norm(item, limit=48) + for item in raw_allowed_modes + if _norm(item, limit=48) in ROUTE_MODES + ] + action_probabilities: Dict[str, float] = {} + raw_probabilities = policy.get("action_probabilities") + if isinstance(raw_probabilities, dict): + for mode in ROUTE_MODES: + value = raw_probabilities.get(mode) + if isinstance(value, (int, float)): + action_probabilities[mode] = round(max(0.0, min(1.0, float(value))), 6) + post_filter_probabilities: Dict[str, float] = {} + raw_post_filter_probabilities = policy.get("post_filter_action_probabilities") + if isinstance(raw_post_filter_probabilities, dict): + for mode in ROUTE_MODES: + value = raw_post_filter_probabilities.get(mode) + if isinstance(value, (int, float)): + post_filter_probabilities[mode] = round(max(0.0, min(1.0, float(value))), 6) + raw_eligible_actions = policy.get("eligible_actions") + eligible_actions = [ + _norm(item, limit=48) + for item in (raw_eligible_actions if isinstance(raw_eligible_actions, (list, tuple)) else []) + if _norm(item, limit=48) in ROUTE_MODES + ] + raw_context = policy.get("decision_context") if isinstance(policy.get("decision_context"), dict) else {} + decision_context = { + "action_mode": _norm(raw_context.get("action_mode") or policy.get("action_mode") or "", limit=32), + "budget_profile": _norm(raw_context.get("budget_profile") or policy.get("budget_profile") or "", limit=32), + "score": raw_context.get("score", policy.get("score")), + "allowed_agent_modes": [ + _norm(item, limit=48) + for item in ( + raw_context.get("allowed_agent_modes") + if isinstance(raw_context.get("allowed_agent_modes"), (list, tuple)) + else allowed_modes + ) + if _norm(item, limit=48) in ROUTE_MODES + ], + } + logging_propensity = policy.get("logging_propensity") + logging_support = policy.get("logging_support") if isinstance(policy.get("logging_support"), dict) else {} + return { + "policy_id": _norm(policy.get("policy_id") or "", limit=80), + "policy_version": _norm(policy.get("policy_version") or "", limit=80), + "feature_schema_version": _norm(policy.get("feature_schema_version") or "", limit=80), + "decision_type": _norm(policy.get("decision_type") or "legacy_unknown", limit=40), + "action_mode": _norm(policy.get("action_mode") or "", limit=32), + "score": policy.get("score"), + "score_before_budget": policy.get("score_before_budget"), + "reason": _norm(policy.get("reason") or "", limit=120), + "reasons": [_norm(item, limit=80) for item in list(policy.get("reasons") or [])[:6]], + "base_selected_agent_mode": _norm(policy.get("base_selected_agent_mode") or "", limit=48), + "selected_agent_mode": _norm(policy.get("selected_agent_mode") or selected_mode, limit=48), + "allowed_agent_modes": allowed_modes, + "eligible_actions": eligible_actions, + "budget_profile": _norm(policy.get("budget_profile") or "", limit=32), + "logging_propensity": ( + round(max(0.0, min(1.0, float(logging_propensity))), 6) + if isinstance(logging_propensity, (int, float)) + else None + ), + "action_probabilities": action_probabilities, + "post_filter_action_probabilities": post_filter_probabilities, + "probability_stage": _norm(policy.get("probability_stage") or "", limit=32), + "support_schema_version": _norm(logging_support.get("schema_version") or "", limit=80), + "candidate_set_hash": _norm(policy.get("candidate_set_hash") or "", limit=80), + "distribution_hash": _norm(policy.get("distribution_hash") or "", limit=80), + "decision_context": decision_context, + } + + +def _route_adaptive_feedback_summary( + rows: Sequence[Dict[str, Any]], + economics: Optional[Dict[str, Any]] = None, + *, + decay: float = ROUTE_FEEDBACK_DECAY, +) -> Dict[str, Any]: + if not rows: + return { + "sample_count": 0, + "quality_sample_count": 0, + "decay": decay, + "weighted_count": 0.0, + "weighted_net": 0.0, + "weighted_positive": 0.0, + "weighted_negative": 0.0, + "quality_score": None, + "quality_cost_score": None, + "effective_sample_size": 0.0, + "quality_lower_bound": None, + "quality_upper_bound": None, + "quality_confidence_width": None, + "quality_cost_lower_bound": None, + "quality_cost_upper_bound": None, + "confidence_level": 0.9, + "confidence_status": "no_evidence", + "interval_kind": ROUTE_EVIDENCE_INTERVAL_KIND, + "coverage_claim": "heuristic_associational_only", + "effective_sample_size_ceiling": round((1.0 + decay) / (1.0 - decay), 3), + "weighted_depth_preference": 0.0, + "weighted_cost_pressure": 0.0, + "weighted_latency_pressure": 0.0, + "preference_direction": None, + "preference_signal": False, + "recent_negative_rate": None, + "regression_signal": False, + } + + weighted_count = 0.0 + squared_weight = 0.0 + weighted_net = 0.0 + weighted_positive = 0.0 + weighted_negative = 0.0 + weighted_depth_preference = 0.0 + weighted_cost_pressure = 0.0 + weighted_latency_pressure = 0.0 + quality_rows: List[Dict[str, Any]] = [] + for idx, row in enumerate(rows): + age = len(rows) - idx - 1 + weight = float(decay) ** age + axes = _feedback_axes(row) + weighted_depth_preference += weight * int(axes["depth_preference"] or 0) + weighted_cost_pressure += weight * int(axes["cost_pressure"] or 0) + weighted_latency_pressure += weight * int(axes["latency_pressure"] or 0) + if axes["quality"] is None: + continue + delta = int(axes["quality"] or 0) + quality_rows.append(row) + weighted_count += weight + squared_weight += weight * weight + weighted_net += weight * delta + if delta > 0: + weighted_positive += weight + elif delta < 0: + weighted_negative += weight + + recent_tail = quality_rows[-3:] + recent_negative = sum(1 for row in recent_tail if int(_feedback_axes(row)["quality"] or 0) < 0) + recent_negative_rate = round(recent_negative / len(recent_tail), 3) if recent_tail else None + quality_score = None + if weighted_count > 0: + quality_score = round(max(0.0, min(1.0, ((weighted_net / weighted_count) + 1.0) / 2.0)), 3) + + confidence = _weighted_wilson_interval( + positive_weight=weighted_positive, + total_weight=weighted_count, + squared_weight=squared_weight, + ) + effective_sample_size = float(confidence["effective_sample_size"] or 0.0) + if effective_sample_size >= 3.0 and len(quality_rows) >= 4: + confidence_status = "established" + elif effective_sample_size >= 1.5 and len(quality_rows) >= 2: + confidence_status = "emerging" + elif quality_rows: + confidence_status = "sparse" + else: + confidence_status = "no_quality_evidence" + + quality_cost_score = None + quality_cost_lower_bound = None + quality_cost_upper_bound = None + avg_cost = None + if isinstance(economics, dict): + raw_cost = economics.get("avg_cost_units") + if isinstance(raw_cost, (int, float)): + avg_cost = float(raw_cost) + if quality_score is not None: + cost_divisor = 1.0 + (max(0.0, avg_cost or 0.0) / 10.0) + quality_cost_score = round(quality_score / cost_divisor, 3) + lower = confidence.get("quality_lower_bound") + upper = confidence.get("quality_upper_bound") + if isinstance(lower, (int, float)): + quality_cost_lower_bound = round(float(lower) / cost_divisor, 3) + if isinstance(upper, (int, float)): + quality_cost_upper_bound = round(float(upper) / cost_divisor, 3) + + if weighted_depth_preference >= ROUTE_PREFERENCE_SIGNAL_FLOOR: + preference_direction = "deeper" + elif weighted_depth_preference <= -ROUTE_PREFERENCE_SIGNAL_FLOOR: + preference_direction = "shallower" + else: + preference_direction = None + + regression_signal = bool( + len(quality_rows) >= 3 + and weighted_negative >= 1.45 + and weighted_net <= -0.25 + and (recent_negative_rate or 0.0) >= 0.667 + ) + return { + "sample_count": len(rows), + "quality_sample_count": len(quality_rows), + "decay": decay, + "weighted_count": round(weighted_count, 3), + "weighted_net": round(weighted_net, 3), + "weighted_positive": round(weighted_positive, 3), + "weighted_negative": round(weighted_negative, 3), + "quality_score": quality_score, + "quality_cost_score": quality_cost_score, + **confidence, + "quality_cost_lower_bound": quality_cost_lower_bound, + "quality_cost_upper_bound": quality_cost_upper_bound, + "confidence_level": 0.9, + "confidence_status": confidence_status, + "interval_kind": ROUTE_EVIDENCE_INTERVAL_KIND, + "coverage_claim": "heuristic_associational_only", + "effective_sample_size_ceiling": round((1.0 + decay) / (1.0 - decay), 3), + "weighted_depth_preference": round(weighted_depth_preference, 3), + "weighted_cost_pressure": round(weighted_cost_pressure, 3), + "weighted_latency_pressure": round(weighted_latency_pressure, 3), + "preference_direction": preference_direction, + "preference_signal": preference_direction is not None, + "recent_negative_rate": recent_negative_rate, + "regression_signal": regression_signal, + } + + class ConversationMemoryStore: def __init__(self, root_dir: Path) -> None: self.root_dir = root_dir.resolve() self.root_dir.mkdir(parents=True, exist_ok=True) def _path_for(self, session_id: str) -> Path: + session_digest = hashlib.sha256(str(session_id).encode("utf-8")).hexdigest() + return self.root_dir / f"{_safe_slug(session_id)}-{session_digest}.json" + + def _legacy_path_for(self, session_id: str) -> Path: return self.root_dir / f"{_safe_slug(session_id)}.json" - def load_session(self, session_id: str) -> Dict[str, Any]: - path = self._path_for(session_id) - if not path.exists(): - now = _now_ts() - return { - "session_id": session_id, - "created_at": now, - "updated_at": now, - "memories": [], - "turns": [], - } + @staticmethod + def _read_session_file(path: Path) -> Optional[Dict[str, Any]]: try: payload = json.loads(path.read_text(encoding="utf-8")) except Exception: - payload = {} - payload.setdefault("session_id", session_id) + return None + return payload if isinstance(payload, dict) else None + + @staticmethod + def _empty_session(session_id: str) -> Dict[str, Any]: + now = _now_ts() + return { + "session_id": session_id, + "created_at": now, + "updated_at": now, + "memories": [], + "turns": [], + "route_feedback": [], + "route_usage": [], + } + + def load_session(self, session_id: str) -> Dict[str, Any]: + path = self._path_for(session_id) + if not path.exists(): + legacy_path = self._legacy_path_for(session_id) + legacy_payload = self._read_session_file(legacy_path) if legacy_path.exists() else None + if legacy_payload is None or legacy_payload.get("session_id") != session_id: + return self._empty_session(session_id) + self.save_session(session_id, legacy_payload) + legacy_path.unlink(missing_ok=True) + payload = self._read_session_file(path) + else: + payload = self._read_session_file(path) + if payload is None or payload.get("session_id") != session_id: + return self._empty_session(session_id) payload.setdefault("created_at", _now_ts()) payload.setdefault("updated_at", payload["created_at"]) payload.setdefault("memories", []) payload.setdefault("turns", []) + payload.setdefault("route_feedback", []) + payload.setdefault("route_usage", []) return payload def save_session(self, session_id: str, payload: Dict[str, Any]) -> None: @@ -108,10 +588,19 @@ def save_session(self, session_id: str, payload: Dict[str, Any]) -> None: payload["session_id"] = session_id payload["updated_at"] = _now_ts() path = self._path_for(session_id) - path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8") + temp_path = path.with_suffix(path.suffix + ".tmp") + try: + temp_path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8") + temp_path.replace(path) + finally: + temp_path.unlink(missing_ok=True) def clear_session(self, session_id: str) -> None: self._path_for(session_id).unlink(missing_ok=True) + legacy_path = self._legacy_path_for(session_id) + legacy_payload = self._read_session_file(legacy_path) if legacy_path.exists() else None + if legacy_payload is not None and legacy_payload.get("session_id") == session_id: + legacy_path.unlink(missing_ok=True) def _extract_memories(self, user_text: str, assistant_text: str) -> List[Dict[str, Any]]: found: List[Dict[str, Any]] = [] @@ -183,6 +672,230 @@ def update( self.save_session(session_id, payload) return payload + def prepare_feedback(self, *, session_id: str, feedback: Dict[str, Any]) -> Dict[str, Any]: + """Normalize one feedback request without mutating the JSON mirror.""" + + payload = self.load_session(session_id) + route_id = _norm(feedback.get("route_id") or "", limit=120) + usage_row = next( + ( + row + for row in reversed(list(payload.get("route_usage") or [])) + if isinstance(row, dict) and route_id and _norm(row.get("route_id") or "", limit=120) == route_id + ), + None, + ) + authoritative = usage_row if isinstance(usage_row, dict) else {} + trace = feedback.get("agent_trace") if isinstance(feedback.get("agent_trace"), dict) else {} + policy = authoritative.get("auto_agent_policy") if authoritative else feedback.get("auto_agent_policy") + if not isinstance(policy, dict): + policy = trace.get("auto_agent_policy") if isinstance(trace.get("auto_agent_policy"), dict) else {} + + rating_raw = str(feedback.get("rating") or "").strip().lower() + rating_pair = FEEDBACK_RATINGS.get(rating_raw) + if rating_pair is None: + raise ValueError("rating must be one of: up, down, good, bad, approve, reject") + rating, score_delta = rating_pair + + intent_raw = str(feedback.get("feedback_intent") or feedback.get("intent") or "").strip().lower() + intent_key = FEEDBACK_INTENT_ALIASES.get(intent_raw, intent_raw) + if not intent_key: + intent_key = "good" if score_delta > 0 else "bad_quality" + if intent_key not in FEEDBACK_INTENTS: + raise ValueError( + "feedback_intent must be one of: good, bad_quality, needs_deeper, too_costly, too_slow" + ) + feedback_axes = dict(FEEDBACK_INTENTS[intent_key]) + feedback_tags: List[str] = [] + raw_tags = feedback.get("feedback_tags") if isinstance(feedback.get("feedback_tags"), list) else [] + for raw_tag in raw_tags[:4]: + tag = FEEDBACK_INTENT_ALIASES.get(str(raw_tag).strip().lower(), str(raw_tag).strip().lower()) + if tag in FEEDBACK_INTENTS and tag != intent_key and tag not in feedback_tags: + feedback_tags.append(tag) + tag_axes = FEEDBACK_INTENTS[tag] + for axis_name in ("depth_preference", "cost_pressure", "latency_pressure"): + feedback_axes[axis_name] = max( + -1, + min(1, int(feedback_axes.get(axis_name) or 0) + int(tag_axes.get(axis_name) or 0)), + ) + quality_score_delta = feedback_axes.get("quality") + quality_score_delta = int(quality_score_delta) if quality_score_delta is not None else 0 + + selected = _norm( + authoritative.get("selected_agent_mode") + or feedback.get("selected_agent_mode") + or feedback.get("resolved_agent_mode") + or trace.get("resolved_agent_mode") + or policy.get("selected_agent_mode") + or "off", + limit=48, + ) + if selected not in ROUTE_MODES: + selected = "off" + + row = { + "ts": _now_ts(), + "route_id": route_id, + "prompt": _norm( + authoritative.get("prompt") or feedback.get("prompt") or feedback.get("message") or "", + limit=1200, + ), + "response": _norm(feedback.get("response") or "", limit=1200), + "selected_agent_mode": selected, + "rating": rating, + "score_delta": quality_score_delta, + "feedback_intent": intent_key, + "feedback_tags": feedback_tags, + "feedback_axes": feedback_axes, + "reason": _norm(feedback.get("reason") or "", limit=240), + "auto_agent_policy": _compact_auto_agent_policy(policy, selected), + "model_key": _norm(authoritative.get("model_key") or feedback.get("model_key") or "", limit=120), + "route_reason": _norm( + authoritative.get("route_reason") or feedback.get("route_reason") or "", + limit=240, + ), + "evidence_source": "server_route_join" if authoritative else "legacy_unjoined", + "feedback_revision": 1, + } + route_economics = _compact_route_economics( + authoritative.get("route_economics") if authoritative else feedback.get("route_economics"), + trace=trace, + policy=policy, + ) + if route_economics: + row["route_economics"] = route_economics + return row + + def commit_feedback( + self, + *, + session_id: str, + feedback_row: Dict[str, Any], + feedback_revision: Optional[int] = None, + ) -> Dict[str, Any]: + """Upsert a prepared row, optionally pinned to its durable revision. + + When ``feedback_revision`` is supplied, retries replace the same mirror + revision and stale retries cannot roll the compatibility view backward. + The durable SQLite ledger remains the source of truth. + """ + + if not isinstance(feedback_row, dict): + raise ValueError("feedback_row must be a JSON object") + row = dict(feedback_row) + route_id = _norm(row.get("route_id") or "", limit=120) + payload = self.load_session(session_id) + rows = list(payload.get("route_feedback") or []) + previous_index = next( + ( + idx + for idx in range(len(rows) - 1, -1, -1) + if isinstance(rows[idx], dict) and route_id and _norm(rows[idx].get("route_id") or "", limit=120) == route_id + ), + None, + ) + previous_revision = ( + int(rows[previous_index].get("feedback_revision") or 1) + if previous_index is not None + else 0 + ) + durable_revision = feedback_revision is not None + if durable_revision: + if isinstance(feedback_revision, bool): + raise ValueError("feedback_revision must be a positive integer") + try: + revision = int(feedback_revision) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("feedback_revision must be a positive integer") from exc + if revision < 1: + raise ValueError("feedback_revision must be a positive integer") + if previous_index is not None and previous_revision > revision: + previous = dict(rows[previous_index]) + return { + "ok": True, + "feedback": previous, + "summary": self.route_feedback_summary(session_id, previous.get("prompt") or ""), + "compatibility_mirror": { + "status": "current_newer_revision", + "durable_revision": revision, + "mirrored_revision": previous_revision, + }, + } + else: + revision = previous_revision + 1 + + row["feedback_revision"] = revision + if durable_revision: + row["durable_feedback_revision"] = revision + idempotent = previous_index is not None and previous_revision == revision + if idempotent: + row["ts"] = rows[previous_index].get("ts", row.get("ts")) + if previous_index is None: + rows.append(row) + else: + rows[previous_index] = row + payload["route_feedback"] = rows[-120:] + self.save_session(session_id, payload) + return { + "ok": True, + "feedback": row, + "summary": self.route_feedback_summary(session_id, row["prompt"]), + "compatibility_mirror": { + "status": "committed", + "durable_revision": revision if durable_revision else None, + "mirrored_revision": revision, + "idempotent": idempotent, + }, + } + + def add_feedback(self, *, session_id: str, feedback: Dict[str, Any]) -> Dict[str, Any]: + """Backward-compatible local feedback write without a durable revision.""" + + row = self.prepare_feedback(session_id=session_id, feedback=feedback) + return self.commit_feedback(session_id=session_id, feedback_row=row) + + def add_route_usage( + self, + *, + session_id: str, + route_id: str, + prompt: str, + selected_agent_mode: str, + route_economics: Dict[str, Any], + auto_agent_policy: Optional[Dict[str, Any]] = None, + route_reason: str = "", + model_key: str = "", + ) -> Dict[str, Any]: + payload = self.load_session(session_id) + mode = _norm(selected_agent_mode or "off", limit=48) + if mode not in ROUTE_MODES: + mode = "off" + policy = auto_agent_policy if isinstance(auto_agent_policy, dict) else {} + compact_economics = _compact_route_economics( + route_economics, + trace={"route_economics": route_economics}, + policy=policy, + ) + row = { + "ts": _now_ts(), + "route_id": _norm(route_id or "", limit=120), + "prompt": _norm(prompt or "", limit=1200), + "selected_agent_mode": mode, + "route_economics": compact_economics, + "auto_agent_policy": _compact_auto_agent_policy(policy, mode), + "model_key": _norm(model_key or "", limit=120), + "route_reason": _norm(route_reason or "", limit=240), + } + rows = list(payload.get("route_usage") or []) + rows.append(row) + payload["route_usage"] = rows[-240:] + self.save_session(session_id, payload) + return { + "ok": True, + "usage": row, + "summary": self.route_usage_summary(session_id), + } + def build_context(self, session_id: str, prompt: str, *, max_memories: int = 5, max_examples: int = 2) -> Dict[str, Any]: payload = self.load_session(session_id) prompt_tokens = _tokens(prompt) @@ -231,6 +944,124 @@ def build_context(self, session_id: str, prompt: str, *, max_memories: int = 5, "raw": payload, } + def route_feedback_summary(self, session_id: str, prompt: str = "", *, max_items: int = 30) -> Dict[str, Any]: + payload = self.load_session(session_id) + rows = [row for row in list(payload.get("route_feedback") or []) if isinstance(row, dict)] + recent = rows[-max_items:] + prompt_tokens = _tokens(prompt) + relevant: List[Dict[str, Any]] = [] + for row in recent: + row_prompt = _norm(row.get("prompt") or "", limit=1200) + overlap = len(prompt_tokens & _tokens(row_prompt)) if prompt_tokens else 0 + if not prompt_tokens or overlap >= 2: + item = dict(row) + item["prompt_overlap"] = overlap + relevant.append(item) + selected_rows = relevant if relevant else recent + mode_scores: Dict[str, Dict[str, Any]] = { + mode: { + "count": 0, + "positive": 0, + "negative": 0, + "net": 0, + "quality_positive": 0, + "quality_negative": 0, + "quality_net": 0, + "depth_preference_net": 0, + "cost_pressure": 0, + "latency_pressure": 0, + } + for mode in ROUTE_MODES + } + mode_rows: Dict[str, List[Dict[str, Any]]] = {mode: [] for mode in ROUTE_MODES} + for row in selected_rows: + mode = str(row.get("selected_agent_mode") or "off") + if mode not in mode_scores: + mode = "off" + mode_rows[mode].append(row) + delta = int(row.get("score_delta") or 0) + mode_scores[mode]["count"] += 1 + mode_scores[mode]["net"] += delta + if delta > 0: + mode_scores[mode]["positive"] += 1 + elif delta < 0: + mode_scores[mode]["negative"] += 1 + axes = _feedback_axes(row) + quality_delta = axes["quality"] + if quality_delta is not None: + mode_scores[mode]["quality_net"] += int(quality_delta) + if quality_delta > 0: + mode_scores[mode]["quality_positive"] += 1 + elif quality_delta < 0: + mode_scores[mode]["quality_negative"] += 1 + mode_scores[mode]["depth_preference_net"] += int(axes["depth_preference"] or 0) + mode_scores[mode]["cost_pressure"] += int(axes["cost_pressure"] or 0) + mode_scores[mode]["latency_pressure"] += int(axes["latency_pressure"] or 0) + for mode, rows_for_mode in mode_rows.items(): + economics = _route_economics_summary(rows_for_mode) + mode_scores[mode]["economics"] = economics + mode_scores[mode]["adaptive"] = _route_adaptive_feedback_summary(rows_for_mode, economics) + economics_summary = _route_economics_summary(selected_rows) + return { + "total_feedback": len(rows), + "recent_feedback": len(recent), + "relevant_feedback": len(relevant), + "used_recent_fallback": bool(prompt_tokens and not relevant and recent), + "mode_scores": mode_scores, + "economics": economics_summary, + "adaptive": _route_adaptive_feedback_summary(selected_rows, economics_summary), + "recent": [ + { + "selected_agent_mode": _norm(row.get("selected_agent_mode") or "off", limit=48), + "rating": _norm(row.get("rating") or "", limit=16), + "score_delta": int(row.get("score_delta") or 0), + "feedback_intent": _norm(row.get("feedback_intent") or "bad_quality", limit=32), + "feedback_tags": [ + _norm(item, limit=32) for item in list(row.get("feedback_tags") or [])[:4] + ], + "feedback_axes": _feedback_axes(row), + "reason": _norm(row.get("reason") or "", limit=120), + "route_economics": ( + row.get("route_economics") + if isinstance(row.get("route_economics"), dict) + else {} + ), + } + for row in rows[-5:] + ], + } + + def route_usage_summary(self, session_id: str, *, max_items: int = 240) -> Dict[str, Any]: + payload = self.load_session(session_id) + rows = [row for row in list(payload.get("route_usage") or []) if isinstance(row, dict)] + recent = rows[-max_items:] + mode_rows: Dict[str, List[Dict[str, Any]]] = {mode: [] for mode in ROUTE_MODES} + for row in recent: + mode = str(row.get("selected_agent_mode") or "off") + if mode not in mode_rows: + mode = "off" + mode_rows[mode].append(row) + return { + "total_routes": len(rows), + "recent_routes": len(recent), + "mode_economics": { + mode: _route_economics_summary(rows_for_mode) + for mode, rows_for_mode in mode_rows.items() + }, + "economics": _route_economics_summary(recent), + "recent": [ + { + "selected_agent_mode": _norm(row.get("selected_agent_mode") or "off", limit=48), + "route_economics": ( + row.get("route_economics") + if isinstance(row.get("route_economics"), dict) + else {} + ), + } + for row in rows[-5:] + ], + } + def session_snapshot(self, session_id: str) -> Dict[str, Any]: payload = self.load_session(session_id) memories = [ @@ -251,6 +1082,9 @@ def session_snapshot(self, session_id: str) -> Dict[str, Any]: "session_id": session_id, "memory_count": len(payload.get("memories") or []), "turn_count": len(turns), + "route_feedback_count": len(payload.get("route_feedback") or []), + "route_feedback": self.route_feedback_summary(session_id), + "route_usage": self.route_usage_summary(session_id), "memories": memories, "recent_turns": recent_turns, "updated_at": payload.get("updated_at"), diff --git a/source/multimodel_runtime.py b/source/multimodel_runtime.py index bac703891..ff62d5f2e 100644 --- a/source/multimodel_runtime.py +++ b/source/multimodel_runtime.py @@ -1,23 +1,27 @@ from __future__ import annotations +import copy import gc import hashlib import io import json import logging +import math import re import threading import time +import uuid import zipfile from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple from urllib.parse import quote from urllib.request import urlopen import torch from PIL import Image +import chat_app import chat_web_app from chat_export import copy_generated_image, render_chat_transcript_image from chat_image_variant_app import ( @@ -35,6 +39,23 @@ discover_model_records, ) from multimodel_memory import ConversationMemoryStore +from route_policy_ledger import ( + OUTCOME_CONTRACT_SCHEMA_VERSION, + SUPPORT_SCHEMA_VERSION, + DecisionNotFoundError, + RoutePolicyLedger, + build_logging_support_envelope, + build_route_outcome_contracts, + hash_session_identity, +) +from route_policy_lab import POLICY_PROFILES, analyze_route_policy +from route_policy_explorer import plan_adjacent_route_study +from route_policy_protocol import ( + audit_route_study_review_bundle, + build_route_study_protocol, + build_route_study_review_bundle_from_input, +) +from route_policy_shadow_registry import RouteShadowAssignmentRegistry from multimodel_tools import ( CmdOpenTool, ToolEvent, @@ -61,11 +82,31 @@ from omni_collective_v6_model import OmniCollectiveEngineV6 from omni_collective_v7_model import OmniCollectiveEngineV7 from omni_collective_v8_model import OmniCollectiveEngineV8 -from omni_collective_v42_model import OmniCollectiveEngineV42 -from omni_collective_v41_model import OmniCollectiveEngineV41 -from omni_collective_v46_model import OmniCollectiveEngineV46 -from omni_collective_v47_model import OmniCollectiveEnginev47, OmniPredictionv47 -from omni_collective_v48_model import OmniCollectiveEnginev48, OmniPredictionv48 +try: + from omni_collective_v42_model import OmniCollectiveEngineV42 +except ImportError: + OmniCollectiveEngineV42 = None + +try: + from omni_collective_v41_model import OmniCollectiveEngineV41 +except ImportError: + OmniCollectiveEngineV41 = None + +try: + from omni_collective_v46_model import OmniCollectiveEngineV46 +except ImportError: + OmniCollectiveEngineV46 = None + +try: + from omni_collective_v47_model import OmniCollectiveEnginev47, OmniPredictionv47 +except ImportError: + OmniCollectiveEnginev47, OmniPredictionv47 = None, None + +try: + from omni_collective_v48_model import OmniCollectiveEnginev48, OmniPredictionv48 +except ImportError: + OmniCollectiveEnginev48, OmniPredictionv48 = None, None + from run import safe_load_state_dict @@ -77,6 +118,7 @@ class ChatResult: route_reason: str response: str = "" timing: Optional[Dict[str, Any]] = None + compute: Optional[Dict[str, Any]] = None image_url: str = "" output_path: str = "" prompt_used: str = "" @@ -92,6 +134,7 @@ def to_dict(self) -> Dict[str, Any]: "route_reason": self.route_reason, "response": self.response, "timing": self.timing or {}, + "compute": self.compute or {}, "image_url": self.image_url, "output_path": self.output_path, "prompt_used": self.prompt_used, @@ -111,6 +154,22 @@ def _trim_text(text: str, limit: int = 320) -> str: return cooked[:limit] +def _coerce_int_setting(value: Any, default: int, *, minimum: int = 0, maximum: int = 100_000) -> int: + try: + cooked = int(float(value)) + except (TypeError, ValueError): + cooked = int(default) + return max(int(minimum), min(int(maximum), cooked)) + + +def _record_route_model_call(settings: Dict[str, Any], count: int = 1) -> None: + """Increment the request-local model invocation counter when one is active.""" + counter = settings.get("_route_model_call_counter") + if not isinstance(counter, dict): + return + counter["count"] = max(0, int(counter.get("count") or 0)) + max(0, int(count)) + + _CHAT_DRIFT_MARKERS = ( "choose omni collective", "common benchmark score", @@ -245,6 +304,84 @@ def _parse_yes_no_section(text: str, label: str) -> Optional[bool]: return None +def _normalize_score_0_1(value: float) -> float: + cooked = float(value) + if cooked > 1.0: + cooked = cooked / (10.0 if cooked <= 10.0 else 100.0) + return max(0.0, min(1.0, cooked)) + + +def _parse_labeled_score_0_1(text: str, label: str) -> Optional[float]: + section = _extract_labeled_section(text, label) + if not section: + return None + fraction = re.search(r"([-+]?\d+(?:\.\d+)?)\s*/\s*([-+]?\d+(?:\.\d+)?)", section) + if fraction: + denominator = float(fraction.group(2)) + if denominator: + return round(_normalize_score_0_1(float(fraction.group(1)) / denominator), 4) + match = re.search(r"([-+]?\d+(?:\.\d+)?)\s*%?", section) + if not match: + return None + value = float(match.group(1)) + if "%" in match.group(0): + value /= 100.0 + return round(_normalize_score_0_1(value), 4) + + +def _first_labeled_score_0_1(text: str, labels: Sequence[str]) -> Optional[float]: + for label in labels: + value = _parse_labeled_score_0_1(text, label) + if value is not None: + return value + return None + + +def _loop_score_threshold(raw_value: Any) -> float: + if raw_value is None or raw_value == "": + return LOOP_AGENT_DEFAULT_SCORE_THRESHOLD + try: + return max(0.55, min(0.98, _normalize_score_0_1(float(raw_value)))) + except (TypeError, ValueError): + return LOOP_AGENT_DEFAULT_SCORE_THRESHOLD + + +def _loop_review_metrics(reviewer_text: str, review_complete: Optional[bool]) -> Dict[str, Any]: + explicit_score = _first_labeled_score_0_1( + reviewer_text, + ("SCORE", "VERIFIER_SCORE", "COMPLETION_SCORE", "QUALITY_SCORE"), + ) + progress = _first_labeled_score_0_1( + reviewer_text, + ("PROGRESS_SCORE", "COMPLETION_SCORE", "COMPLETENESS", "QUALITY_SCORE", "SCORE"), + ) + confidence = _first_labeled_score_0_1(reviewer_text, ("CONFIDENCE", "READY_CONFIDENCE", "CERTAINTY")) + risk = _first_labeled_score_0_1(reviewer_text, ("RISK_SCORE", "RISK", "ERROR_RISK")) + + if progress is None: + progress = 1.0 if review_complete is True else 0.35 if review_complete is False else 0.5 + if confidence is None: + confidence = 0.9 if review_complete is True else 0.45 if review_complete is False else 0.55 + if risk is None: + risk = 0.05 if review_complete is True else 0.45 if review_complete is False else 0.3 + + verifier_score = round((progress * 0.50) + (confidence * 0.30) + ((1.0 - risk) * 0.20), 4) + review_score = round(explicit_score if explicit_score is not None else verifier_score, 4) + evidence = ( + _extract_labeled_section(reviewer_text, "EVIDENCE") + or _extract_labeled_section(reviewer_text, "COMPLETION_EVIDENCE") + or _extract_labeled_section(reviewer_text, "REASON") + ) + return { + "progress_score": round(progress, 4), + "confidence_score": round(confidence, 4), + "risk_score": round(risk, 4), + "review_score": review_score, + "verifier_score": verifier_score, + "completion_evidence": _trim_text(evidence, limit=220), + } + + def _safe_upload_name(filename: str) -> str: cooked = re.sub(r"[^A-Za-z0-9._-]+", "-", str(filename or "").strip()).strip(".-") return cooked[:96] or "upload.png" @@ -254,10 +391,150 @@ def _safe_upload_name(filename: str) -> str: MODEL_STORE_CACHE_TTL_SECONDS = 300.0 LOOP_AGENT_DEFAULT_MAX_STEPS = 4 LOOP_AGENT_HARD_MAX_STEPS = 8 +LOOP_AGENT_DEFAULT_SCORE_THRESHOLD = 0.88 +MODEL_STORE_ARTIFACT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,180}\.zip$", re.IGNORECASE) +AUTO_AGENT_MODE_ORDER = ("off", "collective", "loop", "collective_loop") +AUTO_ROUTE_POLICY_ID = "auto-route-v2" +AUTO_ROUTE_POLICY_VERSION = "2.0.0" +AUTO_ROUTE_FEATURE_SCHEMA_VERSION = "route-context-v1" +AUTO_AGENT_POSITIVE_THRESHOLDS = { + "collective": 1, + "loop": 3, + "collective_loop": 4, +} +AUTO_AGENT_ADAPTIVE_MIN_WEIGHTED_COUNT = 1.2 +AUTO_AGENT_ADAPTIVE_QUALITY_FLOOR = 0.7 +AUTO_AGENT_ADAPTIVE_QUALITY_DELTA = 0.08 +AUTO_AGENT_ADAPTIVE_QUALITY_COST_DELTA = 0.08 +AUTO_AGENT_SELECTION_THRESHOLDS = { + "collective": 2, + "loop": 4, + "collective_loop": 5, +} + + +def _stamp_auto_route_policy( + policy: Dict[str, Any], + *, + selected_agent_mode: str, + action_mode: str, + logging_support: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + support_candidates = ( + logging_support.get("candidates") + if isinstance(logging_support, dict) and isinstance(logging_support.get("candidates"), list) + else [] + ) + allowed_modes = [ + str(candidate.get("action")) + for candidate in support_candidates + if isinstance(candidate, dict) and str(candidate.get("action")) in AUTO_AGENT_MODE_ORDER + ] + if not allowed_modes: + allowed_modes = [ + str(mode) + for mode in list(policy.get("allowed_agent_modes") or []) + if str(mode) in AUTO_AGENT_MODE_ORDER + ] + if selected_agent_mode not in allowed_modes: + allowed_modes.append(selected_agent_mode) + policy["policy_id"] = AUTO_ROUTE_POLICY_ID + policy["policy_version"] = AUTO_ROUTE_POLICY_VERSION + policy["feature_schema_version"] = AUTO_ROUTE_FEATURE_SCHEMA_VERSION + policy["decision_type"] = "deterministic" + policy["action_mode"] = action_mode + policy["eligible_agent_modes"] = list(allowed_modes) + policy["eligible_actions"] = list(allowed_modes) + policy["selected_agent_mode"] = selected_agent_mode + probabilities = { + mode: 1.0 if mode == selected_agent_mode else 0.0 + for mode in AUTO_AGENT_MODE_ORDER + if mode in allowed_modes + } + policy["action_probabilities"] = dict(probabilities) + policy["post_filter_action_probabilities"] = dict(probabilities) + policy["probability_stage"] = "post_filter" + policy["logging_propensity"] = 1.0 + policy["selection_thresholds"] = dict(AUTO_AGENT_SELECTION_THRESHOLDS) + policy["counterfactual_support"] = "none_deterministic_logging" + if isinstance(logging_support, dict): + policy["logging_support"] = dict(logging_support) + policy["candidate_set_hash"] = logging_support.get("candidate_set_hash") + policy["distribution_hash"] = logging_support.get("distribution_hash") + policy["decision_context"] = { + "action_mode": action_mode, + "budget_profile": str(policy.get("budget_profile") or "balanced"), + "score": int(policy.get("score") or 0), + "allowed_agent_modes": list(allowed_modes), + } + return policy + + +AUTO_AGENT_UNCERTAINTY_SIGNAL_PATTERNS = ( + ("audit", re.compile(r"\baudit(?:ing)?\b", re.IGNORECASE)), + ("verification", re.compile(r"\b(verif(?:y|ies|ied|ication)|validat(?:e|es|ed|ion))\b", re.IGNORECASE)), + ("risk", re.compile(r"\b(risk|safety|high[- ]stakes|production)\b", re.IGNORECASE)), + ("correctness", re.compile(r"\b(correctness|correct|accuracy|accurate)\b", re.IGNORECASE)), + ("edge_cases", re.compile(r"\bedge[- ]cases?\b", re.IGNORECASE)), + ("proof", re.compile(r"\b(proof|prove|formal)\b", re.IGNORECASE)), + ("review", re.compile(r"\breview\b", re.IGNORECASE)), + ("regression", re.compile(r"\bregression\b", re.IGNORECASE)), + ("evidence", re.compile(r"\b(source|cite|citation|evidence)\b", re.IGNORECASE)), +) +AUTO_AGENT_BUDGET_PROFILES = { + "fast": { + "label": "Fast", + "cost_preference": "low_latency", + "score_bias": -1, + "max_agent_mode": "collective", + }, + "balanced": { + "label": "Balanced", + "cost_preference": "balanced", + "score_bias": 0, + "max_agent_mode": "collective_loop", + }, + "deep": { + "label": "Deep", + "cost_preference": "quality", + "score_bias": 1, + "max_agent_mode": "collective_loop", + }, + "max": { + "label": "Max", + "cost_preference": "frontier", + "score_bias": 2, + "max_agent_mode": "collective_loop", + }, +} + + +def _validate_model_store_file_name(file_name: str) -> str: + cooked = str(file_name or "").strip() + if not cooked: + raise ValueError("file_name is required") + if ( + cooked in {".", ".."} + or "/" in cooked + or "\\" in cooked + or ":" in cooked + or Path(cooked).name != cooked + or not MODEL_STORE_ARTIFACT_RE.fullmatch(cooked) + ): + raise ValueError(f"Unsafe model store artifact name: {cooked!r}") + return cooked + + +def _is_safe_model_store_manifest_item(item: Dict[str, Any]) -> bool: + try: + _validate_model_store_file_name(item.get("file_name") or "") + return True + except ValueError: + return False def _hf_dataset_file_url(repo_id: str, filename: str) -> str: - return f"https://huggingface.co/datasets/{repo_id}/resolve/main/{quote(filename)}?download=true" + return f"https://huggingface.co/datasets/{repo_id}/resolve/main/{quote(filename, safe='')}?download=true" def _missing_zip_members(archive: zipfile.ZipFile, target: Path) -> List[str]: @@ -460,6 +737,12 @@ def chat(self, session_id: str, prompt: str, settings: Dict[str, Any]) -> ChatRe style_mode=str(settings.get("style_mode") or "auto"), response_temperature=float(settings.get("response_temperature") or 0.08), show_top_responses=int(settings.get("show_top_responses") or 0), + reasoning_cycles=settings.get("reasoning_cycles"), + adaptive_compute=settings.get("adaptive_compute"), + adaptive_exit_tol=settings.get("adaptive_exit_tol"), + adaptive_exit_entropy=settings.get("adaptive_exit_entropy"), + prediction_stability_patience=settings.get("prediction_stability_patience"), + prediction_stability_tol=settings.get("prediction_stability_tol"), ) return ChatResult( kind="text", @@ -468,6 +751,7 @@ def chat(self, session_id: str, prompt: str, settings: Dict[str, Any]) -> ChatRe route_reason=str(settings.get("route_reason") or ""), response=str(payload.get("response") or ""), timing=dict(payload.get("timing_ms") or {}), + compute=dict(payload.get("compute") or {}), prompt_used=effective_prompt, ) @@ -524,7 +808,11 @@ def generate_image(self, session_id: str, prompt: str, settings: Dict[str, Any]) model_key=self.record.key, model_label=self.record.label, route_reason=str(settings.get("route_reason") or ""), - timing={"total_ms": result.get("timing_ms")}, + timing={ + "total_ms": result.get("timing_ms"), + "model_calls": max(1, int(result.get("model_calls") or 1)), + "refiner_model_calls": max(0, int(result.get("refiner_model_calls") or 0)), + }, image_url=str(result.get("image_url") or ""), output_path=str(result.get("output_path") or ""), prompt_used=str(result.get("prompt_used") or effective_prompt), @@ -1179,6 +1467,8 @@ def __init__(self, record: ModelRecord, extracted_dir: Path, generated_dir: Path raise FileNotFoundError(f"Missing omnibus weights/meta for {record.label} in {extracted_dir}") self.weights_path = weights_path.resolve() self.meta_path = meta_path.resolve() + if OmniCollectiveEngineV41 is None: + raise ImportError(f"OmniCollectiveEngineV41 is not available. Please ensure omni_collective_v41_model.py exists.") self.engine = OmniCollectiveEngineV41(weights_path=self.weights_path, meta_path=self.meta_path) def status(self) -> Dict[str, Any]: @@ -1222,6 +1512,8 @@ def __init__(self, record: ModelRecord, extracted_dir: Path, generated_dir: Path raise FileNotFoundError(f"Missing omnibus weights/meta for {record.label} in {extracted_dir}") self.weights_path = weights_path.resolve() self.meta_path = meta_path.resolve() + if OmniCollectiveEngineV42 is None: + raise ImportError(f"OmniCollectiveEngineV42 is not available. Please ensure omni_collective_v42_model.py exists.") self.engine = OmniCollectiveEngineV42(weights_path=self.weights_path, meta_path=self.meta_path) def status(self) -> Dict[str, Any]: @@ -1265,6 +1557,8 @@ def __init__(self, record: ModelRecord, extracted_dir: Path, generated_dir: Path raise FileNotFoundError(f"Missing omnibus weights/meta for {record.label} in {extracted_dir}") self.weights_path = weights_path.resolve() self.meta_path = meta_path.resolve() + if OmniCollectiveEngineV46 is None: + raise ImportError(f"OmniCollectiveEngineV46 is not available. Please ensure omni_collective_v46_model.py exists.") self.engine = OmniCollectiveEngineV46(weights_path=self.weights_path, meta_path=self.meta_path) def status(self) -> Dict[str, Any]: @@ -1324,6 +1618,8 @@ def __init__(self, record: ModelRecord, extracted_dir: Path, generated_dir: Path raise FileNotFoundError(f"Missing omnibus weights/meta for {record.label} in {extracted_dir}") self.weights_path = weights_path.resolve() self.meta_path = meta_path.resolve() + if OmniCollectiveEnginev47 is None: + raise ImportError(f"OmniCollectiveEnginev47 is not available. Please ensure omni_collective_v47_model.py exists.") self.engine = OmniCollectiveEnginev47(weights_path=self.weights_path, meta_path=self.meta_path) def status(self) -> Dict[str, Any]: @@ -1383,6 +1679,8 @@ def __init__(self, record: ModelRecord, extracted_dir: Path, generated_dir: Path raise FileNotFoundError(f"Missing omnibus weights/meta for {record.label} in {extracted_dir}") self.weights_path = weights_path.resolve() self.meta_path = meta_path.resolve() + if OmniCollectiveEnginev48 is None: + raise ImportError(f"OmniCollectiveEnginev48 is not available. Please ensure omni_collective_v48_model.py exists.") self.engine = OmniCollectiveEnginev48(weights_path=self.weights_path, meta_path=self.meta_path) def status(self) -> Dict[str, Any]: @@ -1477,7 +1775,15 @@ def __init__( self._model_store_manifest_cache: Optional[Dict[str, Any]] = None self._model_store_manifest_ts = 0.0 self._model_store_jobs: Dict[str, Dict[str, Any]] = {} - self.memory_store = ConversationMemoryStore(self.extraction_root.parent / "memory") + memory_dir = self.extraction_root.parent / "memory" + self.memory_store = ConversationMemoryStore(memory_dir) + self.route_policy_ledger = RoutePolicyLedger(memory_dir / "route-policy-ledger.sqlite3") + self.route_shadow_registry_path = memory_dir / "route-policy-shadow-registry.sqlite3" + self._route_shadow_registry_cache_signature: Optional[ + Tuple[Tuple[Any, ...], ...] + ] = None + self._route_shadow_registry_cache_snapshot: Optional[Dict[str, Any]] = None + self._route_execution = threading.local() self.web_search = WebSearchTool() self.cmd_open = CmdOpenTool() @@ -1701,8 +2007,9 @@ def model_store_catalog(self, force_refresh: bool = False) -> Dict[str, Any]: for item in manifest.get("models") or []: if not isinstance(item, dict): continue - file_name = str(item.get("file_name") or "").strip() - if not file_name: + try: + file_name = _validate_model_store_file_name(item.get("file_name") or "") + except ValueError: continue details = describe_model_artifact_name(file_name) local_path = self.models_dir / file_name @@ -1749,6 +2056,7 @@ def _set_model_store_job(self, job_id: str, **updates: Any) -> None: self._model_store_jobs[job_id] = payload def _install_model_store_worker(self, job_id: str, file_name: str, expected_size: int) -> None: + file_name = _validate_model_store_file_name(file_name) target = self.models_dir / file_name temp_target = self.models_dir / f"{file_name}.{job_id}.part" try: @@ -1794,15 +2102,14 @@ def _install_model_store_worker(self, job_id: str, file_name: str, expected_size ) def install_model_store_artifact(self, file_name: str) -> Dict[str, Any]: - cooked = str(file_name or "").strip() - if not cooked: - raise ValueError("file_name is required") + cooked = _validate_model_store_file_name(file_name) with self._lock: manifest = self._fetch_model_store_manifest_locked(force_refresh=False) manifest_rows = { - str(item.get("file_name") or "").strip(): item + _validate_model_store_file_name(item.get("file_name") or ""): item for item in (manifest.get("models") or []) if isinstance(item, dict) + and _is_safe_model_store_manifest_item(item) } if cooked not in manifest_rows: raise FileNotFoundError(f"{cooked} is not present in {self.model_store_repo_id}") @@ -1905,7 +2212,12 @@ def _seed_auto_tool_events(self, prompt: str, settings: Dict[str, Any]) -> List[ events: List[ToolEvent] = [] if bool(settings.get("web_search_enabled", False)) and should_offer_web_search(prompt): try: - events.append(self.web_search.search(prompt, max_results=int(settings.get("web_search_results") or 5))) + events.append( + self.web_search.search( + prompt, + max_results=_coerce_int_setting(settings.get("web_search_results"), 5, minimum=1, maximum=20), + ) + ) except Exception: pass if bool(settings.get("cmd_open_enabled", True)) and should_offer_open_cmd(prompt): @@ -1926,10 +2238,13 @@ def _run_web_query_cached( return None if key in tool_cache: return tool_cache[key] - if len(tool_cache) >= int(settings.get("web_search_budget") or 3): + if len(tool_cache) >= _coerce_int_setting(settings.get("web_search_budget"), 3, minimum=0, maximum=100): return None try: - event = self.web_search.search(query, max_results=int(settings.get("web_search_results") or 5)) + event = self.web_search.search( + query, + max_results=_coerce_int_setting(settings.get("web_search_results"), 5, minimum=1, maximum=20), + ) except Exception: return None tool_cache[key] = event @@ -1965,6 +2280,7 @@ def _run_text_model( resolved_record, backend = self.ensure_backend(record.key) local_events: List[ToolEvent] = [] run_settings = dict(settings) + run_settings.pop("_route_model_call_counter", None) effective_route_reason = route_reason if resolved_record.key != record.key: effective_route_reason = ( @@ -1981,6 +2297,7 @@ def _run_text_model( "TOOL:open_cmd: \n" "Use a tool line only when it is explicitly needed, otherwise answer normally." ) + _record_route_model_call(settings) result = backend.chat(session_id, prompt, run_settings) raw_response = str(result.response or "") requests = parse_tool_requests(raw_response) if allow_tool_calls else [] @@ -1997,11 +2314,13 @@ def _run_text_model( local_events.append(tool_event) if local_events: follow_settings = dict(settings) + follow_settings.pop("_route_model_call_counter", None) follow_settings["route_reason"] = route_reason follow_settings["memory_context"] = str(settings.get("memory_context") or "") follow_settings["consultation_context"] = str(settings.get("consultation_context") or "") follow_settings["tool_context"] = format_tool_results(list(tool_cache.values())) follow_settings["tool_instruction"] = "Tool results are already available below. Use them and answer directly." + _record_route_model_call(settings) follow = backend.chat(session_id, prompt, follow_settings) follow.response = strip_tool_calls(follow.response) result = follow @@ -2046,178 +2365,1760 @@ def _build_synthesis_prompt(self, prompt: str, action_mode: str) -> str: def _normalized_agent_mode(self, raw_mode: Any) -> str: cooked = str(raw_mode or "off").strip().lower() aliases = { + "adaptive": "auto", + "smart": "auto", "collective_all": "collective", "panel": "collective", "loop_agent": "loop", "collective_loop_agent": "collective_loop", } normalized = aliases.get(cooked, cooked) - if normalized in {"off", "collective", "loop", "collective_loop"}: + if normalized in {"off", "auto", "collective", "loop", "collective_loop"}: return normalized return "off" - def _format_loop_history(self, steps: Sequence[Dict[str, str]]) -> str: - if not steps: - return "No prior loop work yet." - lines: List[str] = [] - for row in steps: - lines.append(f"Step {row.get('step')}:") - if row.get("goal"): - lines.append(f"Goal: {row['goal']}") - if row.get("worker_excerpt"): - lines.append(f"Worker: {row['worker_excerpt']}") - if row.get("review_note"): - lines.append(f"Reviewer: {row['review_note']}") - if row.get("next_step"): - lines.append(f"Next: {row['next_step']}") - return "\n".join(lines[-20:]) - - def _build_loop_planner_prompt(self, prompt: str, history: str, step_index: int, max_steps: int) -> str: - return ( - "You are the planner sub-agent inside a loop agent.\n" - "Break the user's task into the single highest-value next action.\n" - "Respond using exactly these headings:\n" - "DONE: yes or no\n" - "STEP_GOAL: one concise sentence\n" - "SUCCESS_SIGNAL: how to know the task is complete\n" - "WORKING_NOTES: short reasoning\n\n" - f"Original task:\n{prompt}\n\n" - f"Loop step: {step_index} of {max_steps}\n\n" - f"Work log:\n{history}" - ) + def _normalized_auto_budget_profile(self, raw_profile: Any) -> str: + cooked = str(raw_profile or "balanced").strip().lower() + aliases = { + "default": "balanced", + "normal": "balanced", + "standard": "balanced", + "cheap": "fast", + "economy": "fast", + "latency": "fast", + "quality": "deep", + "thorough": "deep", + "frontier": "max", + "maximum": "max", + "unbounded": "max", + } + normalized = aliases.get(cooked, cooked) + if normalized in AUTO_AGENT_BUDGET_PROFILES: + return normalized + return "balanced" - def _build_loop_worker_prompt( + def _allowed_auto_agent_modes( self, - prompt: str, *, - history: str, - step_goal: str, - step_index: int, - max_steps: int, - ) -> str: - return ( - "You are the worker sub-agent inside a loop agent.\n" - "Execute the next best action for the task. Prefer doing the work over talking about the work.\n" - "Respond using exactly these headings:\n" - "DONE: yes or no\n" - "OUTPUT: the actual deliverable or best current progress\n" - "NEXT_FOCUS: what should happen next if not done\n\n" - f"Original task:\n{prompt}\n\n" - f"Loop step: {step_index} of {max_steps}\n" - f"Current goal:\n{step_goal}\n\n" - f"Work log:\n{history}" + action_mode: str, + allow_collective: bool, + collective_available: bool, + allow_loop: bool, + ) -> List[str]: + modes = ["off"] + if allow_collective and collective_available: + modes.append("collective") + if action_mode != "image" and allow_loop: + modes.append("loop") + if action_mode != "image" and allow_loop and allow_collective and collective_available: + modes.append("collective_loop") + return modes + + def _budget_allowed_auto_agent_modes(self, allowed_modes: Sequence[str], max_agent_mode: str) -> List[str]: + if max_agent_mode not in AUTO_AGENT_MODE_ORDER: + max_agent_mode = "collective_loop" + max_idx = AUTO_AGENT_MODE_ORDER.index(max_agent_mode) + allowed = set(allowed_modes) + return [mode for mode in AUTO_AGENT_MODE_ORDER[: max_idx + 1] if mode in allowed] + + def _effective_agent_mode_for_action(self, agent_mode: str, action_mode: str) -> str: + if action_mode != "image": + return agent_mode + if agent_mode == "loop": + return "off" + if agent_mode == "collective_loop": + return "collective" + return agent_mode + + def _route_latency_tier(self, cost_units: float) -> str: + if cost_units <= 1.5: + return "low" + if cost_units <= 4.0: + return "moderate" + if cost_units <= 10.0: + return "high" + return "frontier" + + def _estimate_route_economics( + self, + *, + selected_agent_mode: str, + action_mode: str, + settings: Dict[str, Any], + auto_agent_policy: Optional[Dict[str, Any]], + ) -> Dict[str, Any]: + loop_budget = _coerce_int_setting( + settings.get("loop_max_steps") or settings.get("loop_budget"), + LOOP_AGENT_DEFAULT_MAX_STEPS, + minimum=2, + maximum=LOOP_AGENT_HARD_MAX_STEPS, + ) + collective_count = int((auto_agent_policy or {}).get("collective_model_count") or 0) + if selected_agent_mode in {"collective", "collective_loop"}: + collective_count = max(2, collective_count) + tool_budget = ( + _coerce_int_setting(settings.get("web_search_budget"), 0, minimum=0, maximum=100) + if bool(settings.get("web_search_enabled", False)) + else 0 ) - def _build_loop_review_prompt( + if selected_agent_mode == "collective": + estimated_model_calls = max(2, collective_count + 1) + planned_steps = 0 + elif selected_agent_mode == "loop": + estimated_model_calls = loop_budget * 3 + planned_steps = loop_budget + elif selected_agent_mode == "collective_loop": + estimated_model_calls = loop_budget * (3 + max(0, collective_count - 1)) + planned_steps = loop_budget + else: + estimated_model_calls = 1 + planned_steps = 0 + + if action_mode == "image": + estimated_model_calls += 1 + + estimated_cost_units = round(float(estimated_model_calls) + (float(tool_budget) * 0.25), 2) + budget_profile = str((auto_agent_policy or {}).get("budget_profile") or "manual") + budget_policy = (auto_agent_policy or {}).get("budget_policy") + cost_preference = "" + if isinstance(budget_policy, dict): + cost_preference = _trim_text(budget_policy.get("cost_preference") or "", limit=80) + return { + "selected_agent_mode": selected_agent_mode, + "action_mode": action_mode, + "budget_profile": budget_profile, + "cost_preference": cost_preference, + "estimated_model_calls": estimated_model_calls, + "estimated_tool_calls": tool_budget, + "estimated_cost_units": estimated_cost_units, + "latency_tier": self._route_latency_tier(estimated_cost_units), + "planned_loop_steps": planned_steps, + "collective_model_count": collective_count, + "requested_reasoning_cycles": settings.get("reasoning_cycles"), + "adaptive_compute": bool(settings.get("adaptive_compute")) if settings.get("adaptive_compute") is not None else None, + } + + def _finalize_route_economics( self, - prompt: str, *, - history: str, - latest_output: str, - success_signal: str, - step_index: int, - max_steps: int, - ) -> str: - return ( - "You are the reviewer sub-agent inside a loop agent.\n" - "Judge whether the user's task is complete enough to stop.\n" - "Respond using exactly these headings:\n" - "COMPLETE: yes or no\n" - "FINAL_RESPONSE: the reply that should be shown to the user right now\n" - "REASON: why you marked it complete or incomplete\n" - "NEXT_STEP: the next action if it is not complete\n\n" - f"Original task:\n{prompt}\n\n" - f"Loop step: {step_index} of {max_steps}\n" - f"Success signal:\n{success_signal or 'Task is complete when the user request has been fully addressed.'}\n\n" - f"Work log:\n{history}\n\n" - f"Latest worker output:\n{latest_output}" - ) + estimate: Dict[str, Any], + trace: Dict[str, Any], + elapsed_ms: float, + actual_model_calls: Optional[int] = None, + ) -> Dict[str, Any]: + loop_steps = list(trace.get("loop_steps") or []) + consultation_rows = list(trace.get("consultation_rows") or []) + consulted_models = list(trace.get("consulted_models") or []) + tool_events = list(trace.get("tool_events") or []) + compute = trace.get("compute") if isinstance(trace.get("compute"), dict) else {} + mode = str(estimate.get("selected_agent_mode") or trace.get("resolved_agent_mode") or trace.get("agent_mode") or "off") + + if mode == "collective_loop": + inferred_model_calls = max(1, len(loop_steps) * 3 + len(consultation_rows)) + elif mode == "loop": + inferred_model_calls = max(1, len(loop_steps) * 3) + elif mode == "collective": + inferred_model_calls = max(1, 1 + max(len(consultation_rows), len(consulted_models))) + else: + inferred_model_calls = 1 + tracked_model_calls = int(actual_model_calls or 0) + resolved_model_calls = tracked_model_calls if tracked_model_calls > 0 else inferred_model_calls + actual_cost_units = round(float(resolved_model_calls) + (float(len(tool_events)) * 0.25), 2) + return { + "estimate": estimate, + "actual": { + "elapsed_ms": round(float(elapsed_ms), 2), + "model_calls": int(resolved_model_calls), + "tool_calls": len(tool_events), + "loop_steps": len(loop_steps), + "consultation_count": max(len(consultation_rows), len(consulted_models)), + "cost_units": actual_cost_units, + "latency_tier": self._route_latency_tier(actual_cost_units), + "compute_applied": bool(compute.get("applied")), + "reasoning_cycles_used": compute.get("cycles_used"), + "compute_exit_reason": compute.get("exit_reason"), + "prediction_confidence_delta": compute.get("prediction_confidence_delta"), + }, + } - def _run_loop_agent_text( + def _auto_session_budget_limit(self, settings: Dict[str, Any]) -> Optional[float]: + for key in ("auto_session_budget_units", "session_route_budget_units", "route_budget_units"): + raw = settings.get(key) + if raw in (None, ""): + continue + try: + value = float(raw) + except (TypeError, ValueError): + continue + if value > 0.0: + return round(min(value, 100_000.0), 3) + return None + + def _auto_session_budget_target_routes(self, settings: Dict[str, Any]) -> Optional[int]: + for key in ("auto_session_budget_target_routes", "session_route_budget_target_routes", "route_budget_target_routes"): + raw = settings.get(key) + if raw in (None, ""): + continue + try: + value = int(float(raw)) + except (TypeError, ValueError): + continue + if value > 1: + return min(value, 10_000) + return None + + def _session_route_budget_snapshot(self, session_id: str, budget_limit: Optional[float]) -> Optional[Dict[str, Any]]: + if budget_limit is None: + return None + usage = self.memory_store.route_usage_summary(session_id) + economics = usage.get("economics") if isinstance(usage.get("economics"), dict) else {} + used = float(economics.get("total_cost_units") or 0.0) + remaining = round(max(0.0, float(budget_limit) - used), 3) + return { + "limit_cost_units": round(float(budget_limit), 3), + "used_cost_units": round(used, 3), + "remaining_cost_units": remaining, + "route_count": int(usage.get("total_routes") or 0), + "recent_route_count": int(usage.get("recent_routes") or 0), + } + + def _apply_auto_session_budget( self, *, - session_id: str, - prompt: str, - chosen_record: ModelRecord, - settings: Dict[str, Any], - route_reason: str, + selected: str, action_mode: str, - memory_bundle: Dict[str, Any], - collective_mode: bool, - ) -> ChatResult: - controller_record = chosen_record if chosen_record.supports_chat else self._default_text_record() - max_steps = int(settings.get("loop_max_steps") or settings.get("loop_budget") or LOOP_AGENT_DEFAULT_MAX_STEPS) - max_steps = max(2, min(LOOP_AGENT_HARD_MAX_STEPS, max_steps)) - tool_cache = { - _trim_text(event.query, limit=220).lower(): event - for event in self._seed_auto_tool_events(prompt, settings) - } - tool_event_rows: List[Dict[str, Any]] = [event.to_dict() for event in list(tool_cache.values())] - consult_rows: List[Dict[str, str]] = [] - consulted_labels: List[str] = [] - skipped_models: List[Dict[str, str]] = [] - loop_steps: List[Dict[str, str]] = [] - final_response = "" - completion_reason = "" - completed = False + settings: Dict[str, Any], + auto_agent_policy: Dict[str, Any], + budget_snapshot: Optional[Dict[str, Any]], + ) -> tuple[str, Dict[str, Any], Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + estimate = self._estimate_route_economics( + selected_agent_mode=selected, + action_mode=action_mode, + settings=settings, + auto_agent_policy=auto_agent_policy, + ) + if not budget_snapshot: + return selected, estimate, None, None + + remaining = float(budget_snapshot.get("remaining_cost_units") or 0.0) + target_routes = self._auto_session_budget_target_routes(settings) + route_count = int(budget_snapshot.get("route_count") or 0) + target_remaining_routes: Optional[int] = None + pacing_cap: Optional[float] = None + target_pacing_active = False + if target_routes is not None: + target_remaining_routes = max(1, int(target_routes) - route_count) + pacing_cap = round(max(0.0, remaining) / float(target_remaining_routes), 3) + target_pacing_active = target_remaining_routes > 1 + allowed_modes = list(auto_agent_policy.get("allowed_agent_modes") or AUTO_AGENT_MODE_ORDER) + original = selected + original_estimated_cost = float(estimate.get("estimated_cost_units") or 0.0) + effective_cap = remaining if not target_pacing_active else min(remaining, float(pacing_cap or 0.0)) + estimates = [ + { + "selected_agent_mode": selected, + "estimated_cost_units": estimate.get("estimated_cost_units"), + "cap_cost_units": round(effective_cap, 3), + } + ] - for step_index in range(1, max_steps + 1): - history = self._format_loop_history(loop_steps) - planner_result, planner_events = self._run_text_model( - controller_record, - session_id=self._session_scope(session_id, controller_record.key, f"loop-plan-{step_index}"), - prompt=self._build_loop_planner_prompt(prompt, history, step_index, max_steps), - settings={ - **settings, - "memory_context": memory_bundle.get("context_block") or "", - }, - route_reason=f"{route_reason} Loop agent planner step {step_index} via {controller_record.label}.", - tool_cache=tool_cache, - allow_tool_calls=True, + while selected != "off" and float(estimate.get("estimated_cost_units") or 0.0) > effective_cap: + next_selected = self._neighbor_auto_agent_mode(selected, allowed_modes, -1) + if next_selected == selected: + break + selected = next_selected + estimate = self._estimate_route_economics( + selected_agent_mode=selected, + action_mode=action_mode, + settings=settings, + auto_agent_policy=auto_agent_policy, ) - tool_event_rows.extend(event.to_dict() for event in planner_events) - step_goal = ( - _extract_labeled_section(planner_result.response, "STEP_GOAL") - or _extract_labeled_section(planner_result.response, "WORKING_NOTES") - or _trim_text(planner_result.response, limit=220) - or "Advance the task toward a complete solution." + estimates.append( + { + "selected_agent_mode": selected, + "estimated_cost_units": estimate.get("estimated_cost_units"), + "cap_cost_units": round(effective_cap, 3), + } ) - success_signal = _extract_labeled_section(planner_result.response, "SUCCESS_SIGNAL") - if collective_mode: - worker_result = self._run_agent_text( - session_id=self._session_scope(session_id, controller_record.key, f"loop-work-{step_index}"), - prompt=self._build_loop_worker_prompt( - prompt, - history=history, - step_goal=step_goal, - step_index=step_index, - max_steps=max_steps, - ), - chosen_record=chosen_record, - settings=settings, - route_reason=f"{route_reason} Loop agent worker step {step_index} used the collective panel.", + estimated_cost = float(estimate.get("estimated_cost_units") or 0.0) + budget_state = dict(budget_snapshot) + if target_routes is not None: + budget_state.update( + { + "target_route_count": int(target_routes), + "target_remaining_routes": int(target_remaining_routes or 1), + "pacing_cap_cost_units": round(float(pacing_cap or 0.0), 3), + "effective_cap_cost_units": round(float(effective_cap), 3), + "pacing_cap_applied": bool(target_pacing_active), + } + ) + budget_state.update( + { + "selected_agent_mode": selected, + "estimated_cost_units": round(estimated_cost, 3), + "would_exceed_remaining": bool(estimated_cost > remaining), + "would_exceed_pacing_cap": bool( + target_pacing_active and estimated_cost > float(pacing_cap or 0.0) + ), + } + ) + if selected == original == "off" and estimated_cost > remaining: + adjustment = { + "direction": "floor", + "from": "off", + "to": "off", + "reason": "session_route_budget_exhausted_single_pass_floor", + "remaining_cost_units": round(remaining, 3), + "limit_cost_units": budget_state["limit_cost_units"], + "used_cost_units": budget_state["used_cost_units"], + "estimated_cost_units": round(estimated_cost, 3), + "candidate_estimates": estimates, + } + return selected, estimate, adjustment, budget_state + if selected == original: + return selected, estimate, None, budget_state + + target_pacing_triggered = ( + target_pacing_active + and original_estimated_cost > float(pacing_cap or 0.0) + and estimated_cost <= remaining + ) + adjustment = { + "direction": "downgrade", + "from": original, + "to": selected, + "reason": ( + "session_route_budget_target_pacing" + if target_pacing_triggered + else "session_route_budget_would_exceed_remaining" + ), + "remaining_cost_units": round(remaining, 3), + "limit_cost_units": budget_state["limit_cost_units"], + "used_cost_units": budget_state["used_cost_units"], + "estimated_cost_units": round(estimated_cost, 3), + "candidate_estimates": estimates, + } + if target_routes is not None: + adjustment.update( + { + "target_route_count": int(target_routes), + "target_remaining_routes": int(target_remaining_routes or 1), + "pacing_cap_cost_units": round(float(pacing_cap or 0.0), 3), + "effective_cap_cost_units": round(float(effective_cap), 3), + } + ) + return selected, estimate, adjustment, budget_state + + def _build_post_filter_logging_support( + self, + *, + selected: str, + action_mode: str, + settings: Dict[str, Any], + auto_agent_policy: Dict[str, Any], + selected_estimate: Dict[str, Any], + ) -> Dict[str, Any]: + """Freeze the final feasible action set after capability and budget filters.""" + + allowed = { + str(mode) + for mode in list(auto_agent_policy.get("allowed_agent_modes") or AUTO_AGENT_MODE_ORDER) + if str(mode) in AUTO_AGENT_MODE_ORDER + } + allowed.add(selected) + budget_state = ( + auto_agent_policy.get("session_budget") + if isinstance(auto_agent_policy.get("session_budget"), dict) + else {} + ) + cap_raw = budget_state.get("effective_cap_cost_units") + if cap_raw is None: + cap_raw = budget_state.get("remaining_cost_units") + try: + effective_cap = float(cap_raw) if cap_raw is not None else None + except (TypeError, ValueError, OverflowError): + effective_cap = None + if effective_cap is not None and (not math.isfinite(effective_cap) or effective_cap < 0.0): + effective_cap = None + + candidates: List[Dict[str, Any]] = [] + exclusions: List[Dict[str, Any]] = [] + for mode in AUTO_AGENT_MODE_ORDER: + reasons: List[str] = [] + if mode not in allowed: + reasons.append( + "action_mode_unsupported" + if action_mode == "image" and mode in {"loop", "collective_loop"} + else "capability_or_policy_filter" + ) + estimate = ( + dict(selected_estimate) + if mode == selected + else self._estimate_route_economics( + selected_agent_mode=mode, action_mode=action_mode, - memory_bundle=memory_bundle, + settings=settings, + auto_agent_policy=auto_agent_policy, ) - consult_rows.extend(list((worker_result.agent_trace or {}).get("consultation_rows") or [])) - consulted_labels.extend(list((worker_result.agent_trace or {}).get("consulted_models") or [])) - skipped_models.extend(list((worker_result.agent_trace or {}).get("skipped_models") or [])) - tool_event_rows.extend(list((worker_result.agent_trace or {}).get("tool_events") or [])) + ) + estimated_cost = float(estimate.get("estimated_cost_units") or 0.0) + if ( + not reasons + and effective_cap is not None + and mode != "off" + and estimated_cost > effective_cap + ): + reasons.append("session_budget_post_filter") + + if mode == selected: + if reasons: + raise RuntimeError("selected route was excluded from its post-filter logging support") + candidates.append( + { + "action": mode, + "selected": True, + "estimated_cost_units": round(estimated_cost, 3), + "estimated_model_calls": int(estimate.get("estimated_model_calls") or 1), + "planned_loop_steps": int(estimate.get("planned_loop_steps") or 0), + "latency_tier": str(estimate.get("latency_tier") or "unknown"), + } + ) + elif reasons: + exclusions.append({"action": mode, "reasons": reasons}) else: - worker_result, worker_events = self._run_text_model( - controller_record, - session_id=self._session_scope(session_id, controller_record.key, f"loop-work-{step_index}"), - prompt=self._build_loop_worker_prompt( - prompt, - history=history, - step_goal=step_goal, - step_index=step_index, - max_steps=max_steps, - ), - settings={ - **settings, - "memory_context": memory_bundle.get("context_block") or "", + candidates.append( + { + "action": mode, + "selected": False, + "estimated_cost_units": round(estimated_cost, 3), + "estimated_model_calls": int(estimate.get("estimated_model_calls") or 1), + "planned_loop_steps": int(estimate.get("planned_loop_steps") or 0), + "latency_tier": str(estimate.get("latency_tier") or "unknown"), + } + ) + + candidates.sort(key=lambda row: AUTO_AGENT_MODE_ORDER.index(str(row["action"]))) + exclusions.sort(key=lambda row: AUTO_AGENT_MODE_ORDER.index(str(row["action"]))) + eligible = [str(row["action"]) for row in candidates] + probabilities = {mode: 1.0 if mode == selected else 0.0 for mode in eligible} + raw_support = { + "schema_version": SUPPORT_SCHEMA_VERSION, + "decision_type": "deterministic", + "probability_stage": "post_filter", + "sampler": { + "name": "threshold_budget_argmax", + "version": "1", + "exploration_rate": 0.0, + "assignment_unit": "route", + "assignment_commitment": None, + }, + "candidates": candidates, + "exclusions": exclusions, + } + return build_logging_support_envelope( + raw_support, + eligible_modes=eligible, + action_probabilities=probabilities, + chosen_mode=selected, + ) + + def _preview_route_alternatives( + self, + *, + selected: str, + action_mode: str, + settings: Dict[str, Any], + auto_agent_policy: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + if auto_agent_policy is not None: + allowed_agent_modes = auto_agent_policy.get("allowed_agent_modes") + raw_candidate_modes = allowed_agent_modes if allowed_agent_modes is not None else AUTO_AGENT_MODE_ORDER + candidate_modes = [ + str(mode) + for mode in list(raw_candidate_modes) + if str(mode) in AUTO_AGENT_MODE_ORDER + ] + else: + candidate_modes = [selected] if selected in AUTO_AGENT_MODE_ORDER else ["off"] + if selected in AUTO_AGENT_MODE_ORDER and selected not in candidate_modes: + candidate_modes.append(selected) + + rows: List[Dict[str, Any]] = [] + seen: set[str] = set() + for mode in sorted(set(candidate_modes), key=lambda item: AUTO_AGENT_MODE_ORDER.index(item)): + if mode in seen: + continue + seen.add(mode) + estimate = self._estimate_route_economics( + selected_agent_mode=mode, + action_mode=action_mode, + settings=settings, + auto_agent_policy=auto_agent_policy, + ) + rows.append( + { + "selected_agent_mode": mode, + "is_selected": mode == selected, + "estimate": estimate, + "estimated_cost_units": estimate.get("estimated_cost_units"), + "estimated_model_calls": estimate.get("estimated_model_calls"), + "planned_loop_steps": estimate.get("planned_loop_steps"), + "latency_tier": estimate.get("latency_tier"), + } + ) + return rows + + def _estimated_route_quality( + self, + *, + mode: str, + action_mode: str, + estimated_cost_units: Optional[float] = None, + auto_agent_policy: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + feedback_summary = ( + auto_agent_policy.get("feedback_summary") + if isinstance(auto_agent_policy, dict) and isinstance(auto_agent_policy.get("feedback_summary"), dict) + else {} + ) + mode_scores = feedback_summary.get("mode_scores") if isinstance(feedback_summary.get("mode_scores"), dict) else {} + stats = mode_scores.get(mode) if isinstance(mode_scores.get(mode), dict) else {} + evidence = self._adaptive_route_evidence(stats) + if evidence is not None: + return { + "estimated_quality_score": round(float(evidence["quality_score"]), 3), + "estimated_quality_cost_score": round(float(evidence["quality_cost_score"]), 3), + "estimated_quality_lower_bound": evidence.get("quality_lower_bound"), + "estimated_quality_upper_bound": evidence.get("quality_upper_bound"), + "risk_adjusted_quality_cost_score": evidence.get("risk_adjusted_quality_cost_score"), + "effective_sample_size": evidence.get("effective_sample_size"), + "confidence_level": evidence.get("confidence_level"), + "confidence_status": evidence.get("confidence_status"), + "quality_source": "adaptive_feedback", + "quality_cost_source": "adaptive_feedback", + "quality_evidence_status": "adaptive_complete", + "quality_evidence": evidence, + } + quality_evidence_status = self._adaptive_route_evidence_status(stats) + + score = int((auto_agent_policy or {}).get("score") or 0) + if mode == "off": + quality = max(0.25, 0.62 - (0.055 * float(score))) + else: + thresholds = {"collective": 2, "loop": 4, "collective_loop": 5} + bases = {"collective": 0.62, "loop": 0.72, "collective_loop": 0.88} + threshold = int(thresholds.get(mode, 99)) + base = float(bases.get(mode, 0.5)) + if action_mode == "image" and mode == "collective": + threshold = 3 + if score >= threshold: + quality = base + (0.04 * float(min(score - threshold, 4))) + else: + quality = base - (0.06 * float(threshold - score)) + estimated_quality_score = round(max(0.0, min(0.99, quality)), 3) + try: + cost = float(estimated_cost_units or 0.0) + except (TypeError, ValueError): + cost = 0.0 + if not math.isfinite(cost): + cost = 0.0 + cost = max(0.0, cost) + return { + "estimated_quality_score": estimated_quality_score, + "estimated_quality_cost_score": round(estimated_quality_score / (1.0 + (cost / 10.0)), 3), + "estimated_quality_lower_bound": None, + "estimated_quality_upper_bound": None, + "risk_adjusted_quality_cost_score": None, + "effective_sample_size": 0.0, + "confidence_level": None, + "confidence_status": "heuristic_prior", + "quality_source": "heuristic_policy", + "quality_cost_source": "heuristic_cost_adjusted", + "quality_evidence_status": quality_evidence_status, + "quality_evidence": None, + } + + def _route_frontier_budget_state(self, auto_agent_policy: Optional[Dict[str, Any]]) -> Dict[str, Any]: + session_budget = ( + auto_agent_policy.get("session_budget") + if isinstance(auto_agent_policy, dict) and isinstance(auto_agent_policy.get("session_budget"), dict) + else None + ) + if not session_budget: + return { + "remaining_cost_units": None, + "pacing_cap_cost_units": None, + "effective_cap_cost_units": None, + "pacing_cap_applied": False, + "limit_cost_units": None, + "used_cost_units": None, + "route_count": None, + "target_route_count": None, + "target_remaining_routes": None, + } + + def optional_float(key: str) -> Optional[float]: + raw = session_budget.get(key) + if raw in (None, ""): + return None + try: + return round(max(0.0, float(raw)), 3) + except (TypeError, ValueError): + return None + + def optional_int(key: str) -> Optional[int]: + raw = session_budget.get(key) + if raw in (None, ""): + return None + try: + return int(float(raw)) + except (TypeError, ValueError): + return None + + remaining = optional_float("remaining_cost_units") + pacing_cap_applied = bool(session_budget.get("pacing_cap_applied")) + pacing_cap = optional_float("pacing_cap_cost_units") if pacing_cap_applied else None + effective_cap = optional_float("effective_cap_cost_units") + if effective_cap is None: + if remaining is not None and pacing_cap is not None: + effective_cap = min(remaining, pacing_cap) + elif remaining is not None: + effective_cap = remaining + elif pacing_cap is not None: + effective_cap = pacing_cap + return { + "remaining_cost_units": remaining, + "pacing_cap_cost_units": pacing_cap, + "effective_cap_cost_units": effective_cap, + "pacing_cap_applied": pacing_cap is not None, + "limit_cost_units": optional_float("limit_cost_units"), + "used_cost_units": optional_float("used_cost_units"), + "route_count": optional_int("route_count"), + "target_route_count": optional_int("target_route_count"), + "target_remaining_routes": optional_int("target_remaining_routes"), + } + + def _annotate_route_frontier( + self, + *, + alternatives: List[Dict[str, Any]], + selected: str, + action_mode: str, + auto_agent_policy: Optional[Dict[str, Any]], + ) -> Dict[str, Any]: + if not alternatives: + return { + "recommended_agent_mode": selected, + "recommended_reason": "no_alternatives_available", + "recommended_budget_blocker": None, + "ranked_modes": [], + "pareto_modes": [], + "budget_feasible_pareto_modes": [], + "budget_fit_count": 0, + "budget_blockers": {"remaining_budget": 0, "pacing_cap": 0, "effective_cap": 0, "none": 0}, + "budget_cap_cost_units": None, + "remaining_cost_units": None, + "pacing_cap_cost_units": None, + "effective_cap_cost_units": None, + "pacing_cap_applied": False, + "recommended_estimated_quality_cost_score": None, + "selected_estimated_quality_cost_score": None, + "minimum_route_floor": False, + "selected_budget_fit": None, + "selected_budget_blocker": None, + "selected_fits_remaining_budget": None, + "selected_fits_pacing_cap": None, + "selected_frontier_rank": None, + "selected_matches_recommendation": False, + "selected_agent_mode": selected, + } + + budget_profile = str((auto_agent_policy or {}).get("budget_profile") or "manual") + cost_penalty = { + "fast": 0.35, + "balanced": 0.06, + "deep": 0.025, + "max": 0.0, + "manual": 0.03, + }.get(budget_profile, 0.06) + budget_state = self._route_frontier_budget_state(auto_agent_policy) + effective_cap = budget_state["effective_cap_cost_units"] + remaining_cap = budget_state["remaining_cost_units"] + pacing_cap = budget_state["pacing_cap_cost_units"] + session_budget_adjustment = ( + auto_agent_policy.get("session_budget_adjustment") + if isinstance(auto_agent_policy, dict) and isinstance(auto_agent_policy.get("session_budget_adjustment"), dict) + else {} + ) + minimum_route_floor = ( + selected == "off" + and session_budget_adjustment.get("reason") == "session_route_budget_exhausted_single_pass_floor" + ) + max_cost = max(float(row.get("estimated_cost_units") or 0.0) for row in alternatives) or 1.0 + + adaptive_quality_cost_available = False + for row in alternatives: + mode = str(row.get("selected_agent_mode") or "off") + cost = max(0.0, float(row.get("estimated_cost_units") or 0.0)) + quality = self._estimated_route_quality( + mode=mode, + action_mode=action_mode, + estimated_cost_units=cost, + auto_agent_policy=auto_agent_policy, + ) + adaptive_quality_cost_available = ( + adaptive_quality_cost_available or quality.get("quality_source") == "adaptive_feedback" + ) + fits_remaining = True if remaining_cap is None else cost <= (float(remaining_cap) + 1e-9) + fits_pacing = None if pacing_cap is None else cost <= (float(pacing_cap) + 1e-9) + budget_fit = True if effective_cap is None else cost <= (float(effective_cap) + 1e-9) + budget_overage = 0.0 if effective_cap is None else max(0.0, cost - float(effective_cap)) + if not fits_remaining: + budget_blocker = "remaining_budget" + elif fits_pacing is False: + budget_blocker = "pacing_cap" + elif not budget_fit: + budget_blocker = "effective_cap" + else: + budget_blocker = None + row.update( + { + "estimated_quality_score": quality["estimated_quality_score"], + "estimated_quality_cost_score": quality["estimated_quality_cost_score"], + "estimated_quality_lower_bound": quality["estimated_quality_lower_bound"], + "estimated_quality_upper_bound": quality["estimated_quality_upper_bound"], + "risk_adjusted_quality_cost_score": quality["risk_adjusted_quality_cost_score"], + "effective_sample_size": quality["effective_sample_size"], + "confidence_level": quality["confidence_level"], + "confidence_status": quality["confidence_status"], + "quality_source": quality["quality_source"], + "quality_cost_source": quality["quality_cost_source"], + "quality_evidence_status": quality["quality_evidence_status"], + "quality_evidence": quality["quality_evidence"], + "budget_fit": bool(budget_fit), + "fits_remaining_budget": bool(fits_remaining), + "fits_pacing_cap": fits_pacing, + "effective_cap_cost_units": effective_cap, + "minimum_route_floor": bool(minimum_route_floor and mode == "off"), + "budget_blocker": budget_blocker, + "budget_overage_units": round(budget_overage, 3), + } + ) + + for row in alternatives: + cost = max(0.0, float(row.get("estimated_cost_units") or 0.0)) + budget_fit = bool(row.get("budget_fit")) + budget_overage = float(row.get("budget_overage_units") or 0.0) + quality_score = float(row.get("estimated_quality_score") or 0.0) + if adaptive_quality_cost_available and budget_profile != "max": + adaptive_utility = row.get("risk_adjusted_quality_cost_score") + frontier_score = float( + adaptive_utility + if isinstance(adaptive_utility, (int, float)) + else row.get("estimated_quality_cost_score") or 0.0 + ) + (0.02 * quality_score) + else: + frontier_score = quality_score - (cost_penalty * (cost / max_cost)) + if not budget_fit: + frontier_score -= 0.5 + min(0.5, budget_overage / max_cost) + row["frontier_score"] = round(frontier_score, 4) + + for row in alternatives: + row_quality = float(row.get("estimated_quality_score") or 0.0) + row_cost = float(row.get("estimated_cost_units") or 0.0) + dominated = False + for other in alternatives: + if other is row: + continue + other_quality = float(other.get("estimated_quality_score") or 0.0) + other_cost = float(other.get("estimated_cost_units") or 0.0) + if ( + other_quality >= row_quality + and other_cost <= row_cost + and (other_quality > row_quality or other_cost < row_cost) + ): + dominated = True + break + row["pareto_frontier"] = not dominated + + feasible_for_pareto = [row for row in alternatives if bool(row.get("budget_fit"))] + for row in alternatives: + if not bool(row.get("budget_fit")): + row["budget_feasible_pareto_frontier"] = False + continue + row_quality = float(row.get("estimated_quality_score") or 0.0) + row_cost = float(row.get("estimated_cost_units") or 0.0) + dominated = False + for other in feasible_for_pareto: + if other is row: + continue + other_quality = float(other.get("estimated_quality_score") or 0.0) + other_cost = float(other.get("estimated_cost_units") or 0.0) + if ( + other_quality >= row_quality + and other_cost <= row_cost + and (other_quality > row_quality or other_cost < row_cost) + ): + dominated = True + break + row["budget_feasible_pareto_frontier"] = not dominated + + ranked = sorted( + alternatives, + key=lambda row: ( + bool(row.get("budget_fit")), + float(row.get("frontier_score") or 0.0), + float(row.get("estimated_quality_score") or 0.0), + -float(row.get("estimated_cost_units") or 0.0), + -AUTO_AGENT_MODE_ORDER.index(str(row.get("selected_agent_mode") or "off")), + ), + reverse=True, + ) + for idx, row in enumerate(ranked, start=1): + row["frontier_rank"] = idx + + feasible = [row for row in ranked if bool(row.get("budget_fit"))] + if feasible: + recommended = feasible[0] + else: + recommended = min( + ranked, + key=lambda row: ( + float(row.get("budget_overage_units") or 0.0), + float(row.get("estimated_cost_units") or 0.0), + -float(row.get("estimated_quality_score") or 0.0), + AUTO_AGENT_MODE_ORDER.index(str(row.get("selected_agent_mode") or "off")), + ), + ) + recommended_mode = str(recommended.get("selected_agent_mode") or selected) + selected_row = next((row for row in alternatives if row.get("selected_agent_mode") == selected), None) + cheapest_feasible = min(feasible, key=lambda row: float(row.get("estimated_cost_units") or 0.0)) if feasible else None + highest_quality_feasible = ( + max(feasible, key=lambda row: float(row.get("estimated_quality_score") or 0.0)) if feasible else None + ) + if not feasible: + remaining_fit_count = sum(1 for row in alternatives if bool(row.get("fits_remaining_budget"))) + reason = ( + "no_pacing_feasible_route" + if remaining_fit_count > 0 and pacing_cap is not None + else "no_budget_feasible_route" + ) + elif recommended_mode == selected: + reason = "selected_route_is_frontier_recommended" + elif adaptive_quality_cost_available and recommended.get("quality_source") == "adaptive_feedback": + reason = "adaptive_quality_cost_frontier_recommended" + elif float(recommended.get("estimated_cost_units") or 0.0) < float((selected_row or {}).get("estimated_cost_units") or 0.0): + reason = "cheaper_feasible_quality_cost_tradeoff" + else: + reason = "higher_quality_feasible_tradeoff" + budget_blockers = { + "remaining_budget": sum(1 for row in alternatives if row.get("budget_blocker") == "remaining_budget"), + "pacing_cap": sum(1 for row in alternatives if row.get("budget_blocker") == "pacing_cap"), + "effective_cap": sum(1 for row in alternatives if row.get("budget_blocker") == "effective_cap"), + "none": sum(1 for row in alternatives if row.get("budget_blocker") is None), + } + + return { + "selected_agent_mode": selected, + "recommended_agent_mode": recommended_mode, + "recommended_reason": reason, + "recommended_budget_blocker": recommended.get("budget_blocker"), + "recommended_estimated_cost_units": recommended.get("estimated_cost_units"), + "recommended_estimated_quality_score": recommended.get("estimated_quality_score"), + "recommended_estimated_quality_cost_score": recommended.get("estimated_quality_cost_score"), + "recommended_risk_adjusted_quality_cost_score": recommended.get("risk_adjusted_quality_cost_score"), + "recommended_quality_lower_bound": recommended.get("estimated_quality_lower_bound"), + "recommended_quality_upper_bound": recommended.get("estimated_quality_upper_bound"), + "budget_profile": budget_profile, + "budget_blockers": budget_blockers, + "budget_cap_cost_units": effective_cap, + "remaining_cost_units": remaining_cap, + "pacing_cap_cost_units": pacing_cap, + "effective_cap_cost_units": effective_cap, + "pacing_cap_applied": bool(budget_state["pacing_cap_applied"]), + "limit_cost_units": budget_state["limit_cost_units"], + "used_cost_units": budget_state["used_cost_units"], + "route_count": budget_state["route_count"], + "target_route_count": budget_state["target_route_count"], + "target_remaining_routes": budget_state["target_remaining_routes"], + "minimum_route_floor": bool(minimum_route_floor), + "budget_fit_count": len(feasible), + "selected_budget_fit": bool(selected_row.get("budget_fit")) if selected_row else None, + "selected_budget_blocker": selected_row.get("budget_blocker") if selected_row else None, + "selected_estimated_quality_cost_score": ( + selected_row.get("estimated_quality_cost_score") if selected_row else None + ), + "selected_risk_adjusted_quality_cost_score": ( + selected_row.get("risk_adjusted_quality_cost_score") if selected_row else None + ), + "selected_quality_lower_bound": ( + selected_row.get("estimated_quality_lower_bound") if selected_row else None + ), + "selected_quality_upper_bound": ( + selected_row.get("estimated_quality_upper_bound") if selected_row else None + ), + "selected_fits_remaining_budget": ( + bool(selected_row.get("fits_remaining_budget")) if selected_row else None + ), + "selected_fits_pacing_cap": selected_row.get("fits_pacing_cap") if selected_row else None, + "selected_frontier_rank": selected_row.get("frontier_rank") if selected_row else None, + "selected_matches_recommendation": recommended_mode == selected, + "cheapest_feasible_agent_mode": ( + cheapest_feasible.get("selected_agent_mode") if cheapest_feasible else None + ), + "highest_quality_feasible_agent_mode": ( + highest_quality_feasible.get("selected_agent_mode") if highest_quality_feasible else None + ), + "pareto_modes": [ + str(row.get("selected_agent_mode")) + for row in alternatives + if bool(row.get("pareto_frontier")) + ], + "budget_feasible_pareto_modes": [ + str(row.get("selected_agent_mode")) + for row in alternatives + if bool(row.get("budget_feasible_pareto_frontier")) + ], + "ranked_modes": [ + { + "selected_agent_mode": str(row.get("selected_agent_mode") or "off"), + "frontier_rank": int(row.get("frontier_rank") or 0), + "frontier_score": row.get("frontier_score"), + "estimated_quality_score": row.get("estimated_quality_score"), + "estimated_quality_cost_score": row.get("estimated_quality_cost_score"), + "estimated_quality_lower_bound": row.get("estimated_quality_lower_bound"), + "estimated_quality_upper_bound": row.get("estimated_quality_upper_bound"), + "risk_adjusted_quality_cost_score": row.get("risk_adjusted_quality_cost_score"), + "effective_sample_size": row.get("effective_sample_size"), + "confidence_level": row.get("confidence_level"), + "confidence_status": row.get("confidence_status"), + "quality_source": row.get("quality_source"), + "quality_cost_source": row.get("quality_cost_source"), + "quality_evidence_status": row.get("quality_evidence_status"), + "estimated_cost_units": row.get("estimated_cost_units"), + "budget_fit": bool(row.get("budget_fit")), + "fits_remaining_budget": bool(row.get("fits_remaining_budget")), + "fits_pacing_cap": row.get("fits_pacing_cap"), + "budget_blocker": row.get("budget_blocker"), + "budget_overage_units": row.get("budget_overage_units"), + "pareto_frontier": bool(row.get("pareto_frontier")), + "budget_feasible_pareto_frontier": bool(row.get("budget_feasible_pareto_frontier")), + } + for row in ranked + ], + } + + def _neighbor_auto_agent_mode(self, selected: str, allowed_modes: Sequence[str], direction: int) -> str: + ordered = [mode for mode in AUTO_AGENT_MODE_ORDER if mode in set(allowed_modes)] + if selected not in ordered: + return selected + idx = ordered.index(selected) + next_idx = idx + direction + if next_idx < 0 or next_idx >= len(ordered): + return selected + return ordered[next_idx] + + def _auto_route_uncertainty_signals(self, prompt: str) -> List[str]: + signals: List[str] = [] + raw_prompt = str(prompt or "") + for name, pattern in AUTO_AGENT_UNCERTAINTY_SIGNAL_PATTERNS: + if pattern.search(raw_prompt): + signals.append(name) + return signals + + def _auto_route_score_to_mode(self, mode: str, action_mode: str) -> Optional[int]: + if action_mode == "image" and mode == "collective": + return 3 + return AUTO_AGENT_SELECTION_THRESHOLDS.get(mode) + + def _auto_route_confidence( + self, + *, + selected: str, + score: int, + allowed_modes: Sequence[str], + action_mode: str, + uncertainty_signals: Sequence[str], + ) -> Dict[str, Any]: + next_mode = self._neighbor_auto_agent_mode(selected, allowed_modes, 1) + next_threshold = self._auto_route_score_to_mode(next_mode, action_mode) if next_mode != selected else None + score_to_next_mode: Optional[int] = None + if next_threshold is not None: + score_to_next_mode = max(0, int(next_threshold) - int(score)) + + level = "high" + if score_to_next_mode == 1 and uncertainty_signals: + level = "low" + elif score_to_next_mode == 1: + level = "medium" + + return { + "level": level, + "score": int(score), + "selected_agent_mode": selected, + "next_agent_mode": next_mode if next_mode != selected else None, + "next_agent_mode_threshold": next_threshold, + "score_to_next_mode": score_to_next_mode, + "uncertainty_signals": list(uncertainty_signals)[:6], + "adjusted": False, + } + + def _maybe_apply_auto_uncertainty_margin( + self, + *, + selected: str, + score: int, + allowed_modes: Sequence[str], + action_mode: str, + budget_profile: str, + route_confidence: Dict[str, Any], + uncertainty_signals: Sequence[str], + ) -> tuple[str, Optional[Dict[str, Any]], Dict[str, Any]]: + if ( + budget_profile == "fast" + or action_mode == "image" + or len(uncertainty_signals) < 2 + or route_confidence.get("score_to_next_mode") != 1 + ): + return selected, None, route_confidence + + upgraded = self._neighbor_auto_agent_mode(selected, allowed_modes, 1) + if upgraded == selected: + return selected, None, route_confidence + + adjustment = { + "direction": "upgrade", + "from": selected, + "to": upgraded, + "reason": "borderline_score_with_uncertainty_signals", + "score": int(score), + "score_to_next_mode": route_confidence.get("score_to_next_mode"), + "uncertainty_signals": list(uncertainty_signals)[:6], + "budget_profile": budget_profile, + } + adjusted_confidence = dict(route_confidence) + adjusted_confidence.update( + { + "adjusted": True, + "adjusted_from": selected, + "selected_agent_mode": upgraded, + "adjustment_reason": adjustment["reason"], + "level": "medium", + } + ) + return upgraded, adjustment, adjusted_confidence + + def _adaptive_route_evidence(self, stats: Any) -> Optional[Dict[str, Any]]: + if not isinstance(stats, dict): + return None + adaptive = stats.get("adaptive") if isinstance(stats.get("adaptive"), dict) else {} + economics = stats.get("economics") if isinstance(stats.get("economics"), dict) else {} + try: + economics_samples = int(economics.get("sample_count") or 0) + avg_cost = float(economics.get("avg_cost_units")) + except (TypeError, ValueError): + return None + try: + weighted_count = float(adaptive.get("weighted_count") or 0.0) + quality_score = float(adaptive.get("quality_score")) + quality_cost_score = float(adaptive.get("quality_cost_score")) + except (TypeError, ValueError): + return None + if ( + economics_samples < 1 + or not math.isfinite(avg_cost) + or not math.isfinite(weighted_count) + or not math.isfinite(quality_score) + or not math.isfinite(quality_cost_score) + or weighted_count < AUTO_AGENT_ADAPTIVE_MIN_WEIGHTED_COUNT + or quality_score < 0.0 + or quality_cost_score < 0.0 + or avg_cost < 0.0 + or bool(adaptive.get("regression_signal")) + ): + return None + quality_lower_bound = adaptive.get("quality_lower_bound") + quality_upper_bound = adaptive.get("quality_upper_bound") + quality_cost_lower_bound = adaptive.get("quality_cost_lower_bound") + quality_cost_upper_bound = adaptive.get("quality_cost_upper_bound") + effective_sample_size = adaptive.get("effective_sample_size") + try: + quality_lower_bound = float(quality_lower_bound) + quality_upper_bound = float(quality_upper_bound) + quality_cost_lower_bound = float(quality_cost_lower_bound) + quality_cost_upper_bound = float(quality_cost_upper_bound) + effective_sample_size = float(effective_sample_size) + except (TypeError, ValueError): + quality_lower_bound = None + quality_upper_bound = None + quality_cost_lower_bound = None + quality_cost_upper_bound = None + effective_sample_size = 0.0 + confidence_status = str(adaptive.get("confidence_status") or "no_evidence") + return { + "sample_count": int(adaptive.get("sample_count") or stats.get("count") or 0), + "economics_sample_count": economics_samples, + "avg_cost_units": round(avg_cost, 3), + "weighted_count": round(weighted_count, 3), + "quality_score": round(quality_score, 3), + "quality_cost_score": round(quality_cost_score, 3), + "quality_lower_bound": ( + round(quality_lower_bound, 3) if isinstance(quality_lower_bound, float) and math.isfinite(quality_lower_bound) else None + ), + "quality_upper_bound": ( + round(quality_upper_bound, 3) if isinstance(quality_upper_bound, float) and math.isfinite(quality_upper_bound) else None + ), + "quality_cost_lower_bound": ( + round(quality_cost_lower_bound, 3) + if isinstance(quality_cost_lower_bound, float) and math.isfinite(quality_cost_lower_bound) + else None + ), + "quality_cost_upper_bound": ( + round(quality_cost_upper_bound, 3) + if isinstance(quality_cost_upper_bound, float) and math.isfinite(quality_cost_upper_bound) + else None + ), + "effective_sample_size": ( + round(effective_sample_size, 3) if math.isfinite(effective_sample_size) else 0.0 + ), + "confidence_level": adaptive.get("confidence_level"), + "confidence_status": confidence_status, + "risk_adjusted_quality_cost_score": ( + round(quality_cost_lower_bound, 3) + if confidence_status == "established" + and isinstance(quality_cost_lower_bound, float) + and math.isfinite(quality_cost_lower_bound) + else round(quality_cost_score, 3) + ), + "weighted_net": adaptive.get("weighted_net"), + "decay": adaptive.get("decay"), + } + + def _adaptive_route_evidence_status(self, stats: Any) -> str: + if not isinstance(stats, dict): + return "heuristic_prior" + adaptive = stats.get("adaptive") if isinstance(stats.get("adaptive"), dict) else {} + economics = stats.get("economics") if isinstance(stats.get("economics"), dict) else {} + sample_count = 0 + try: + sample_count = int(adaptive.get("sample_count") or stats.get("count") or 0) + except (TypeError, ValueError): + sample_count = 0 + economics_sample_count = 0 + try: + economics_sample_count = int(economics.get("sample_count") or 0) + except (TypeError, ValueError): + economics_sample_count = 0 + if sample_count <= 0 and economics_sample_count <= 0: + return "heuristic_prior" + if bool(adaptive.get("regression_signal")): + return "adaptive_regression_blocked" + try: + avg_cost = float(economics.get("avg_cost_units")) + except (TypeError, ValueError): + return "adaptive_incomplete_cost_evidence" + if economics_sample_count < 1 or not math.isfinite(avg_cost) or avg_cost < 0.0: + return "adaptive_incomplete_cost_evidence" + try: + weighted_count = float(adaptive.get("weighted_count") or 0.0) + quality_score = float(adaptive.get("quality_score")) + quality_cost_score = float(adaptive.get("quality_cost_score")) + except (TypeError, ValueError): + return "adaptive_incomplete_feedback_evidence" + if ( + not math.isfinite(weighted_count) + or not math.isfinite(quality_score) + or not math.isfinite(quality_cost_score) + or weighted_count < AUTO_AGENT_ADAPTIVE_MIN_WEIGHTED_COUNT + or quality_score < 0.0 + or quality_cost_score < 0.0 + ): + return "adaptive_incomplete_feedback_evidence" + return "adaptive_complete" + + def _prefer_adaptive_neighbor_route( + self, + *, + selected: str, + score: int, + mode_scores: Dict[str, Any], + allowed_modes: Sequence[str], + budget_profile: str, + ) -> tuple[str, Optional[Dict[str, Any]]]: + selected_idx = AUTO_AGENT_MODE_ORDER.index(selected) if selected in AUTO_AGENT_MODE_ORDER else 0 + selected_stats = mode_scores.get(selected) if isinstance(mode_scores.get(selected), dict) else {} + selected_evidence = self._adaptive_route_evidence(selected_stats) + selected_quality = float((selected_evidence or {}).get("quality_score", 0.5)) + selected_confident = bool((selected_evidence or {}).get("confidence_status") == "established") + selected_quality_cost = float( + (selected_evidence or {}).get( + "quality_cost_upper_bound" if selected_confident else "quality_cost_score", + 0.5, + ) + ) + selected_has_evidence = selected_evidence is not None + + candidates: List[Dict[str, Any]] = [] + for direction in (-1, 1): + candidate = self._neighbor_auto_agent_mode(selected, allowed_modes, direction) + if candidate == selected: + continue + stats = mode_scores.get(candidate) if isinstance(mode_scores.get(candidate), dict) else {} + evidence = self._adaptive_route_evidence(stats) + if evidence is None: + continue + candidate_idx = AUTO_AGENT_MODE_ORDER.index(candidate) + moving_deeper = candidate_idx > selected_idx + threshold = int(AUTO_AGENT_POSITIVE_THRESHOLDS.get(candidate, 99)) + if moving_deeper and score < threshold: + continue + candidate_quality = float(evidence["quality_score"]) + use_confidence_bound = selected_confident and evidence.get("confidence_status") == "established" + candidate_quality_cost = float( + evidence["risk_adjusted_quality_cost_score" if use_confidence_bound else "quality_cost_score"] + ) + if candidate_quality < AUTO_AGENT_ADAPTIVE_QUALITY_FLOOR: + continue + + quality_delta = round(candidate_quality - selected_quality, 3) + quality_cost_delta = round(candidate_quality_cost - selected_quality_cost, 3) + if budget_profile == "max": + if quality_delta < AUTO_AGENT_ADAPTIVE_QUALITY_DELTA: + continue + rank_score = candidate_quality + elif selected_has_evidence: + if ( + quality_cost_delta < AUTO_AGENT_ADAPTIVE_QUALITY_COST_DELTA + and quality_delta < AUTO_AGENT_ADAPTIVE_QUALITY_DELTA + ): + continue + rank_score = candidate_quality_cost + max(0.0, quality_delta * 0.5) + else: + if moving_deeper: + if candidate_quality < 0.78: + continue + elif quality_cost_delta < AUTO_AGENT_ADAPTIVE_QUALITY_COST_DELTA: + continue + rank_score = candidate_quality_cost + + candidates.append( + { + "mode": candidate, + "direction": "upgrade" if moving_deeper else "downgrade", + "rank_score": rank_score, + "quality_delta": quality_delta, + "quality_cost_delta": quality_cost_delta, + "evidence": evidence, + "utility_source": ( + "confidence_lower_bound" + if use_confidence_bound + else "adaptive_mean" + ), + } + ) + + if not candidates: + return selected, None + chosen = max(candidates, key=lambda item: (float(item["rank_score"]), float(item["evidence"]["weighted_count"]))) + if chosen["mode"] == selected: + return selected, None + evidence = chosen["evidence"] + return str(chosen["mode"]), { + "direction": chosen["direction"], + "from": selected, + "to": str(chosen["mode"]), + "reason": "adaptive_quality_cost_preferred_neighbor", + "budget_profile": budget_profile, + "quality_score": evidence.get("quality_score"), + "quality_cost_score": evidence.get("quality_cost_score"), + "risk_adjusted_quality_cost_score": evidence.get("risk_adjusted_quality_cost_score"), + "quality_lower_bound": evidence.get("quality_lower_bound"), + "quality_upper_bound": evidence.get("quality_upper_bound"), + "effective_sample_size": evidence.get("effective_sample_size"), + "confidence_status": evidence.get("confidence_status"), + "utility_source": chosen.get("utility_source"), + "quality_delta": chosen["quality_delta"], + "quality_cost_delta": chosen["quality_cost_delta"], + "weighted_count": evidence.get("weighted_count"), + "weighted_net": evidence.get("weighted_net"), + "decay": evidence.get("decay"), + } + + def _apply_auto_route_feedback( + self, + *, + selected: str, + score: int, + feedback_summary: Dict[str, Any], + allowed_modes: Sequence[str], + budget_profile: str = "balanced", + ) -> tuple[str, Optional[Dict[str, Any]]]: + mode_scores = feedback_summary.get("mode_scores") if isinstance(feedback_summary, dict) else {} + if not isinstance(mode_scores, dict): + return selected, None + + selected_stats = mode_scores.get(selected) if isinstance(mode_scores.get(selected), dict) else {} + selected_adaptive = selected_stats.get("adaptive") if isinstance(selected_stats.get("adaptive"), dict) else {} + selected_net = int(selected_stats.get("quality_net", selected_stats.get("net")) or 0) + selected_negative = int(selected_stats.get("quality_negative", selected_stats.get("negative")) or 0) + + if not bool(feedback_summary.get("used_recent_fallback")) and bool(selected_adaptive.get("preference_signal")): + preference_direction = str(selected_adaptive.get("preference_direction") or "") + direction = 1 if preference_direction == "deeper" else -1 if preference_direction == "shallower" else 0 + preferred = self._neighbor_auto_agent_mode(selected, allowed_modes, direction) if direction else selected + if preferred != selected: + cost_pressure = float(selected_adaptive.get("weighted_cost_pressure") or 0.0) + latency_pressure = float(selected_adaptive.get("weighted_latency_pressure") or 0.0) + if direction > 0: + reason = "explicit_feedback_requested_deeper_route" + elif latency_pressure > cost_pressure: + reason = "explicit_feedback_requested_lower_latency_route" + else: + reason = "explicit_feedback_requested_lower_cost_route" + return preferred, { + "direction": "upgrade" if direction > 0 else "downgrade", + "from": selected, + "to": preferred, + "reason": reason, + "preference_direction": preference_direction, + "weighted_depth_preference": selected_adaptive.get("weighted_depth_preference"), + "weighted_cost_pressure": selected_adaptive.get("weighted_cost_pressure"), + "weighted_latency_pressure": selected_adaptive.get("weighted_latency_pressure"), + "feedback_scope": "prompt_session", + } + + if ( + not bool(feedback_summary.get("used_recent_fallback")) + and selected_negative >= 2 + and selected_net <= -2 + ): + downgraded = self._neighbor_auto_agent_mode(selected, allowed_modes, -1) + if downgraded != selected: + return downgraded, { + "direction": "downgrade", + "from": selected, + "to": downgraded, + "reason": "recent_session_feedback_rejected_selected_route", + "net": selected_net, + "negative": selected_negative, + } + + if ( + selected != "off" + and not bool(feedback_summary.get("used_recent_fallback")) + and bool(selected_adaptive.get("regression_signal")) + ): + downgraded = self._neighbor_auto_agent_mode(selected, allowed_modes, -1) + if downgraded != selected: + return downgraded, { + "direction": "downgrade", + "from": selected, + "to": downgraded, + "reason": "recent_weighted_feedback_regression", + "weighted_net": selected_adaptive.get("weighted_net"), + "weighted_negative": selected_adaptive.get("weighted_negative"), + "quality_score": selected_adaptive.get("quality_score"), + "recent_negative_rate": selected_adaptive.get("recent_negative_rate"), + "decay": selected_adaptive.get("decay"), + } + + if not bool(feedback_summary.get("used_recent_fallback")): + adjusted, adaptive_adjustment = self._prefer_adaptive_neighbor_route( + selected=selected, + score=score, + mode_scores=mode_scores, + allowed_modes=allowed_modes, + budget_profile=budget_profile, + ) + if adjusted != selected: + return adjusted, adaptive_adjustment + + upgraded = self._neighbor_auto_agent_mode(selected, allowed_modes, 1) + if not bool(feedback_summary.get("used_recent_fallback")) and upgraded != selected: + upgraded_stats = mode_scores.get(upgraded) if isinstance(mode_scores.get(upgraded), dict) else {} + upgraded_net = int(upgraded_stats.get("quality_net", upgraded_stats.get("net")) or 0) + upgraded_positive = int(upgraded_stats.get("quality_positive", upgraded_stats.get("positive")) or 0) + threshold = int(AUTO_AGENT_POSITIVE_THRESHOLDS.get(upgraded, 99)) + if upgraded_positive >= 2 and upgraded_net >= 2 and score >= threshold: + return upgraded, { + "direction": "upgrade", + "from": selected, + "to": upgraded, + "reason": "recent_session_feedback_preferred_deeper_route", + "net": upgraded_net, + "positive": upgraded_positive, + } + + if selected != "off" and budget_profile != "max": + selected_economics = ( + selected_stats.get("economics") + if isinstance(selected_stats.get("economics"), dict) + else {} + ) + economics = selected_economics + if isinstance(economics, dict): + sample_count = int(economics.get("sample_count") or 0) + relevant_feedback = int(feedback_summary.get("relevant_feedback") or 0) + used_recent_fallback = bool(feedback_summary.get("used_recent_fallback")) + avg_cost = float(economics.get("avg_cost_units") or 0.0) + avg_elapsed_ms = float(economics.get("avg_elapsed_ms") or 0.0) + budget_limits = { + "fast": {"cost": 3.0, "elapsed_ms": 2500.0}, + "balanced": {"cost": 6.0, "elapsed_ms": 8000.0}, + "deep": {"cost": 10.0, "elapsed_ms": 15000.0}, + } + limits = budget_limits.get(budget_profile, budget_limits["balanced"]) + selected_positive = int(selected_stats.get("quality_positive", selected_stats.get("positive")) or 0) + pressure_reasons: List[str] = [] + if avg_cost >= float(limits["cost"]): + pressure_reasons.append("cost") + if avg_elapsed_ms >= float(limits["elapsed_ms"]): + pressure_reasons.append("latency") + if ( + sample_count >= 1 + and relevant_feedback >= 2 + and not used_recent_fallback + and pressure_reasons + and selected_positive <= 0 + and selected_net <= 0 + ): + downgraded = self._neighbor_auto_agent_mode(selected, allowed_modes, -1) + if downgraded != selected: + return downgraded, { + "direction": "downgrade", + "from": selected, + "to": downgraded, + "reason": "session_route_economics_exceeded_budget_health", + "budget_profile": budget_profile, + "economic_pressure": pressure_reasons, + "pressure_scope": selected, + "sample_count": sample_count, + "relevant_feedback": relevant_feedback, + "avg_cost_units": round(avg_cost, 3), + "avg_elapsed_ms": round(avg_elapsed_ms, 3), + "cost_limit": limits["cost"], + "elapsed_ms_limit": limits["elapsed_ms"], + } + + return selected, None + + def _resolve_auto_agent_mode( + self, + *, + session_id: str, + prompt: str, + action_mode: str, + settings: Dict[str, Any], + chosen_record: ModelRecord, + ) -> Dict[str, Any]: + compute_policy = chat_app.estimate_auto_reasoning_cycles(prompt) + score = int(compute_policy.get("score") or 0) + reasons = list(compute_policy.get("reasons") or []) + lowered = str(prompt or "").lower() + + if bool(settings.get("web_search_enabled", False)) and should_offer_web_search(prompt): + score += 1 + reasons.append("fresh_knowledge") + if settings.get("uploaded_image_path"): + score += 1 + reasons.append("uploaded_artifact") + if re.search(r"\b(latest|recent|state of the art|sota|paper|github|cite|source)\b", lowered): + score += 1 + reasons.append("external_evidence") + if re.search(r"\b(end to end|whole project|multi[- ]agent|sub[- ]agent|regression|production)\b", lowered): + score += 1 + reasons.append("workflow_depth") + + consultants = self._collective_consultants(settings=settings, chosen_record=chosen_record) + collective_model_count = len(consultants) + collective_available = collective_model_count >= 2 + allow_collective = settings.get("auto_agent_collective", True) is not False + allow_loop = settings.get("auto_agent_loop", True) is not False + budget_profile = self._normalized_auto_budget_profile( + settings.get("auto_agent_budget") or settings.get("auto_route_budget") or settings.get("budget_profile") + ) + budget_policy = dict(AUTO_AGENT_BUDGET_PROFILES[budget_profile]) + score_before_budget = score + score = max(0, score + int(budget_policy.get("score_bias") or 0)) + if budget_policy.get("score_bias"): + reasons.append(f"budget_{budget_profile}") + + base_allowed_modes = self._allowed_auto_agent_modes( + action_mode=action_mode, + allow_collective=allow_collective, + collective_available=collective_available, + allow_loop=allow_loop, + ) + allowed_modes = self._budget_allowed_auto_agent_modes( + base_allowed_modes, + str(budget_policy.get("max_agent_mode") or "collective_loop"), + ) + + if action_mode == "image": + if allow_collective and collective_available and score >= 3: + selected = "collective" + reason = "complex_image_or_multimodal_prompt" + else: + selected = "off" + reason = "image_generation_prefers_single_pass" + elif score >= 5 and "collective_loop" in allowed_modes: + selected = "collective_loop" + reason = "high_complexity_with_collective_capacity" + elif score >= 4 and "loop" in allowed_modes: + selected = "loop" + reason = "high_complexity_iterative_work" + elif score >= 2 and "collective" in allowed_modes: + selected = "collective" + reason = "moderate_complexity_panel_synthesis" + else: + selected = "off" + reason = "low_complexity_single_pass" + if selected not in allowed_modes: + selected = self._neighbor_auto_agent_mode(allowed_modes[-1], allowed_modes, 0) if allowed_modes else "off" + reason = "budget_profile_limited_route_depth" + base_selected_agent_mode = selected + uncertainty_signals = self._auto_route_uncertainty_signals(prompt) + route_confidence = self._auto_route_confidence( + selected=selected, + score=score, + allowed_modes=allowed_modes, + action_mode=action_mode, + uncertainty_signals=uncertainty_signals, + ) + uncertainty_adjustment: Optional[Dict[str, Any]] = None + adjusted, uncertainty_adjustment, route_confidence = self._maybe_apply_auto_uncertainty_margin( + selected=selected, + score=score, + allowed_modes=allowed_modes, + action_mode=action_mode, + budget_profile=budget_profile, + route_confidence=route_confidence, + uncertainty_signals=uncertainty_signals, + ) + if adjusted != selected: + selected = adjusted + reason = str(uncertainty_adjustment.get("reason") or reason) + feedback_summary = self.memory_store.route_feedback_summary(session_id, prompt) + feedback_adjustment: Optional[Dict[str, Any]] = None + if int(feedback_summary.get("total_feedback") or 0) > 0: + adjusted, feedback_adjustment = self._apply_auto_route_feedback( + selected=selected, + score=score, + feedback_summary=feedback_summary, + allowed_modes=allowed_modes, + budget_profile=budget_profile, + ) + if adjusted != selected: + selected = adjusted + reason = str(feedback_adjustment.get("reason") or reason) + route_confidence["feedback_adjusted"] = True + route_confidence["selected_agent_mode"] = selected + + policy = { + "mode": "auto", + "selected_agent_mode": selected, + "base_selected_agent_mode": base_selected_agent_mode, + "reason": reason, + "score": score, + "score_before_budget": score_before_budget, + "reasons": reasons or ["simple_prompt"], + "reasoning_cycles": int(compute_policy.get("cycles") or 0), + "collective_model_count": collective_model_count, + "collective_available": collective_available, + "loop_available": action_mode != "image", + "budget_profile": budget_profile, + "budget_policy": budget_policy, + "allowed_agent_modes": allowed_modes, + "route_confidence": route_confidence, + "uncertainty_adjustment": uncertainty_adjustment, + "feedback_summary": feedback_summary, + "feedback_adjustment": feedback_adjustment, + } + return _stamp_auto_route_policy( + policy, + selected_agent_mode=selected, + action_mode=action_mode, + ) + + def _format_loop_history(self, steps: Sequence[Dict[str, Any]]) -> str: + if not steps: + return "No prior loop work yet." + lines: List[str] = [] + for row in steps: + lines.append(f"Step {row.get('step')}:") + if row.get("goal"): + lines.append(f"Goal: {row['goal']}") + if row.get("worker_excerpt"): + lines.append(f"Worker: {row['worker_excerpt']}") + if row.get("review_note"): + lines.append(f"Reviewer: {row['review_note']}") + if row.get("review_score") is not None: + lines.append(f"Verifier score: {row['review_score']}") + if row.get("next_step"): + lines.append(f"Next: {row['next_step']}") + return "\n".join(lines[-20:]) + + def _build_loop_planner_prompt(self, prompt: str, history: str, step_index: int, max_steps: int) -> str: + return ( + "You are the planner sub-agent inside a loop agent.\n" + "Break the user's task into the single highest-value next action.\n" + "Respond using exactly these headings:\n" + "DONE: yes or no\n" + "STEP_GOAL: one concise sentence\n" + "SUCCESS_SIGNAL: how to know the task is complete\n" + "WORKING_NOTES: short reasoning\n\n" + f"Original task:\n{prompt}\n\n" + f"Loop step: {step_index} of {max_steps}\n\n" + f"Work log:\n{history}" + ) + + def _build_loop_worker_prompt( + self, + prompt: str, + *, + history: str, + step_goal: str, + step_index: int, + max_steps: int, + ) -> str: + return ( + "You are the worker sub-agent inside a loop agent.\n" + "Execute the next best action for the task. Prefer doing the work over talking about the work.\n" + "Respond using exactly these headings:\n" + "DONE: yes or no\n" + "OUTPUT: the actual deliverable or best current progress\n" + "NEXT_FOCUS: what should happen next if not done\n\n" + f"Original task:\n{prompt}\n\n" + f"Loop step: {step_index} of {max_steps}\n" + f"Current goal:\n{step_goal}\n\n" + f"Work log:\n{history}" + ) + + def _build_loop_review_prompt( + self, + prompt: str, + *, + history: str, + latest_output: str, + success_signal: str, + step_index: int, + max_steps: int, + ) -> str: + return ( + "You are the reviewer sub-agent inside a loop agent.\n" + "Judge whether the user's task is complete enough to stop.\n" + "Respond using exactly these headings:\n" + "COMPLETE: yes or no\n" + "SCORE: number from 0.0 to 1.0 where 1.0 means fully complete and safe to stop\n" + "CONFIDENCE: number from 0.0 to 1.0\n" + "RISK_SCORE: number from 0.0 to 1.0 where 1.0 means high error or missing-work risk\n" + "FINAL_RESPONSE: the reply that should be shown to the user right now\n" + "REASON: why you marked it complete or incomplete\n" + "EVIDENCE: concrete evidence behind the score\n" + "NEXT_STEP: the next action if it is not complete\n\n" + f"Original task:\n{prompt}\n\n" + f"Loop step: {step_index} of {max_steps}\n" + f"Success signal:\n{success_signal or 'Task is complete when the user request has been fully addressed.'}\n\n" + f"Work log:\n{history}\n\n" + f"Latest worker output:\n{latest_output}" + ) + + def _run_loop_agent_text( + self, + *, + session_id: str, + prompt: str, + chosen_record: ModelRecord, + settings: Dict[str, Any], + route_reason: str, + action_mode: str, + memory_bundle: Dict[str, Any], + collective_mode: bool, + ) -> ChatResult: + controller_record = chosen_record if chosen_record.supports_chat else self._default_text_record() + max_steps = _coerce_int_setting( + settings.get("loop_max_steps") or settings.get("loop_budget"), + LOOP_AGENT_DEFAULT_MAX_STEPS, + minimum=2, + maximum=LOOP_AGENT_HARD_MAX_STEPS, + ) + score_threshold = _loop_score_threshold( + settings.get("loop_score_threshold") or settings.get("loop_completion_threshold") + ) + tool_cache = { + _trim_text(event.query, limit=220).lower(): event + for event in self._seed_auto_tool_events(prompt, settings) + } + tool_event_rows: List[Dict[str, Any]] = [event.to_dict() for event in list(tool_cache.values())] + consult_rows: List[Dict[str, str]] = [] + consulted_labels: List[str] = [] + skipped_models: List[Dict[str, str]] = [] + loop_steps: List[Dict[str, Any]] = [] + final_response = "" + completion_reason = "" + stop_reason_code = "budget_exhausted" + loop_stop_score: Optional[float] = None + completed = False + + for step_index in range(1, max_steps + 1): + history = self._format_loop_history(loop_steps) + planner_result, planner_events = self._run_text_model( + controller_record, + session_id=self._session_scope(session_id, controller_record.key, f"loop-plan-{step_index}"), + prompt=self._build_loop_planner_prompt(prompt, history, step_index, max_steps), + settings={ + **settings, + "memory_context": memory_bundle.get("context_block") or "", + }, + route_reason=f"{route_reason} Loop agent planner step {step_index} via {controller_record.label}.", + tool_cache=tool_cache, + allow_tool_calls=True, + ) + tool_event_rows.extend(event.to_dict() for event in planner_events) + step_goal = ( + _extract_labeled_section(planner_result.response, "STEP_GOAL") + or _extract_labeled_section(planner_result.response, "WORKING_NOTES") + or _trim_text(planner_result.response, limit=220) + or "Advance the task toward a complete solution." + ) + success_signal = _extract_labeled_section(planner_result.response, "SUCCESS_SIGNAL") + + if collective_mode: + worker_result = self._run_agent_text( + session_id=self._session_scope(session_id, controller_record.key, f"loop-work-{step_index}"), + prompt=self._build_loop_worker_prompt( + prompt, + history=history, + step_goal=step_goal, + step_index=step_index, + max_steps=max_steps, + ), + chosen_record=chosen_record, + settings=settings, + route_reason=f"{route_reason} Loop agent worker step {step_index} used the collective panel.", + action_mode=action_mode, + memory_bundle=memory_bundle, + ) + consult_rows.extend(list((worker_result.agent_trace or {}).get("consultation_rows") or [])) + consulted_labels.extend(list((worker_result.agent_trace or {}).get("consulted_models") or [])) + skipped_models.extend(list((worker_result.agent_trace or {}).get("skipped_models") or [])) + tool_event_rows.extend(list((worker_result.agent_trace or {}).get("tool_events") or [])) + else: + worker_result, worker_events = self._run_text_model( + controller_record, + session_id=self._session_scope(session_id, controller_record.key, f"loop-work-{step_index}"), + prompt=self._build_loop_worker_prompt( + prompt, + history=history, + step_goal=step_goal, + step_index=step_index, + max_steps=max_steps, + ), + settings={ + **settings, + "memory_context": memory_bundle.get("context_block") or "", }, route_reason=f"{route_reason} Loop agent worker step {step_index} via {controller_record.label}.", tool_cache=tool_cache, @@ -2254,6 +4155,7 @@ def _run_loop_agent_text( review_complete = _parse_yes_no_section(reviewer_result.response, "COMPLETE") if review_complete is None: review_complete = _parse_yes_no_section(reviewer_result.response, "DONE") + review_metrics = _loop_review_metrics(reviewer_result.response, review_complete) review_reason = ( _extract_labeled_section(reviewer_result.response, "REASON") or _extract_labeled_section(reviewer_result.response, "MISSING") @@ -2268,6 +4170,24 @@ def _run_loop_agent_text( or worker_output or _trim_text(reviewer_result.response, limit=1800) ) + score_stop = ( + review_complete is None + and review_metrics["review_score"] >= score_threshold + and review_metrics["risk_score"] <= 0.35 + ) + step_stop_reason = "" + stop_decision = "continue" + if review_complete is True: + stop_decision = "stop" + step_stop_reason = "reviewer_complete" + elif score_stop: + stop_decision = "stop" + step_stop_reason = "score_threshold" + elif review_complete is False: + step_stop_reason = "reviewer_continue" + else: + step_stop_reason = "score_below_threshold" + loop_steps.append( { "step": str(step_index), @@ -2275,13 +4195,30 @@ def _run_loop_agent_text( "worker_excerpt": _trim_text(worker_output, limit=320), "review_note": _trim_text(review_reason or "", limit=220), "next_step": _trim_text(next_step, limit=220), + "review_complete": review_complete, + "review_score": review_metrics["review_score"], + "loop_score": review_metrics["review_score"], + "verifier_score": review_metrics["verifier_score"], + "progress_score": review_metrics["progress_score"], + "confidence_score": review_metrics["confidence_score"], + "risk_score": review_metrics["risk_score"], + "completion_evidence": review_metrics["completion_evidence"], + "stop_decision": stop_decision, + "stop_reason_code": step_stop_reason, } ) final_response = final_candidate - completion_reason = review_reason or success_signal or f"Loop agent finished step {step_index}." + completion_reason = ( + review_reason + or review_metrics["completion_evidence"] + or success_signal + or f"Loop agent finished step {step_index}." + ) + loop_stop_score = review_metrics["review_score"] - if review_complete is True: + if stop_decision == "stop": completed = True + stop_reason_code = step_stop_reason break if not completed and final_response: @@ -2289,6 +4226,10 @@ def _run_loop_agent_text( f"{final_response}\n\nLoop agent note: the loop budget ended before it confidently marked the task complete." ) completion_reason = completion_reason or "Loop budget reached before completion." + elif completed and stop_reason_code == "score_threshold": + completion_reason = completion_reason or ( + f"Reviewer score met the {score_threshold:.2f} completion threshold." + ) deduped_consulted = list(dict.fromkeys(label for label in consulted_labels if label)) result = ChatResult( @@ -2300,7 +4241,13 @@ def _run_loop_agent_text( + (" with collective worker consultations." if collective_mode else f" on {controller_record.label}.") ), response=final_response or "Loop agent did not produce a final answer.", - timing={"loop_steps": len(loop_steps), "loop_max_steps": max_steps}, + timing={ + "loop_steps": len(loop_steps), + "loop_max_steps": max_steps, + "loop_stop_step": len(loop_steps), + "loop_stop_score": loop_stop_score, + "loop_stop_reason_code": stop_reason_code, + }, prompt_used=prompt, ) result.agent_trace = { @@ -2314,6 +4261,10 @@ def _run_loop_agent_text( "loop_completed": completed, "loop_completion_reason": completion_reason, "loop_budget": max_steps, + "loop_score_threshold": score_threshold, + "loop_stop_reason_code": stop_reason_code, + "loop_stop_step": len(loop_steps), + "loop_stop_score": loop_stop_score, "loop_controller_model": controller_record.label, "loop_worker_mode": "collective" if collective_mode else "single_model", } @@ -2470,7 +4421,13 @@ def _run_agent_image( f"back to {_record.label}." ) final_settings["route_reason"] = effective_route_reason - image_result = backend.generate_image(session_id, planner_result.response or prompt, final_settings) + generation_settings = dict(final_settings) + generation_settings.pop("_route_model_call_counter", None) + image_result = backend.generate_image(session_id, planner_result.response or prompt, generation_settings) + _record_route_model_call( + final_settings, + max(1, int((image_result.timing or {}).get("model_calls") or 1)), + ) image_result.agent_trace = { "agent_mode": "collective_panel", "memory_notes": list(memory_bundle.get("memory_notes") or []), @@ -2498,6 +4455,7 @@ def status(self) -> Dict[str, Any]: "exports_dir": str(self.exports_dir), "extraction_root": str(self.extraction_root), "memory_status": self.memory_store.global_status(), + "route_policy_ledger": self.route_policy_ledger.report(), "active_backend_status": self._backend.status() if self._backend is not None else None, } @@ -2505,6 +4463,249 @@ def session_memory_snapshot(self, session_id: str) -> Dict[str, Any]: with self._lock: return self.memory_store.session_snapshot(session_id) + def record_route_feedback(self, *, session_id: str, feedback: Dict[str, Any]) -> Dict[str, Any]: + with self._lock: + route_id = str(feedback.get("route_id") or "").strip() + try: + durable_decision = self.route_policy_ledger.get_decision(route_id) + except (DecisionNotFoundError, KeyError): + durable_decision = None + if durable_decision is not None: + if durable_decision.get("session_hash") != hash_session_identity(session_id): + raise ValueError("route_id does not belong to this session") + if durable_decision.get("status") != "completed": + raise ValueError("feedback requires a successfully completed route") + if durable_decision is None: + # Legacy feedback can predate the durable decision ledger. It + # remains descriptive in the JSON compatibility store and is + # never eligible for readiness certification. + result = self.memory_store.add_feedback(session_id=session_id, feedback=feedback) + result["durable_feedback"] = { + "status": "legacy_unjoined", + "eligible_for_readiness": False, + } + return result + + row = self.memory_store.prepare_feedback(session_id=session_id, feedback=feedback) + durable_payload = { + "rating": row.get("rating"), + "feedback_intent": row.get("feedback_intent"), + "feedback_tags": list(row.get("feedback_tags") or []), + "feedback_axes": dict(row.get("feedback_axes") or {}), + "reason": row.get("reason"), + "source": "explicit_user_route_feedback", + "observation_status": "observed", + } + # Commit the source-of-truth row first. Its content-derived + # idempotency key makes an identical retry return the same revision. + durable_feedback = self.route_policy_ledger.record_feedback(route_id, durable_payload) + durable_revision = max(1, int(durable_feedback.get("revision") or 1)) + try: + result = self.memory_store.commit_feedback( + session_id=session_id, + feedback_row=row, + feedback_revision=durable_revision, + ) + except Exception as exc: + # A compatibility-mirror failure must not turn an already + # durable user acknowledgement into an API failure. An + # identical retry safely reconciles the pinned revision. + logging.exception("Route feedback compatibility mirror is pending reconciliation") + accepted_row = dict(row) + accepted_row.update( + { + "feedback_revision": durable_revision, + "durable_feedback_revision": durable_revision, + } + ) + result = { + "ok": True, + "feedback": accepted_row, + "summary": None, + "compatibility_mirror": { + "status": "pending_reconciliation", + "durable_revision": durable_revision, + "error_category": type(exc).__name__, + }, + } + result["durable_feedback"] = { + "status": "committed", + "revision": durable_revision, + "idempotent": bool(durable_feedback.get("idempotent")), + "eligible_for_readiness": True, + } + return result + + def route_health_snapshot(self, session_id: str) -> Dict[str, Any]: + with self._lock: + summary = self.memory_store.route_feedback_summary(session_id) + summary["route_usage"] = self.memory_store.route_usage_summary(session_id) + return summary + + def route_shadow_registry_snapshot(self) -> Dict[str, Any]: + """Return the isolated shadow-assignment registry without changing it.""" + + registry_path = self.route_shadow_registry_path + with self._lock: + if not registry_path.is_file(): + self._route_shadow_registry_cache_signature = None + self._route_shadow_registry_cache_snapshot = None + return { + "ok": True, + "available": False, + "status": "not_initialized", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + + signature_before = self._route_shadow_registry_signature(registry_path) + if ( + signature_before == self._route_shadow_registry_cache_signature + and self._route_shadow_registry_cache_snapshot is not None + ): + return copy.deepcopy(self._route_shadow_registry_cache_snapshot) + + snapshot = RouteShadowAssignmentRegistry(registry_path, read_only=True).snapshot() + result = { + **snapshot, + "available": True, + "status": "verified" if snapshot.get("ok") else "verification_failed", + "registry_location": f"memory/{registry_path.name}", + "read_only": True, + } + signature_after = self._route_shadow_registry_signature(registry_path) + if signature_before == signature_after: + self._route_shadow_registry_cache_signature = signature_after + self._route_shadow_registry_cache_snapshot = copy.deepcopy(result) + else: + # A concurrent registry write may have landed during the audit. + # The transactional snapshot is still coherent, but do not reuse it. + self._route_shadow_registry_cache_signature = None + self._route_shadow_registry_cache_snapshot = None + return result + + @staticmethod + def _route_shadow_registry_signature( + registry_path: Path, + ) -> Tuple[Tuple[Any, ...], ...]: + """Fingerprint durable SQLite state without reading or mutating the registry.""" + + rows: List[Tuple[Any, ...]] = [] + for path in (registry_path, Path(f"{registry_path}-wal")): + try: + stat = path.stat() + except FileNotFoundError: + rows.append((False,)) + else: + # SQLite may materialize an empty WAL while opening a read-only + # connection. It carries no committed frames, so normalize it + # to the same signature as an absent WAL. + if path != registry_path and stat.st_size == 0: + rows.append((False,)) + continue + rows.append( + ( + True, + stat.st_dev, + stat.st_ino, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ) + ) + return tuple(rows) + + def route_policy_lab_snapshot(self, session_id: str, profile: str = "balanced") -> Dict[str, Any]: + with self._lock: + payload = self.memory_store.load_session(session_id) + durable_evidence = self.route_policy_ledger.policy_evidence_snapshot( + session_id=session_id, + policy_name=AUTO_ROUTE_POLICY_ID, + policy_version=AUTO_ROUTE_POLICY_VERSION, + limit=1000, + ) + durable_usage = list(durable_evidence.get("usage_rows") or []) + durable_feedback = list(durable_evidence.get("feedback_rows") or []) + if durable_usage: + usage_rows = durable_usage + feedback_rows = durable_feedback + expected_contexts = durable_evidence.get("expected_context_by_route_id") + evidence_source = "durable_sqlite_ledger" + lifecycle_source: Optional[Mapping[str, Any]] = durable_evidence + else: + usage_rows = list(payload.get("route_usage") or []) + feedback_rows = list(payload.get("route_feedback") or []) + expected_contexts = None + evidence_source = "legacy_json_compatibility" + lifecycle_source = None + report = analyze_route_policy( + usage_rows, + feedback_rows, + profile=profile, + expected_policy_id=AUTO_ROUTE_POLICY_ID, + expected_policy_version=AUTO_ROUTE_POLICY_VERSION, + expected_feature_schema=AUTO_ROUTE_FEATURE_SCHEMA_VERSION, + expected_context_by_route_id=( + expected_contexts if isinstance(expected_contexts, Mapping) else None + ), + durable_lifecycle=lifecycle_source, + ) + report["evidence_source"] = evidence_source + report["durable_evidence_window"] = durable_evidence.get("analysis_window") + outcome_contract_maturity = durable_evidence.get("outcome_contract_maturity") + report["outcome_contract_maturity"] = ( + dict(outcome_contract_maturity) + if isinstance(outcome_contract_maturity, Mapping) + else {} + ) + report["compatibility_view"] = { + "usage_rows": len(payload.get("route_usage") or []), + "feedback_rows": len(payload.get("route_feedback") or []), + "used_for_analysis": not bool(durable_usage), + "used_for_readiness": False, + "eligible_for_readiness": False, + } + report["available_profiles"] = [item.as_dict() for item in POLICY_PROFILES.values()] + report["active_logging_policy"] = { + "policy_id": AUTO_ROUTE_POLICY_ID, + "policy_version": AUTO_ROUTE_POLICY_VERSION, + "feature_schema_version": AUTO_ROUTE_FEATURE_SCHEMA_VERSION, + "decision_type": "deterministic", + "counterfactual_support": "none", + "support_schema_version": SUPPORT_SCHEMA_VERSION, + "failure_observation": "two_phase_started_and_terminal", + "durable_store": "sqlite_wal", + "retention": { + "durable_decisions": "uncapped", + "compatibility_usage_rows": 240, + "compatibility_feedback_rows": 120, + }, + } + durable = durable_evidence.get("lifecycle") or self.route_policy_ledger.report( + session_id=session_id, + policy_name=AUTO_ROUTE_POLICY_ID, + policy_version=AUTO_ROUTE_POLICY_VERSION, + ) + report["durable_ledger"] = durable + report["warnings"].append( + "Durable failure counts are descriptive runtime diagnostics; they do not create counterfactual support." + ) + report["warnings"].append( + "Outcome-contract maturity is diagnostic-only telemetry; it is not a policy-value estimator and " + "cannot authorize route promotion." + ) + if int((durable.get("counts") or {}).get("inflight") or 0) > 0: + report["warnings"].append( + "In-flight rows can include active work or interrupted processes and require terminal reconciliation." + ) + return report + def three_d_model_view(self) -> Dict[str, Any]: record = self.record_map.get("three_d_generation_micro_v1") if record is None: @@ -2624,6 +4825,20 @@ def export_chat_image( "saved_path": str(target), } + @staticmethod + def _route_error_category(exc: Exception) -> str: + if isinstance(exc, TimeoutError): + return "timeout" + if isinstance(exc, MemoryError): + return "resource_exhausted" + if isinstance(exc, FileNotFoundError): + return "model_artifact_missing" + if isinstance(exc, PermissionError): + return "permission_error" + if isinstance(exc, (KeyError, ValueError)): + return "validation_error" + return "runtime_error" + def handle_prompt( self, *, @@ -2632,8 +4847,65 @@ def handle_prompt( model_key: str, action_mode: str, settings: Dict[str, Any], + ) -> Dict[str, Any]: + self._route_execution.route_id = "" + self._route_execution.ledger_started = False + self._route_execution.started_at_perf = None + try: + return self._handle_prompt_impl( + session_id=session_id, + prompt=prompt, + model_key=model_key, + action_mode=action_mode, + settings=settings, + ) + except Exception as exc: + route_id = str(getattr(self._route_execution, "route_id", "") or "") + if route_id and bool(getattr(self._route_execution, "ledger_started", False)): + try: + failure_economics: Optional[Dict[str, float]] = None + started_at_perf = getattr(self._route_execution, "started_at_perf", None) + if ( + isinstance(started_at_perf, (int, float)) + and not isinstance(started_at_perf, bool) + and math.isfinite(float(started_at_perf)) + ): + completed_at_perf = time.perf_counter() + if math.isfinite(completed_at_perf) and completed_at_perf >= float(started_at_perf): + failure_economics = { + "elapsed_ms": round( + (completed_at_perf - float(started_at_perf)) * 1000.0, + 2, + ) + } + self.route_policy_ledger.complete_decision( + route_id, + success=False, + executed_mode=str(getattr(self._route_execution, "selected_mode", "") or "") or None, + actual_economics=failure_economics, + error_category=self._route_error_category(exc), + error_message=str(exc)[:500], + ) + except Exception: + logging.exception("Route decision failure completion could not be persisted") + raise + finally: + self._route_execution.route_id = "" + self._route_execution.ledger_started = False + self._route_execution.selected_mode = "" + self._route_execution.started_at_perf = None + + def _handle_prompt_impl( + self, + *, + session_id: str, + prompt: str, + model_key: str, + action_mode: str, + settings: Dict[str, Any], ) -> Dict[str, Any]: with self._lock: + route_started = time.perf_counter() requested_key = model_key or self.selected_model_key or "auto" if requested_key == "auto": chosen_record, route_reason = choose_auto_model( @@ -2658,6 +4930,8 @@ def handle_prompt( resolved_action = "text" settings = dict(settings or {}) + route_model_call_counter = {"count": 0} + settings["_route_model_call_counter"] = route_model_call_counter settings.setdefault("memory_enabled", True) settings.setdefault("agent_mode", "off") settings.setdefault("web_search_enabled", False) @@ -2667,14 +4941,162 @@ def handle_prompt( settings.setdefault("loop_max_steps", LOOP_AGENT_DEFAULT_MAX_STEPS) memory_bundle = self._prepare_memory_bundle(session_id, prompt, settings) agent_mode = self._normalized_agent_mode(settings.get("agent_mode")) - collective_enabled = agent_mode in {"collective", "collective_loop"} - loop_enabled = agent_mode in {"loop", "collective_loop"} - if loop_enabled and resolved_action == "image": - loop_enabled = False + requested_agent_mode = agent_mode + auto_agent_policy: Optional[Dict[str, Any]] = None + if agent_mode == "auto": + auto_agent_policy = self._resolve_auto_agent_mode( + session_id=session_id, + prompt=prompt, + action_mode=resolved_action, + settings=settings, + chosen_record=chosen_record, + ) + if settings.get("reasoning_cycles") in (None, "", 0, "0", "model", "default"): + settings["reasoning_cycles"] = int(auto_agent_policy.get("reasoning_cycles") or 1) + if settings.get("adaptive_compute") is None: + settings["adaptive_compute"] = True + auto_agent_policy["runtime_compute_request"] = { + "reasoning_cycles": settings.get("reasoning_cycles"), + "adaptive_compute": bool(settings.get("adaptive_compute")), + "source": "auto_route_policy", + } + agent_mode = str(auto_agent_policy["selected_agent_mode"]) + settings["agent_mode"] = agent_mode + budget_profile = str(auto_agent_policy.get("budget_profile") or "balanced") + budget_phrase = "" if budget_profile == "balanced" else f" under the {budget_profile} budget profile" + route_reason = ( + f"{route_reason} Auto orchestration selected {agent_mode} " + f"because {auto_agent_policy['reason']}{budget_phrase}." + ) + effective_agent_mode = self._effective_agent_mode_for_action(agent_mode, resolved_action) + if effective_agent_mode != agent_mode: + previous_agent_mode = agent_mode + agent_mode = effective_agent_mode + settings["agent_mode"] = agent_mode + if auto_agent_policy is not None: + auto_agent_policy["pre_action_agent_mode"] = previous_agent_mode + auto_agent_policy["selected_agent_mode"] = agent_mode + route_confidence = auto_agent_policy.get("route_confidence") + if isinstance(route_confidence, dict): + route_confidence["selected_agent_mode"] = agent_mode + route_confidence["action_adjusted"] = True route_reason = ( f"{route_reason} Loop agent currently supports text and vision replies; " - "image generation stayed on the normal one-pass path." + f"image generation used {agent_mode} instead of {previous_agent_mode}." + ) + route_economics_estimate = self._estimate_route_economics( + selected_agent_mode=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + ) + if auto_agent_policy is not None: + budget_limit = self._auto_session_budget_limit(settings) + budget_snapshot = self._session_route_budget_snapshot(session_id, budget_limit) + adjusted_agent_mode, route_economics_estimate, session_budget_adjustment, session_budget = ( + self._apply_auto_session_budget( + selected=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + budget_snapshot=budget_snapshot, + ) + ) + if session_budget is not None: + auto_agent_policy["session_budget"] = session_budget + if session_budget_adjustment is not None: + previous_agent_mode = agent_mode + agent_mode = adjusted_agent_mode + settings["agent_mode"] = agent_mode + auto_agent_policy["pre_session_budget_agent_mode"] = previous_agent_mode + auto_agent_policy["selected_agent_mode"] = agent_mode + auto_agent_policy["reason"] = str(session_budget_adjustment.get("reason") or auto_agent_policy["reason"]) + auto_agent_policy["session_budget_adjustment"] = session_budget_adjustment + route_confidence = auto_agent_policy.get("route_confidence") + if isinstance(route_confidence, dict): + route_confidence["session_budget_adjusted"] = True + route_confidence["selected_agent_mode"] = agent_mode + if previous_agent_mode == agent_mode: + route_reason = ( + f"{route_reason} Session budget is exhausted; Auto kept {agent_mode} " + "because single-pass execution is the minimum route." + ) + else: + if session_budget_adjustment.get("reason") == "session_route_budget_target_pacing": + route_reason = ( + f"{route_reason} Session budget target pacing downgraded {previous_agent_mode} " + f"to {agent_mode} to preserve budget for the configured route horizon." + ) + else: + route_reason = ( + f"{route_reason} Session budget pacing downgraded {previous_agent_mode} to {agent_mode} " + f"because the estimated route cost would exceed the remaining budget." + ) + else: + auto_agent_policy["session_budget_adjustment"] = None + auto_agent_policy["route_economics_estimate"] = dict(route_economics_estimate) + logging_support = self._build_post_filter_logging_support( + selected=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + selected_estimate=route_economics_estimate, + ) + _stamp_auto_route_policy( + auto_agent_policy, + selected_agent_mode=agent_mode, + action_mode=resolved_action, + logging_support=logging_support, ) + route_id = uuid.uuid4().hex + ledger_policy = auto_agent_policy if isinstance(auto_agent_policy, dict) else {} + ledger_context = dict(ledger_policy.get("decision_context") or {}) + ledger_context.update( + { + "action_mode": resolved_action, + "requested_model_key": requested_key, + "selected_model_key": chosen_record.key, + "requested_agent_mode": requested_agent_mode, + "selected_agent_mode": agent_mode, + "feedback_observation_status": "pending", + } + ) + eligible_modes = list( + ledger_policy.get("eligible_agent_modes") + or ledger_policy.get("eligible_actions") + or [agent_mode] + ) + action_probabilities = dict( + ledger_policy.get("post_filter_action_probabilities") + or ledger_policy.get("action_probabilities") + or {agent_mode: 1.0} + ) + self.route_policy_ledger.begin_decision( + session_id=session_id, + policy_name=str(ledger_policy.get("policy_id") or "explicit-route-v1"), + policy_version=str(ledger_policy.get("policy_version") or "1.0.0"), + policy_schema_version=str( + ledger_policy.get("feature_schema_version") or AUTO_ROUTE_FEATURE_SCHEMA_VERSION + ), + decision_context=ledger_context, + eligible_modes=eligible_modes, + chosen_mode=agent_mode, + action_probabilities=action_probabilities, + logging_support=( + ledger_policy.get("logging_support") + if isinstance(ledger_policy.get("logging_support"), dict) + else None + ), + estimated_economics=route_economics_estimate, + outcome_contracts=build_route_outcome_contracts(commitment_source="caller"), + route_id=route_id, + ) + self._route_execution.route_id = route_id + self._route_execution.ledger_started = True + self._route_execution.selected_mode = agent_mode + self._route_execution.started_at_perf = route_started + collective_enabled = agent_mode in {"collective", "collective_loop"} + loop_enabled = agent_mode in {"loop", "collective_loop"} if loop_enabled: result = self._run_loop_agent_text( @@ -2728,7 +5150,13 @@ def handle_prompt( if resolved_action == "image": if not record.supports_image: raise RuntimeError(f"{record.label} does not support image generation.") - result = backend.generate_image(session_id, prompt, base_settings) + generation_settings = dict(base_settings) + generation_settings.pop("_route_model_call_counter", None) + result = backend.generate_image(session_id, prompt, generation_settings) + _record_route_model_call( + base_settings, + max(1, int((result.timing or {}).get("model_calls") or 1)), + ) result.agent_trace = { "agent_mode": "off", "memory_notes": list(memory_bundle.get("memory_notes") or []), @@ -2736,7 +5164,13 @@ def handle_prompt( } else: if not record.supports_chat and record.supports_image: - result = backend.generate_image(session_id, prompt, base_settings) + generation_settings = dict(base_settings) + generation_settings.pop("_route_model_call_counter", None) + result = backend.generate_image(session_id, prompt, generation_settings) + _record_route_model_call( + base_settings, + max(1, int((result.timing or {}).get("model_calls") or 1)), + ) result.agent_trace = { "agent_mode": "off", "memory_notes": list(memory_bundle.get("memory_notes") or []), @@ -2758,9 +5192,42 @@ def handle_prompt( "tool_events": [event.to_dict() for event in list(tool_cache.values())], } + elapsed_ms = (time.perf_counter() - route_started) * 1000.0 + trace = dict(result.agent_trace or {}) + trace["route_id"] = route_id + if result.compute: + trace["compute"] = dict(result.compute) + route_economics = self._finalize_route_economics( + estimate=route_economics_estimate, + trace=trace, + elapsed_ms=elapsed_ms, + actual_model_calls=int(route_model_call_counter.get("count") or 0), + ) + trace["route_economics"] = route_economics + if auto_agent_policy: + trace["requested_agent_mode"] = "auto" + trace["resolved_agent_mode"] = agent_mode + auto_agent_policy["route_economics_actual"] = dict(route_economics["actual"]) + trace["auto_agent_policy"] = auto_agent_policy + result.agent_trace = trace + timing = dict(result.timing or {}) + timing["route_elapsed_ms"] = round(float(elapsed_ms), 2) + timing.setdefault("route_cost_units", route_economics["actual"]["cost_units"]) + result.timing = timing + assistant_summary = result.response or result.prompt_used or result.refined_prompt or "" tools_for_memory = list((result.agent_trace or {}).get("tool_events") or []) consultants_for_memory = list((result.agent_trace or {}).get("consultation_rows") or []) + self.memory_store.add_route_usage( + session_id=session_id, + route_id=route_id, + prompt=prompt, + selected_agent_mode=agent_mode, + route_economics=route_economics, + auto_agent_policy=auto_agent_policy, + route_reason=result.route_reason, + model_key=result.model_key, + ) self.memory_store.update( session_id=session_id, user_text=prompt, @@ -2778,4 +5245,377 @@ def handle_prompt( payload["active_model_key"] = result.model_key payload["active_model_label"] = result.model_label payload["active_model_kind"] = self.record_map[result.model_key].kind + payload["route_id"] = route_id + # The durable ledger stores a flat terminal economics record. Do + # not wrap the already-partitioned trace object a second time or + # replay will miss cost_units/elapsed_ms at the expected level. + ledger_actual_economics = dict(route_economics.get("actual") or {}) + ledger_actual_economics["executed_model_key"] = result.model_key + # A terminal-ledger failure is not a model failure. Clear the + # exception bridge before the atomic pending -> completed update so + # a persistence error leaves an honest in-flight row for repair. + self._route_execution.route_id = "" + self._route_execution.ledger_started = False + ledger_row = self.route_policy_ledger.complete_decision( + route_id, + success=True, + executed_mode=agent_mode, + actual_economics=ledger_actual_economics, + ) + payload_trace = payload.get("agent_trace") + if isinstance(payload_trace, dict): + payload_trace["route_ledger"] = { + "schema_version": ledger_row.get("ledger_schema_version"), + "status": ledger_row.get("status"), + "session_sequence": ledger_row.get("session_sequence"), + "feedback_status": ledger_row.get("feedback_status"), + "decision_type": ledger_row.get("decision_type"), + "probability_stage": ledger_row.get("probability_stage"), + "candidate_set_hash": ledger_row.get("candidate_set_hash"), + "distribution_hash": ledger_row.get("distribution_hash"), + } return payload + + def preview_route_plan( + self, + *, + session_id: str, + prompt: str, + model_key: str, + action_mode: str, + settings: Dict[str, Any], + ) -> Dict[str, Any]: + with self._lock: + requested_key = model_key or self.selected_model_key or "auto" + if requested_key == "auto": + chosen_record, route_reason = choose_auto_model( + self.records, + prompt, + action_mode=action_mode, + uploaded_image_path=str((settings or {}).get("uploaded_image_path") or ""), + ) + if chosen_record is None: + raise RuntimeError("No local models were discovered.") + else: + chosen_record = self.record_map.get(requested_key) + if chosen_record is None: + raise KeyError(f"Unknown model key: {requested_key}") + route_reason = f"Manual selection kept {chosen_record.label}." + + resolved_action = action_mode + if resolved_action == "auto": + resolved_action = "image" if chosen_record.supports_image and not chosen_record.supports_chat else "text" + + settings = dict(settings or {}) + settings.setdefault("agent_mode", "off") + settings.setdefault("web_search_enabled", False) + settings.setdefault("cmd_open_enabled", True) + settings.setdefault("web_search_budget", 3) + settings.setdefault("web_search_results", 5) + settings.setdefault("loop_max_steps", LOOP_AGENT_DEFAULT_MAX_STEPS) + + requested_agent_mode = self._normalized_agent_mode(settings.get("agent_mode")) + agent_mode = requested_agent_mode + auto_agent_policy: Optional[Dict[str, Any]] = None + if agent_mode == "auto": + auto_agent_policy = self._resolve_auto_agent_mode( + session_id=session_id, + prompt=prompt, + action_mode=resolved_action, + settings=settings, + chosen_record=chosen_record, + ) + if settings.get("reasoning_cycles") in (None, "", 0, "0", "model", "default"): + settings["reasoning_cycles"] = int(auto_agent_policy.get("reasoning_cycles") or 1) + if settings.get("adaptive_compute") is None: + settings["adaptive_compute"] = True + auto_agent_policy["runtime_compute_request"] = { + "reasoning_cycles": settings.get("reasoning_cycles"), + "adaptive_compute": bool(settings.get("adaptive_compute")), + "source": "auto_route_policy", + } + agent_mode = str(auto_agent_policy["selected_agent_mode"]) + settings["agent_mode"] = agent_mode + budget_profile = str(auto_agent_policy.get("budget_profile") or "balanced") + budget_phrase = "" if budget_profile == "balanced" else f" under the {budget_profile} budget profile" + route_reason = ( + f"{route_reason} Auto orchestration selected {agent_mode} " + f"because {auto_agent_policy['reason']}{budget_phrase}." + ) + effective_agent_mode = self._effective_agent_mode_for_action(agent_mode, resolved_action) + if effective_agent_mode != agent_mode: + previous_agent_mode = agent_mode + agent_mode = effective_agent_mode + settings["agent_mode"] = agent_mode + if auto_agent_policy is not None: + auto_agent_policy["pre_action_agent_mode"] = previous_agent_mode + auto_agent_policy["selected_agent_mode"] = agent_mode + route_confidence = auto_agent_policy.get("route_confidence") + if isinstance(route_confidence, dict): + route_confidence["selected_agent_mode"] = agent_mode + route_confidence["action_adjusted"] = True + route_reason = ( + f"{route_reason} Loop agent currently supports text and vision replies; " + f"image generation would use {agent_mode} instead of {previous_agent_mode}." + ) + + route_economics_estimate = self._estimate_route_economics( + selected_agent_mode=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + ) + if auto_agent_policy is not None: + budget_limit = self._auto_session_budget_limit(settings) + budget_snapshot = self._session_route_budget_snapshot(session_id, budget_limit) + adjusted_agent_mode, route_economics_estimate, session_budget_adjustment, session_budget = ( + self._apply_auto_session_budget( + selected=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + budget_snapshot=budget_snapshot, + ) + ) + if session_budget is not None: + auto_agent_policy["session_budget"] = session_budget + if session_budget_adjustment is not None: + previous_agent_mode = agent_mode + agent_mode = adjusted_agent_mode + settings["agent_mode"] = agent_mode + auto_agent_policy["pre_session_budget_agent_mode"] = previous_agent_mode + auto_agent_policy["selected_agent_mode"] = agent_mode + auto_agent_policy["reason"] = str(session_budget_adjustment.get("reason") or auto_agent_policy["reason"]) + auto_agent_policy["session_budget_adjustment"] = session_budget_adjustment + route_confidence = auto_agent_policy.get("route_confidence") + if isinstance(route_confidence, dict): + route_confidence["session_budget_adjusted"] = True + route_confidence["selected_agent_mode"] = agent_mode + if previous_agent_mode == agent_mode: + route_reason = ( + f"{route_reason} Session budget is exhausted; Auto kept {agent_mode} " + "because single-pass execution is the minimum route." + ) + elif session_budget_adjustment.get("reason") == "session_route_budget_target_pacing": + route_reason = ( + f"{route_reason} Session budget target pacing downgraded {previous_agent_mode} " + f"to {agent_mode} to preserve budget for the configured route horizon." + ) + else: + route_reason = ( + f"{route_reason} Session budget pacing downgraded {previous_agent_mode} to {agent_mode} " + f"because the estimated route cost would exceed the remaining budget." + ) + else: + auto_agent_policy["session_budget_adjustment"] = None + auto_agent_policy["route_economics_estimate"] = dict(route_economics_estimate) + logging_support = self._build_post_filter_logging_support( + selected=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + selected_estimate=route_economics_estimate, + ) + _stamp_auto_route_policy( + auto_agent_policy, + selected_agent_mode=agent_mode, + action_mode=resolved_action, + logging_support=logging_support, + ) + + collective_enabled = agent_mode in {"collective", "collective_loop"} + loop_enabled = agent_mode in {"loop", "collective_loop"} and resolved_action != "image" + route_alternatives = self._preview_route_alternatives( + selected=agent_mode, + action_mode=resolved_action, + settings=settings, + auto_agent_policy=auto_agent_policy, + ) + route_frontier = self._annotate_route_frontier( + alternatives=route_alternatives, + selected=agent_mode, + action_mode=resolved_action, + auto_agent_policy=auto_agent_policy, + ) + return { + "ok": True, + "dry_run": True, + "session_id": session_id, + "requested_model_key": requested_key, + "active_model_key": chosen_record.key, + "active_model_label": chosen_record.label, + "active_model_kind": chosen_record.kind, + "action_mode": resolved_action, + "requested_agent_mode": requested_agent_mode, + "selected_agent_mode": agent_mode, + "route_reason": route_reason, + "route_economics_estimate": route_economics_estimate, + "route_alternatives": route_alternatives, + "route_frontier": route_frontier, + "auto_agent_policy": auto_agent_policy, + "execution_plan": { + "collective_enabled": collective_enabled, + "loop_enabled": loop_enabled, + "single_pass": not collective_enabled and not loop_enabled, + "will_write_memory": False, + "will_run_inference": False, + }, + } + + def preview_route_study( + self, + *, + session_id: str, + prompt: str, + model_key: str, + action_mode: str, + settings: Dict[str, Any], + exploration_rate: float = 0.10, + planned_routes: int = 2_000, + scenario_confidence: float = 0.95, + assumed_feedback_rate: float = 0.30, + target_observed_labels: int = 20, + target_policy_profile: str = "balanced", + protocol_design_mode: str = "sticky_session_cluster", + carryover_scope: str = "unknown", + interference_scope: str = "unknown", + temporal_variation: str = "unknown", + planned_clusters: int = 200, + max_routes_per_cluster: int = 20, + analysis_every_clusters: int = 50, + block_length_routes: int = 20, + washout_routes: int = 0, + ) -> Dict[str, Any]: + """Rehearse a prompt-specific adjacent-route charter without side effects. + + The normal preview path performs the final capability and session-budget + filtering. This method projects that immutable support into a distinct + study cohort, but deliberately never samples, executes, or persists an + assignment. Rehearsed probabilities must not be treated as executed + logging propensities. + """ + + preview_settings = dict(settings or {}) + if self._normalized_agent_mode(preview_settings.get("agent_mode")) != "auto": + raise ValueError("Adjacent-route study rehearsal requires Auto Router mode") + preview = self.preview_route_plan( + session_id=session_id, + prompt=prompt, + model_key=model_key, + action_mode=action_mode, + settings=preview_settings, + ) + policy = preview.get("auto_agent_policy") + if not isinstance(policy, Mapping): + raise ValueError("Auto Router did not produce a route policy for study rehearsal") + support = policy.get("logging_support") + if not isinstance(support, Mapping): + raise ValueError("Auto Router did not produce final post-filter support") + study = plan_adjacent_route_study( + str(preview.get("selected_agent_mode") or ""), + list(support.get("candidates") or []), + list(support.get("exclusions") or []), + source_contract={ + "policy_id": policy.get("policy_id"), + "policy_version": policy.get("policy_version"), + "feature_schema_version": policy.get("feature_schema_version"), + "support_schema_version": support.get("schema_version"), + "candidate_set_hash": support.get("candidate_set_hash"), + "distribution_hash": support.get("distribution_hash"), + "outcome_contract_schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + }, + exploration_rate=exploration_rate, + planned_routes=planned_routes, + scenario_confidence=scenario_confidence, + assumed_feedback_rate=assumed_feedback_rate, + target_observed_labels=target_observed_labels, + ) + protocol_preflight = None + protocol_preflight_reason = "no_eligible_adjacent_support" + if study["charter"]["enrollment"]["eligible"] is True: + protocol_preflight = build_route_study_protocol( + study, + target_policy_profile=target_policy_profile, + design_mode=protocol_design_mode, + carryover_scope=carryover_scope, + interference_scope=interference_scope, + temporal_variation=temporal_variation, + population_rule_id="interactive-auto-route-opt-in", + population_rule_version="1", + cluster_key_schema_version="session-hash-v1", + planned_clusters=planned_clusters, + max_routes_per_cluster=max_routes_per_cluster, + analysis_every_clusters=analysis_every_clusters, + block_length_routes=block_length_routes, + washout_routes=washout_routes, + ) + protocol_preflight_reason = "draft_for_independent_review" + return { + "ok": True, + "dry_run": True, + "session_id": session_id, + "requested_model_key": preview.get("requested_model_key"), + "active_model_key": preview.get("active_model_key"), + "active_model_label": preview.get("active_model_label"), + "action_mode": preview.get("action_mode"), + "baseline_agent_mode": preview.get("selected_agent_mode"), + "route_reason": preview.get("route_reason"), + "deterministic_support": { + "policy_id": policy.get("policy_id"), + "policy_version": policy.get("policy_version"), + "feature_schema_version": policy.get("feature_schema_version"), + "support_schema_version": support.get("schema_version"), + "candidate_set_hash": policy.get("candidate_set_hash"), + "distribution_hash": policy.get("distribution_hash"), + "outcome_contract_schema_version": OUTCOME_CONTRACT_SCHEMA_VERSION, + }, + "route_study": study, + "route_protocol_preflight": protocol_preflight, + "route_protocol_preflight_reason": protocol_preflight_reason, + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } + + def build_route_protocol_review_bundle(self, payload: Mapping[str, Any]) -> Dict[str, Any]: + """Build and fully reconstruct a prompt-free multi-stratum review bundle.""" + + bundle = build_route_study_review_bundle_from_input(payload) + verification = audit_route_study_review_bundle(bundle) + return { + "ok": True, + "dry_run": True, + "route_protocol_review_bundle": bundle, + "verification": verification, + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } + + def audit_route_protocol_review_bundle(self, bundle: Any) -> Dict[str, Any]: + """Perform full source-bound reconstruction without touching runtime state.""" + + verification = audit_route_study_review_bundle(bundle) + return { + "ok": True, + "dry_run": True, + "verification": verification, + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } diff --git a/source/supermix_multimodel_web_app.py b/source/supermix_multimodel_web_app.py index 68681033c..004fd942b 100644 --- a/source/supermix_multimodel_web_app.py +++ b/source/supermix_multimodel_web_app.py @@ -18,6 +18,8 @@ app = Flask(__name__) manager: UnifiedModelManager | None = None +MAX_ROUTE_REVIEW_BUNDLE_REQUEST_BYTES = 2 * 1024 * 1024 +MAX_ROUTE_REVIEW_BUNDLE_WEB_STRATA = 100 def build_app(unified_manager: UnifiedModelManager) -> Flask: global manager @@ -25,6 +27,68 @@ def build_app(unified_manager: UnifiedModelManager) -> Flask: return app +def _validate_route_review_request_size(value: Any) -> None: + try: + size = len( + json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + ) + except (TypeError, ValueError) as exc: + raise ValueError("route protocol review request must be finite canonical JSON") from exc + if size > MAX_ROUTE_REVIEW_BUNDLE_REQUEST_BYTES: + raise ValueError("route protocol review request exceeds the 2 MiB browser limit") + + +def _read_strict_route_review_json() -> Any: + """Parse the integrity-sensitive review surface without last-key-wins JSON.""" + + raw = request.get_data(cache=True) + if not raw: + raise ValueError("route protocol review request body must contain JSON") + if len(raw) > MAX_ROUTE_REVIEW_BUNDLE_REQUEST_BYTES: + raise ValueError("route protocol review request exceeds the 2 MiB browser limit") + + def reject_duplicate_keys(pairs): + value = {} + for key, item in pairs: + if key in value: + raise ValueError( + f"route protocol review JSON contains duplicate object key: {key}" + ) + value[key] = item + return value + + def reject_non_finite(token): + raise ValueError( + f"route protocol review JSON contains non-finite number: {token}" + ) + + try: + return json.loads( + raw.decode("utf-8-sig"), + object_pairs_hook=reject_duplicate_keys, + parse_constant=reject_non_finite, + ) + except UnicodeDecodeError as exc: + raise ValueError("route protocol review request must be UTF-8 JSON") from exc + except json.JSONDecodeError as exc: + raise ValueError( + f"route protocol review request is not valid JSON: {exc.msg}" + ) from exc + + +def _validate_route_review_strata(value: Any) -> None: + if not isinstance(value, list) or not value: + raise ValueError("route protocol review requires a non-empty study_plans list") + if len(value) > MAX_ROUTE_REVIEW_BUNDLE_WEB_STRATA: + raise ValueError("route protocol review supports at most 100 browser strata") + + # ─── Benchmark graph embed helper ─────────────────────────────────────────── def _bench_graph_b64() -> str: """Return base64-encoded PNG of the benchmark graph, or empty string.""" @@ -240,6 +304,14 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: .model-pill.v46 { background:rgba(52,211,153,0.13); border-color:rgba(52,211,153,0.45); color:#86efac; box-shadow:0 0 20px rgba(52,211,153,0.13); } + .panel-toggle { display:none; align-items:center; gap:8px; min-height:38px; padding:8px 12px; + border:1px solid rgba(56,189,248,.28); border-radius:12px; + background:rgba(56,189,248,.09); color:#bae6fd; font-size:11px; + font-weight:900; letter-spacing:.06em; text-transform:uppercase; } + .panel-toggle:hover { background:rgba(56,189,248,.16); border-color:rgba(56,189,248,.5); } + .panel-toggle:focus-visible, .panel-close:focus-visible { + outline:2px solid var(--blue); outline-offset:2px; + } /* ── Thread ─────────────────────────────────────────────────────── */ .thread { padding:40px 14%; overflow-y:auto; display:flex; @@ -329,6 +401,109 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: border:1.5px solid var(--teal); display:flex; align-items:center; justify-content:center; font-size:10px; font-weight:800; color:var(--teal); } + .trace-summary { display:flex; flex-wrap:wrap; gap:10px; align-items:center; } + .trace-pill { display:inline-flex; margin:0 6px 4px 0; + border:1px solid rgba(255,255,255,.10); border-radius:999px; + padding:5px 9px; color:var(--muted); background:rgba(255,255,255,.03); } + .trace-score { color:var(--green); border-color:rgba(52,211,153,.28); } + .route-feedback { margin-top:12px; display:flex; align-items:center; gap:8px; flex-wrap:wrap; } + .route-feedback-label { color:var(--muted); font-size:11px; font-weight:900; + text-transform:uppercase; letter-spacing:.12em; margin-right:2px; } + .route-feedback button { border:1px solid var(--border); border-radius:999px; + background:rgba(255,255,255,.04); color:var(--muted); + padding:6px 10px; font-size:11px; font-weight:900; } + .route-feedback button:hover { color:var(--text); border-color:rgba(129,140,248,.4); + background:rgba(129,140,248,.09); } + .route-feedback button:disabled { opacity:.5; cursor:default; } + .route-health { margin-top:10px; display:flex; align-items:center; gap:7px; + flex-wrap:wrap; color:var(--muted); } + .route-health span { display:inline-flex; border:1px solid rgba(56,189,248,.18); + background:rgba(56,189,248,.06); border-radius:999px; + padding:6px 9px; font-size:10.5px; font-weight:900; } + .policy-lab { margin-top:14px; padding:13px; border:1px solid rgba(129,140,248,.2); + border-radius:14px; background:linear-gradient(145deg,rgba(129,140,248,.08),rgba(56,189,248,.035)); } + .policy-lab-head { display:flex; align-items:center; justify-content:space-between; gap:8px; } + .policy-lab-heading { display:flex; align-items:center; flex-wrap:wrap; gap:7px; min-width:0; } + .policy-lab-title { font-size:11px; font-weight:900; letter-spacing:.12em; text-transform:uppercase; } + .policy-lab-source { padding:3px 7px; border:1px solid rgba(56,189,248,.3); border-radius:999px; + background:rgba(56,189,248,.08); color:#7dd3fc; font-size:9px; font-weight:800; + letter-spacing:.04em; text-transform:uppercase; } + .policy-lab-controls { display:flex; gap:6px; align-items:center; } + .policy-lab select,.policy-lab button { border:1px solid var(--border); border-radius:8px; + background:rgba(8,12,25,.75); color:var(--text); + padding:5px 7px; font-size:10.5px; font-weight:800; } + .policy-lab button:disabled { opacity:.55; cursor:wait; } + .policy-lab-metrics { margin-top:10px; display:grid; grid-template-columns:repeat(2,minmax(0,1fr)); gap:7px; } + .policy-lab-metric { border:1px solid rgba(255,255,255,.07); border-radius:9px; + padding:8px; background:rgba(0,0,0,.16); } + .policy-lab-metric b { display:block; font-size:13px; color:var(--text); } + .policy-lab-metric span { display:block; margin-top:2px; color:var(--muted); font-size:9.5px; } + .policy-lab-gate { margin-top:9px; font-size:10.5px; line-height:1.45; color:#fbbf24; } + .policy-lab-readiness { margin-top:9px; padding:9px; border:1px solid rgba(255,255,255,.07); + border-radius:10px; background:rgba(0,0,0,.14); } + .policy-lab-readiness-title { margin-bottom:7px; color:var(--muted); font-size:9px; + font-weight:900; letter-spacing:.12em; text-transform:uppercase; } + .policy-lab-checks { display:grid; grid-template-columns:1fr; gap:4px; } + .policy-lab-check { display:flex; align-items:flex-start; gap:6px; padding:5px 6px; + border-radius:7px; color:var(--muted); font-size:9.5px; line-height:1.35; } + .policy-lab-check::before { content:'?'; flex:0 0 14px; height:14px; border-radius:50%; + display:inline-flex; align-items:center; justify-content:center; + background:rgba(148,163,184,.12); color:#cbd5e1; font-size:8px; font-weight:900; } + .policy-lab-check[data-state="pass"] { background:rgba(52,211,153,.06); color:#a7f3d0; } + .policy-lab-check[data-state="pass"]::before { content:'\2713'; background:rgba(52,211,153,.18); color:#6ee7b7; } + .policy-lab-check[data-state="fail"] { background:rgba(251,113,133,.06); color:#fecdd3; } + .policy-lab-check[data-state="fail"]::before { content:'\00D7'; background:rgba(251,113,133,.17); color:#fda4af; } + .policy-lab-blockers { margin-top:7px; color:#fcd34d; font-size:9.5px; line-height:1.45; } + .policy-lab-warning { margin-top:7px; padding:7px 8px; border:1px solid rgba(251,191,36,.32); + border-radius:8px; background:rgba(251,191,36,.08); color:#fde68a; + font-size:10px; font-weight:700; line-height:1.45; } + .policy-lab-note { margin-top:6px; font-size:9.5px; line-height:1.45; color:var(--muted); } + .route-study { margin-top:12px; padding:13px; border:1px solid rgba(45,212,191,.24); + border-radius:14px; background:linear-gradient(145deg,rgba(45,212,191,.075),rgba(56,189,248,.025)); } + .route-study-head { display:flex; align-items:flex-start; justify-content:space-between; gap:8px; } + .route-study-heading { min-width:0; } + .route-study-title { font-size:11px; font-weight:900; letter-spacing:.12em; text-transform:uppercase; } + .route-study-badge { display:inline-flex; margin-top:5px; padding:3px 7px; border:1px solid rgba(45,212,191,.32); + border-radius:999px; background:rgba(45,212,191,.08); color:#99f6e4; + font-size:8.5px; font-weight:900; letter-spacing:.08em; text-transform:uppercase; } + .route-study button { border:1px solid rgba(45,212,191,.28); border-radius:8px; + background:rgba(8,12,25,.75); color:#ccfbf1; padding:6px 8px; + font-size:10.5px; font-weight:900; } + .route-study button:disabled { opacity:.55; cursor:wait; } + .route-study-controls { margin-top:10px; display:grid; grid-template-columns:repeat(2,minmax(0,1fr)); gap:7px; } + .route-study-control { display:flex; flex-direction:column; gap:4px; min-width:0; color:var(--muted); + font-size:8.5px; font-weight:900; letter-spacing:.06em; text-transform:uppercase; } + .route-study-control input,.route-study-control select { width:100%; min-width:0; border:1px solid var(--border); + border-radius:8px; background:rgba(8,12,25,.75); + color:var(--text); padding:6px 7px; font-size:10.5px; } + .route-study-status { margin-top:9px; padding:8px; border:1px solid rgba(251,191,36,.28); + border-radius:9px; background:rgba(251,191,36,.065); color:#fde68a; + font-size:10px; font-weight:800; line-height:1.45; } + .route-study-metrics { margin-top:8px; display:grid; grid-template-columns:repeat(2,minmax(0,1fr)); gap:7px; } + .route-study-metric { padding:8px; border:1px solid rgba(255,255,255,.07); border-radius:9px; + background:rgba(0,0,0,.15); min-width:0; } + .route-study-metric b { display:block; color:var(--text); font-size:12.5px; overflow-wrap:anywhere; } + .route-study-metric span { display:block; margin-top:2px; color:var(--muted); font-size:9px; line-height:1.35; } + .route-study-dist { margin-top:8px; display:flex; flex-wrap:wrap; gap:5px; } + .route-study-chip { padding:5px 7px; border:1px solid rgba(56,189,248,.22); border-radius:999px; + background:rgba(56,189,248,.06); color:#bae6fd; font-size:9px; font-weight:800; } + .route-study-chip[data-state="unresolved"] { border-color:rgba(251,113,133,.3); background:rgba(251,113,133,.07); color:#fecdd3; } + .route-study-chip[data-state="declared_unvalidated"] { border-color:rgba(251,191,36,.3); background:rgba(251,191,36,.07); color:#fde68a; } + .route-study-chip[data-state="drafted_unsealed"],.route-study-chip[data-state="drafted_unvalidated"] { + border-color:rgba(167,139,250,.3); background:rgba(167,139,250,.07); color:#ddd6fe; + } + .route-study-note { margin-top:7px; color:var(--muted); font-size:9.5px; line-height:1.45; } + .route-study-campaign { margin-top:10px; padding-top:10px; border-top:1px solid rgba(255,255,255,.08); } + .route-study-campaign-head { display:flex; align-items:center; justify-content:space-between; gap:8px; + color:var(--text); font-size:10px; font-weight:900; } + .route-study-actions { margin-top:7px; display:flex; flex-wrap:wrap; gap:6px; align-items:center; } + .route-study-file-label { display:inline-flex; align-items:center; border:1px solid rgba(45,212,191,.28); + border-radius:8px; background:rgba(8,12,25,.75); color:#ccfbf1; + padding:6px 8px; font-size:10.5px; font-weight:900; cursor:pointer; } + .route-study-file-label input { position:absolute; width:1px; height:1px; overflow:hidden; + clip:rect(0,0,0,0); white-space:nowrap; } + .route-study-inventory { margin-top:7px; display:flex; flex-wrap:wrap; gap:5px; } + .route-study-inventory button { text-align:left; overflow-wrap:anywhere; } /* ── Composer ───────────────────────────────────────────────────── */ .compose-wrap { padding:0 14% 32px; flex-shrink: 0; z-index: 20; @@ -385,6 +560,12 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: .panel { background:rgba(15, 23, 42, 0.4); border-left:1px solid var(--border); backdrop-filter: var(--glass); display:flex; flex-direction:column; overflow:hidden; } .panel-tabs { display:flex; border-bottom:1px solid var(--border); background: rgba(0,0,0,0.1); } + .panel-close { display:none; flex:0 0 48px; align-items:center; justify-content:center; + color:var(--muted); border-left:1px solid var(--border); font-size:22px; } + .panel-close:hover { color:var(--text); background:rgba(255,255,255,.06); } + .panel-backdrop { display:none; position:fixed; inset:0; z-index:70; padding:0; + background:rgba(2,6,23,.68); backdrop-filter:blur(4px); opacity:0; + pointer-events:none; transition:opacity .25s ease; } .ptab { flex:1; padding:18px 8px; font-size:11px; font-weight:800; text-align:center; text-transform:uppercase; letter-spacing:.14em; color:var(--muted); cursor:pointer; border:none; @@ -421,6 +602,8 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: .mode-card.on { background:rgba(56,189,248,.08); border-color:rgba(56,189,248,.5); box-shadow:0 8px 24px rgba(0,0,0,0.2), inset 0 0 20px rgba(56,189,248,0.03); } + .mode-card.on[data-mode="auto"] { background:rgba(129,140,248,.08); + border-color:rgba(129,140,248,.5); } .mode-card.on[data-mode="collective"] { background:rgba(45,212,191,.08); border-color:rgba(45,212,191,.5); } .mode-card.on[data-mode="loop"] { background:rgba(245,158,11,.08); @@ -428,6 +611,7 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: .mc-title { font-size:14.5px; font-weight:800; margin-bottom:6px; display:flex; align-items:center; gap:10px; } .mc-dot { width:9px; height:9px; border-radius:50%; flex-shrink:0; } + .mc-dot.auto { background:#818cf8; box-shadow:0 0 8px #818cf8; } .mc-dot.std { background:var(--green); box-shadow:0 0 8px var(--green); } .mc-dot.col { background:var(--teal); box-shadow:0 0 8px var(--teal); } .mc-dot.loop { background:var(--amber); box-shadow:0 0 8px var(--amber); } @@ -463,14 +647,26 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: @media (max-width: 1100px) { .shell { grid-template-columns:64px 1fr; } - .panel { position:fixed; right:0; top:0; bottom:0; width:min(390px,92vw); z-index:80; } + .panel-toggle, .panel-close { display:inline-flex; } + .panel-backdrop { display:block; } + .panel { position:fixed; right:0; top:0; bottom:0; width:min(390px,92vw); z-index:80; + background:rgba(8,15,30,.97); box-shadow:-24px 0 70px rgba(0,0,0,.48); + transform:translateX(105%); opacity:0; visibility:hidden; pointer-events:none; + transition:transform .28s cubic-bezier(.16,1,.3,1), opacity .2s ease, visibility 0s linear .28s; } + .panel.is-open { transform:none; opacity:1; visibility:visible; pointer-events:auto; + transition-delay:0s; } + .shell.panel-open .panel-backdrop { opacity:1; pointer-events:auto; } .thread, .compose-wrap { padding-left:7%; padding-right:7%; } } @media (max-width: 760px) { .shell { grid-template-columns:1fr; } .rail { display:none; } - .panel { display:none; } .wk-header { padding:0 18px; } + .wk-header > div:first-child { min-width:0; gap:10px !important; } + .wk-title { font-size:17px; white-space:nowrap; overflow:hidden; text-overflow:ellipsis; } + .model-pill { display:none; } + .panel-toggle span { display:none; } + .panel-toggle { min-width:40px; padding:8px; justify-content:center; } .thread { padding:24px 18px; gap:22px; } .compose-wrap { padding:0 18px 18px; } .msg { max-width:100%; } @@ -521,7 +717,14 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]:
V46 Champion
-
+
+ +
@@ -569,6 +772,11 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: +
+ @@ -620,6 +829,42 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: +
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+ + + + + +
+ +
+
0 / 0exact usage-feedback joins
+
-candidate action agreement
+
-matched observed approval
+
-matched cost / latency
+
0 / 0 / 0completed / failed / in flight
+
-terminal feedback coverage
+
ESS 0.0 / 20target-policy overlap
+
-weakest target action
+
0 / 12readiness checks passed
+
0 / 0quality outcomes observed
+
0 / 0routes with precommitted outcome set
+
waitingevidence source
+
+
Shadow only - waiting for joined evidence.
+
+
Readiness matrix
+
+
Waiting for durable evidence.
+
+
Associational replay only. Changed actions receive no imputed reward.
+ +
+
+
+
Bounded Exposure Rehearsal
+
Rehearsal only - execution off
+
+ +
+
+ + + + + + + + + + + + +
+
+ Preview a prompt to rehearse exact post-filter propensities. No route will run and no evidence will be written. +
+
+
-incumbent route
+
-minimum alternate propensity
+
-routes for target on every alternate
+
-same-stratum expected cost units
+
-latency-tier envelope
+
-draft charter fingerprint
+
-stateful design screen
+
-frozen target-policy draft
+
-independent-review state
+
-protocol draft fingerprint
+
+
+
+
+
+ Multi-stratum semantic review + 0 strata +
+
+ + + + + +
+
+
+
not builtsemantic verification
+
-review bundle fingerprint
+
+
+ Inventory only. Strata are never pooled, weighted, assigned, executed, or promoted. +
+
+
+
+ Shadow assignment registry - read only + +
+
+
0sealed campaigns
+
not loadedappend-only event chain
+
0 / 0verified / committed assignments
+
not initializedcampaign states
+
+
+
+ Refresh to inspect the isolated local registry. This browser cannot seal, assign, reveal, activate, or promote a route policy. +
+
+
+ Same-support rehearsal only. ESS, policy value, live assignment, and promotion remain unavailable. +
+
@@ -676,6 +1110,7 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: +
@@ -691,11 +1126,15 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: const sessionId = ([1e7]+-1e3+-4e3+-8e2+-1e11).replace(/[018]/g, c => (c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> c/4).toString(16)); - let agentMode = 'off'; + let agentMode = 'auto'; let currentUpload = null; let currentUpUrl = ''; let loopStep = 0; let catalogByKey = {}; + let lastRouteFeedback = null; + let latestRouteStudy = null; + let routeStudyStrata = []; + let latestRouteReviewBundle = null; async function api(path, body=null) { const opts = body @@ -732,6 +1171,32 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: } } + const compactPanelMedia = window.matchMedia('(max-width: 1100px)'); + + function setPanelOpen(open, options = {}) { + const panel = el('controlPanel'); + const toggle = el('panelToggle'); + const compact = compactPanelMedia.matches; + const nextOpen = compact && Boolean(open); + el('shell').classList.toggle('panel-open', nextOpen); + panel.classList.toggle('is-open', nextOpen); + panel.setAttribute('aria-hidden', compact && !nextOpen ? 'true' : 'false'); + panel.inert = compact && !nextOpen; + toggle.setAttribute('aria-expanded', String(nextOpen)); + toggle.setAttribute('aria-label', nextOpen ? 'Close control panel' : 'Open control panel'); + if (nextOpen && options.focus !== false) { + const activeTab = panel.querySelector('.ptab.on') || panel.querySelector('.ptab'); + if (activeTab) activeTab.focus(); + } else if (!nextOpen && options.restoreFocus) { + toggle.focus(); + } + } + + function openPanelTab(name) { + switchPtab(name); + setPanelOpen(true); + } + // ── Tabs (rail) ───────────────────────────────────────────────────── qsa('.rail-item[data-tab]').forEach(btn => { btn.onclick = () => { @@ -740,7 +1205,11 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: const ptab = btn.dataset.tab === 'bench' ? 'bench' : btn.dataset.tab === 'settings' ? 'model' : 'model'; - switchPtab(ptab); + if (btn.dataset.tab === 'bench' || btn.dataset.tab === 'settings') openPanelTab(ptab); + else { + switchPtab(ptab); + setPanelOpen(false); + } }; }); @@ -749,6 +1218,17 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: btn.onclick = () => switchPtab(btn.dataset.ptab); }); + el('panelToggle').onclick = () => setPanelOpen(!el('controlPanel').classList.contains('is-open')); + el('panelClose').onclick = () => setPanelOpen(false, { restoreFocus: true }); + el('panelBackdrop').onclick = () => setPanelOpen(false, { restoreFocus: true }); + document.addEventListener('keydown', event => { + if (event.key === 'Escape' && el('controlPanel').classList.contains('is-open')) { + setPanelOpen(false, { restoreFocus: true }); + } + }); + compactPanelMedia.addEventListener('change', () => setPanelOpen(false, { focus: false })); + setPanelOpen(false, { focus: false }); + function switchPtab(name) { qsa('.ptab').forEach(b => b.classList.toggle('on', b.dataset.ptab === name)); ['model','mode','bench'].forEach(t => @@ -762,7 +1242,7 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: qsa('.mode-card').forEach(x => x.classList.remove('on')); c.classList.add('on'); agentMode = c.dataset.mode; - el('loopPanel').style.display = agentMode==='loop' ? 'block' : 'none'; + el('loopPanel').style.display = ['auto','loop','collective_loop'].includes(agentMode) ? 'block' : 'none'; el('loopSteps').innerHTML = ''; loopStep = 0; updateModePill(); @@ -772,14 +1252,16 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: function updateModePill() { const pill = el('modePill'); - const labels = { off:'Standard', collective:'Collective', loop:'Autonomous' }; + const labels = { auto:'Auto', off:'Standard', collective:'Collective', loop:'Autonomous', collective_loop:'Collective Loop' }; if (agentMode === 'off') { pill.style.display='none'; return; } pill.style.display='block'; pill.textContent = labels[agentMode] || agentMode; - pill.style.color = agentMode==='collective' ? 'var(--teal)' : 'var(--amber)'; - pill.style.borderColor = agentMode==='collective' + const isCollective = agentMode==='collective' || agentMode==='collective_loop'; + const isAuto = agentMode==='auto'; + pill.style.color = isAuto ? '#818cf8' : isCollective ? 'var(--teal)' : 'var(--amber)'; + pill.style.borderColor = isAuto ? 'rgba(129,140,248,.4)' : isCollective ? 'rgba(45,212,191,.4)' : 'rgba(245,158,11,.4)'; - pill.style.background = agentMode==='collective' + pill.style.background = isAuto ? 'rgba(129,140,248,.08)' : isCollective ? 'rgba(45,212,191,.08)' : 'rgba(245,158,11,.08)'; } @@ -836,6 +1318,110 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: .replace(/>/g,'>').replace(/"/g,'"'); } + function scorePct(v) { + const n = Number(v); + if (!Number.isFinite(n)) return null; + const bounded = Math.max(0, Math.min(1, n)); + return Math.round(bounded * 100); + } + + function stopReasonLabel(code) { + const labels = { + reviewer_complete: 'Reviewer complete', + score_threshold: 'Score threshold', + budget_exhausted: 'Budget exhausted', + reviewer_continue: 'Reviewer continue', + score_below_threshold: 'Score below threshold' + }; + return labels[code] || String(code || '').replace(/_/g, ' '); + } + + function autoPolicyPills(policy) { + if (!policy) return []; + const selected = policy.selected_agent_mode || policy.resolved_agent_mode || 'off'; + const bits = [`Auto ${escHtml(selected)}`]; + if (policy.budget_profile) bits.push(`Budget ${escHtml(policy.budget_profile)}`); + if (policy.session_budget) { + const b = policy.session_budget; + bits.push(`Session ${escHtml(b.remaining_cost_units)} / ${escHtml(b.limit_cost_units)}`); + } + if (policy.score !== undefined) bits.push(`Difficulty ${escHtml(policy.score)}`); + if (policy.score_before_budget !== undefined && policy.score_before_budget !== policy.score) { + bits.push(`Score ${escHtml(policy.score_before_budget)} -> ${escHtml(policy.score)}`); + } + if (policy.reason) bits.push(`${escHtml(stopReasonLabel(policy.reason))}`); + if (Array.isArray(policy.reasons) && policy.reasons.length) { + bits.push(`${escHtml(policy.reasons.slice(0,3).join(', '))}`); + } + if (policy.feedback_adjustment) { + const adj = policy.feedback_adjustment; + const label = adj.reason === 'recent_weighted_feedback_regression' + ? 'Adaptive' + : (adj.reason === 'adaptive_quality_cost_preferred_neighbor' ? 'Pareto' : 'Feedback'); + bits.push(`${label} ${escHtml(adj.from || '')} -> ${escHtml(adj.to || '')}`); + } else if (policy.feedback_summary && policy.feedback_summary.total_feedback) { + bits.push(`Feedback ${escHtml(policy.feedback_summary.total_feedback)}`); + } + if (policy.uncertainty_adjustment) { + const adj = policy.uncertainty_adjustment; + bits.push(`Uncertain ${escHtml(adj.from || '')} -> ${escHtml(adj.to || '')}`); + } + if (policy.session_budget_adjustment) { + const adj = policy.session_budget_adjustment; + bits.push(`Paced ${escHtml(adj.from || '')} -> ${escHtml(adj.to || '')}`); + } + return bits; + } + + function routeEconomicsPills(economics) { + if (!economics) return []; + const estimate = economics.estimate || {}; + const actual = economics.actual || {}; + const bits = []; + if (estimate.estimated_cost_units !== undefined) { + bits.push(`Cost ~${escHtml(estimate.estimated_cost_units)}`); + } + if (estimate.estimated_model_calls !== undefined) { + bits.push(`Planned calls ${escHtml(estimate.estimated_model_calls)}`); + } + if (actual.elapsed_ms !== undefined) { + bits.push(`Elapsed ${escHtml(actual.elapsed_ms)}ms`); + } + if (actual.model_calls !== undefined) { + bits.push(`Calls ${escHtml(actual.model_calls)}`); + } + return bits; + } + + function computePills(compute) { + if (!compute || !Object.keys(compute).length) return []; + const bits = []; + const requested = compute.requested_reasoning_cycles ?? compute.selected_reasoning_cycles; + const used = compute.cycles_used; + if (requested !== undefined && requested !== null) { + bits.push(`Compute ${escHtml(requested)} requested`); + } + if (used !== undefined && used !== null) { + bits.push(`${escHtml(used)} cycles used`); + } + if (compute.reasoning_budget_mode) { + bits.push(`${escHtml(compute.reasoning_budget_mode)} budget`); + } + if (compute.adaptive_compute !== undefined) { + bits.push(`Adaptive ${compute.adaptive_compute ? 'on' : 'off'}`); + } + if (compute.exit_reason) { + bits.push(`Exit ${escHtml(stopReasonLabel(compute.exit_reason))}`); + } + if (compute.prediction_confidence_delta !== undefined && compute.prediction_confidence_delta !== null) { + bits.push(`Prediction drift ${escHtml(compute.prediction_confidence_delta)}`); + } + if (compute.applied === false && compute.supported === false) { + bits.push('Compute controls unsupported'); + } + return bits; + } + function buildTrace(trace) { const wrapper = document.createElement('div'); wrapper.className = 'trace'; @@ -847,13 +1433,36 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: hdr.innerHTML = ` Autonomous Logic Chain — ${trace.loop_steps.length} cycles`; wrapper.appendChild(hdr); + const summaryBits = autoPolicyPills(trace.auto_agent_policy); + summaryBits.push(...routeEconomicsPills(trace.route_economics)); + summaryBits.push(...computePills(trace.compute)); + const stopScore = scorePct(trace.loop_stop_score); + if (trace.loop_stop_reason_code) summaryBits.push(`${escHtml(stopReasonLabel(trace.loop_stop_reason_code))}`); + if (stopScore != null) summaryBits.push(`Stop score ${stopScore}%`); + if (trace.loop_budget != null) summaryBits.push(`Budget ${escHtml(trace.loop_steps.length)}/${escHtml(trace.loop_budget)}`); + if (trace.loop_completion_reason) summaryBits.push(`${escHtml(String(trace.loop_completion_reason).slice(0,160))}`); + if (summaryBits.length) { + const summary = document.createElement('div'); + summary.className = 'trace-body trace-summary'; + summary.innerHTML = summaryBits.join(''); + wrapper.appendChild(summary); + } + const body = document.createElement('div'); body.className = 'trace-body'; trace.loop_steps.forEach(s => { + const stepScore = scorePct(s.review_score ?? s.loop_score); + const scoreHtml = stepScore == null ? '' : `Score ${stepScore}%`; + const stopHtml = s.stop_decision === 'stop' + ? `${escHtml(stopReasonLabel(s.stop_reason_code))}` + : ''; + const note = s.review_note || s.completion_evidence || s.next_step || ''; body.innerHTML += `
${s.step}
${escHtml(s.goal||'Strategy Initialization')}
- ${escHtml((s.worker_excerpt||'').slice(0,140))}...
+ ${escHtml((s.worker_excerpt||'').slice(0,140))}... +
${scoreHtml}${stopHtml}
+ ${note ? `
${escHtml(String(note).slice(0,160))}
` : ''}
`; }); wrapper.appendChild(body); @@ -883,7 +1492,41 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: const body = document.createElement('div'); body.className = 'trace-body'; body.style.fontSize = '12px'; - body.innerHTML = 'Expert weights synthesized from: ' + trace.consulted_models.join(', '); + const policyBits = autoPolicyPills(trace.auto_agent_policy) + .concat(routeEconomicsPills(trace.route_economics), computePills(trace.compute)); + body.innerHTML = (policyBits.length ? `
${policyBits.join('')}
` : '') + + 'Expert weights synthesized from: ' + trace.consulted_models.join(', '); + wrapper.appendChild(body); + } else if (trace.auto_agent_policy) { + const hdr = document.createElement('div'); + hdr.className = 'trace-hdr'; + hdr.style.color = '#818cf8'; + hdr.textContent = 'Adaptive Router'; + wrapper.appendChild(hdr); + const body = document.createElement('div'); + body.className = 'trace-body trace-summary'; + body.innerHTML = autoPolicyPills(trace.auto_agent_policy) + .concat(routeEconomicsPills(trace.route_economics), computePills(trace.compute)).join(''); + wrapper.appendChild(body); + } else if (trace.route_economics) { + const hdr = document.createElement('div'); + hdr.className = 'trace-hdr'; + hdr.style.color = 'var(--blue)'; + hdr.textContent = 'Route Economics'; + wrapper.appendChild(hdr); + const body = document.createElement('div'); + body.className = 'trace-body trace-summary'; + body.innerHTML = routeEconomicsPills(trace.route_economics).concat(computePills(trace.compute)).join(''); + wrapper.appendChild(body); + } else if (trace.compute && Object.keys(trace.compute).length) { + const hdr = document.createElement('div'); + hdr.className = 'trace-hdr'; + hdr.style.color = 'var(--teal)'; + hdr.textContent = 'Adaptive Compute'; + wrapper.appendChild(hdr); + const body = document.createElement('div'); + body.className = 'trace-body trace-summary'; + body.innerHTML = computePills(trace.compute).join(''); wrapper.appendChild(body); } @@ -916,7 +1559,788 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: }); } + function setRouteFeedbackVisible(show, label='Route') { + const box = el('routeFeedback'); + if (!box) return; + box.style.display = show ? 'flex' : 'none'; + const labelEl = el('routeFeedbackLabel'); + if (labelEl) labelEl.textContent = label; + ['routeGoodBtn','routeBadBtn','routeDeeperBtn','routeCostBtn','routeSlowBtn'].forEach(id => { + const btn = el(id); + if (btn) btn.disabled = false; + }); + } + + function routeMetricText(value, digits=1) { + const n = Number(value); + if (!Number.isFinite(n)) return '-'; + return n.toFixed(digits).replace(/\.0$/, ''); + } + + function routeQualityText(adaptive) { + if (!adaptive) return 'Quality -'; + const q = Number(adaptive.quality_score); + if (!Number.isFinite(q)) return 'Quality -'; + const suffix = adaptive.regression_signal ? ' risk' : ''; + return `Quality ${Math.round(Math.max(0, Math.min(1, q)) * 100)}%${suffix}`; + } + + function routeConfidenceText(adaptive) { + if (!adaptive) return 'Recent evidence -'; + const lowerRaw = adaptive.quality_lower_bound; + const upperRaw = adaptive.quality_upper_bound; + if (lowerRaw === null || lowerRaw === undefined || upperRaw === null || upperRaw === undefined) return 'Recent evidence -'; + const lower = Number(lowerRaw); + const upper = Number(upperRaw); + if (!Number.isFinite(lower) || !Number.isFinite(upper)) return 'Recent evidence -'; + const status = adaptive.confidence_status ? ` ${adaptive.confidence_status}` : ''; + return `Heuristic ${Math.round(Math.max(0, lower) * 100)}-${Math.round(Math.min(1, upper) * 100)}%${status}`; + } + + function routePreferenceText(adaptive) { + if (!adaptive) return 'Preference -'; + if (adaptive.preference_direction === 'deeper') return 'Preference deeper'; + if (adaptive.preference_direction === 'shallower') { + const cost = Number(adaptive.weighted_cost_pressure) || 0; + const latency = Number(adaptive.weighted_latency_pressure) || 0; + return latency > cost ? 'Preference faster' : 'Preference cheaper'; + } + return 'Preference neutral'; + } + + function renderRouteHealth(summary) { + const box = el('routeHealth'); + if (!box) return; + const usage = summary && summary.route_usage ? summary.route_usage : {}; + const usageEconomics = usage && usage.economics ? usage.economics : {}; + const feedbackEconomics = summary && summary.economics ? summary.economics : {}; + const economics = usageEconomics.sample_count ? usageEconomics : feedbackEconomics; + const adaptive = summary && summary.adaptive ? summary.adaptive : {}; + const total = summary && summary.total_feedback !== undefined ? Number(summary.total_feedback) : NaN; + const routeTotal = usage && usage.total_routes !== undefined ? Number(usage.total_routes) : NaN; + const samples = economics.sample_count ? Number(economics.sample_count) : 0; + const shownTotal = Number.isFinite(routeTotal) && routeTotal ? routeTotal : Number.isFinite(total) ? total : samples; + if (!shownTotal && !samples) { + box.style.display = 'none'; + return; + } + box.style.display = 'flex'; + el('routeHealthCount').textContent = `Routes ${shownTotal}`; + el('routeHealthQuality').textContent = routeQualityText(adaptive); + el('routeHealthConfidence').textContent = routeConfidenceText(adaptive); + el('routeHealthPreference').textContent = routePreferenceText(adaptive); + el('routeHealthCost').textContent = `Avg cost ${routeMetricText(economics.avg_cost_units, 2)}`; + el('routeHealthLatency').textContent = `Avg ms ${routeMetricText(economics.avg_elapsed_ms, 0)}`; + } + + async function refreshRouteHealth() { + try { + const result = await api('/api/route_health', { session_id: sessionId }); + renderRouteHealth(result.route_health || result.summary || {}); + } catch (_) {} + } + + function policyLabPercent(value) { + if (value === null || value === undefined || value === '') return '-'; + const n = Number(value); + return Number.isFinite(n) ? `${Math.round(Math.max(0, Math.min(1, n)) * 100)}%` : '-'; + } + + function policyLabLabel(value) { + return String(value || '').replace(/_/g, ' ').trim(); + } + + const policyLabCheckLabels = { + candidate_delta_present: 'Candidate changes at least one action', + population_integrity_complete: 'Durable population is unique and evaluable', + execution_integrity_complete: 'Execution state and chosen action reconcile', + logging_integrity_complete: 'Logging metadata and decision fingerprint verify', + minimum_overlap_routes_met: 'Minimum overlap route count', + target_probability_floor_met: 'Target probability floor', + global_overlap_ess_met: 'Global overlap ESS floor', + per_action_overlap_met: 'Per-action overlap ESS floor', + outcome_evidence_integrity: 'Outcome contracts and eligible quality evidence verify', + quality_observation_ready: 'Quality observation process', + durable_lifecycle_present: 'Durable lifecycle evidence', + lifecycle_reconciled: 'Lifecycle fully reconciled', + }; + + function renderPolicyLabChecks(checks, blockers) { + const container = el('policyLabChecks'); + container.replaceChildren(); + const knownChecks = checks && typeof checks === 'object' ? checks : {}; + Object.entries(policyLabCheckLabels).forEach(([key, label]) => { + const item = document.createElement('div'); + const known = Object.prototype.hasOwnProperty.call(knownChecks, key); + item.className = 'policy-lab-check'; + item.dataset.state = known ? (knownChecks[key] ? 'pass' : 'fail') : 'unknown'; + item.textContent = label; + container.appendChild(item); + }); + const blockerList = Array.isArray(blockers) ? blockers.filter(Boolean) : []; + const blockerBox = el('policyLabBlockers'); + if (blockerList.length) { + blockerBox.textContent = `Blockers: ${blockerList.map(policyLabLabel).join(' | ')}`; + } else if (Object.keys(knownChecks).length) { + blockerBox.textContent = 'No readiness blockers; validated external OPE is still required.'; + } else { + blockerBox.textContent = 'Blockers unavailable until replay succeeds.'; + } + } + + function renderPolicyLabUnavailable(message) { + const defaults = { + policyLabJoined: '0 / 0', policyLabAgreement: '-', policyLabApproval: '-', + policyLabEconomics: '-', policyLabLifecycle: '0 / 0 / 0', policyLabFeedbackCoverage: '-', + policyLabOverlapEss: 'ESS - / -', policyLabWeakestAction: '-', + policyLabReadinessChecks: `0 / ${Object.keys(policyLabCheckLabels).length}`, + policyLabOutcomeCoverage: '0 / 0', + policyLabContractCoverage: '0 / 0', + policyLabEvidenceSource: 'unavailable', + }; + Object.entries(defaults).forEach(([id, value]) => { el(id).textContent = value; }); + el('policyLabGate').textContent = `Policy replay unavailable: ${message}`; + el('policyLabNote').textContent = 'No Policy Lab metrics are current. Shadow-only gating remains in force.'; + renderPolicyLabChecks({}, []); + } + + function renderPolicyLab(report) { + const support = report && report.support ? report.support : {}; + const usage = support.usage || {}; + const agreement = report && report.candidate_action_agreement ? report.candidate_action_agreement : {}; + const matched = report && report.matched_observed ? report.matched_observed : {}; + const gate = report && report.promotion_gate ? report.promotion_gate : {}; + const propensity = report && report.propensity_readiness ? report.propensity_readiness : {}; + const readiness = report && report.evaluation_readiness ? report.evaluation_readiness : {}; + const thresholds = readiness.thresholds || {}; + const overlap = readiness.target_overlap || {}; + const outcome = readiness.outcome_observation || {}; + const maturity = report && report.outcome_contract_maturity && typeof report.outcome_contract_maturity === 'object' + ? report.outcome_contract_maturity + : {}; + const maturityByOutcome = maturity.by_outcome && typeof maturity.by_outcome === 'object' + ? maturity.by_outcome + : {}; + const qualityMaturity = maturityByOutcome.user_quality_rating && typeof maturityByOutcome.user_quality_rating === 'object' + ? maturityByOutcome.user_quality_rating + : {}; + const durable = report && report.durable_ledger ? report.durable_ledger : {}; + const lifecycle = durable.counts || {}; + const feedbackCoverage = durable.feedback_coverage || {}; + const evidenceSource = String(report && report.evidence_source ? report.evidence_source : 'unknown'); + const durableEvidence = evidenceSource === 'durable_sqlite_ledger'; + el('policyLabJoined').textContent = `${Number(support.exact_joined_route_ids) || 0} / ${Number(usage.unique_route_ids) || 0}`; + el('policyLabAgreement').textContent = policyLabPercent(agreement.agreement_rate); + const approval = policyLabPercent(matched.approval_rate); + el('policyLabApproval').textContent = matched.quality_sample_count ? `${approval} (n=${matched.quality_sample_count})` : '-'; + const hasCost = matched.avg_cost_units !== null && matched.avg_cost_units !== undefined && matched.avg_cost_units !== ''; + const hasLatency = matched.avg_elapsed_ms !== null && matched.avg_elapsed_ms !== undefined && matched.avg_elapsed_ms !== ''; + const cost = hasCost ? Number(matched.avg_cost_units) : NaN; + const latency = hasLatency ? Number(matched.avg_elapsed_ms) : NaN; + el('policyLabEconomics').textContent = Number.isFinite(cost) || Number.isFinite(latency) + ? `${Number.isFinite(cost) ? cost.toFixed(2) : '-'} / ${Number.isFinite(latency) ? Math.round(latency) + ' ms' : '-'}` + : '-'; + el('policyLabLifecycle').textContent = + `${Number(lifecycle.completed) || 0} / ${Number(lifecycle.failed) || 0} / ${Number(lifecycle.inflight) || 0}`; + el('policyLabFeedbackCoverage').textContent = policyLabPercent(feedbackCoverage.terminal_coverage_rate); + const overlapEss = Number(overlap.effective_sample_size) || 0; + const globalEssFloor = Number(thresholds.minimum_global_effective_sample_size) || 20; + const actionEssFloor = Number(thresholds.minimum_per_action_effective_sample_size) || 10; + const weakestAction = policyLabLabel(overlap.weakest_target_action || 'none'); + const weakestEss = Number(overlap.weakest_action_effective_sample_size) || 0; + el('policyLabOverlapEss').textContent = `ESS ${overlapEss.toFixed(1)} / ${globalEssFloor.toFixed(0)}`; + el('policyLabWeakestAction').textContent = overlap.weakest_target_action + ? `${weakestAction} ${weakestEss.toFixed(1)} / ${actionEssFloor.toFixed(0)}` + : '-'; + const totalChecks = Number(gate.total_checks) || Object.keys(policyLabCheckLabels).length; + el('policyLabReadinessChecks').textContent = `${Number(gate.passed_checks) || 0} / ${totalChecks}`; + el('policyLabOutcomeCoverage').textContent = + `${Number(outcome.quality_observed_routes) || 0} / ${Number(outcome.evaluable_routes) || 0}`; + el('policyLabContractCoverage').textContent = + `${Number(maturity.precommitted_routes) || 0} / ${Number(maturity.included_routes) || 0}`; + el('policyLabEvidenceSource').textContent = durableEvidence ? 'durable SQLite' : policyLabLabel(evidenceSource); + const gateStatus = policyLabLabel(gate.status || 'blocked'); + const gateReason = policyLabLabel(gate.reason_code || 'no evidence'); + const blockers = Array.isArray(gate.blocking_reason_codes) ? gate.blocking_reason_codes : []; + const moreBlockers = Math.max(0, blockers.length - 1); + renderPolicyLabChecks(gate.checks || {}, blockers); + el('policyLabGate').textContent = + `${gateStatus} · ${Number(gate.passed_checks) || 0}/${totalChecks} checks · ${gateReason}` + + `${moreBlockers ? ` + ${moreBlockers} more` : ''} · ${policyLabLabel(gate.deployment || 'shadow only')}`; + const valid = Number(propensity.valid_routes) || 0; + const checked = Number(propensity.checked_evaluable_usage_routes) || 0; + const started = Number(lifecycle.started) || 0; + const minTarget = overlap.minimum_target_probability === null || overlap.minimum_target_probability === undefined + ? '-' + : Number(overlap.minimum_target_probability).toFixed(3); + el('policyLabNote').textContent = + `${durableEvidence ? 'Durable ledger evidence.' : 'Nondurable compatibility evidence; readiness is blocked.'} ` + + `Durable routes ${started}; target overlap ESS ${overlapEss.toFixed(1)}/${globalEssFloor.toFixed(0)}; ` + + `minimum target propensity ${minTarget}; quality observed ${Number(outcome.quality_observed_routes) || 0}/${Number(outcome.evaluable_routes) || 0}. ` + + `Outcome contracts ${Number(maturity.precommitted_routes) || 0}/${Number(maturity.included_routes) || 0} precommitted; ` + + `quality events ${Number(qualityMaturity.observed_event_count) || 0}, mature contracts ${Number(qualityMaturity.mature_contract_count) || 0}. ` + + `Missing feedback stays unknown; maturity is diagnostic only; ${valid}/${checked} usage rows validate logging metadata. No policy value was estimated.`; + } + + async function refreshPolicyLab() { + const button = el('policyLabRefresh'); + if (button) button.disabled = true; + try { + const result = await api('/api/route_policy_lab', { + session_id: sessionId, + profile: el('policyLabProfile').value, + }); + renderPolicyLab(result.policy_lab || {}); + } catch (err) { + renderPolicyLabUnavailable(err.message); + } finally { + if (button) button.disabled = false; + } + } + + function routeStudyNumber(value, digits = 0) { + const number = Number(value); + if (!Number.isFinite(number)) return '-'; + return number.toLocaleString(undefined, { + minimumFractionDigits: digits, + maximumFractionDigits: digits, + }); + } + + function renderRouteStudyUnavailable(message) { + latestRouteStudy = null; + el('routeBundleAdd').disabled = true; + const defaults = { + routeStudyBaseline: '-', routeStudyFloor: '-', routeStudyTraffic: '-', + routeStudyCost: '-', routeStudyLatency: '-', routeStudyCharter: '-', + routeProtocolMode: '-', routeProtocolPolicy: '-', routeProtocolReview: '-', + routeProtocolHash: '-', + }; + Object.entries(defaults).forEach(([id, value]) => { el(id).textContent = value; }); + el('routeStudyDistribution').replaceChildren(); + el('routeProtocolBlockers').replaceChildren(); + el('routeStudyStatus').textContent = message; + el('routeStudyNote').textContent = + 'No study probabilities are current. Execution, evidence writes, policy value, and promotion remain unavailable.'; + } + + function renderRouteShadowRegistry(snapshot) { + const campaigns = Array.isArray(snapshot && snapshot.campaigns) ? snapshot.campaigns : []; + const chain = snapshot && snapshot.event_chain ? snapshot.event_chain : null; + const available = snapshot && snapshot.available === true; + const committed = campaigns.reduce((total, row) => total + (Number(row.commitment_count) || 0), 0); + const matched = campaigns.reduce((total, row) => total + (Number(row.matched_assignment_count) || 0), 0); + const processed = campaigns.reduce((total, row) => total + (Number(row.processed_reveal_count) || 0), 0); + const mismatched = campaigns.reduce((total, row) => total + (Number(row.mismatched_assignment_count) || 0), 0); + const states = [...new Set(campaigns.map(row => policyLabLabel((row || {}).state || 'unknown')))]; + el('routeShadowRegistryCount').textContent = routeStudyNumber(snapshot && snapshot.campaign_count, 0); + el('routeShadowRegistryChain').textContent = chain + ? `${chain.ok ? 'verified' : 'failed'} - ${routeStudyNumber(chain.verified_events, 0)} events` + : 'not initialized'; + el('routeShadowRegistryAssignments').textContent = `${routeStudyNumber(matched, 0)} / ${routeStudyNumber(committed, 0)}`; + el('routeShadowRegistryState').textContent = states.length ? states.join(', ') : 'not initialized'; + + const inventory = el('routeShadowRegistryCampaigns'); + inventory.replaceChildren(); + campaigns.forEach(row => { + const chip = document.createElement('span'); + chip.className = 'route-study-chip'; + chip.dataset.state = Number(row.mismatched_assignment_count) > 0 ? 'unresolved' : ''; + chip.textContent = `${row.campaign_id || 'campaign'} - ${policyLabLabel(row.state || 'unknown')} - ${routeStudyNumber(row.matched_assignment_count, 0)}/${routeStudyNumber(row.commitment_count, 0)} matched`; + inventory.appendChild(chip); + }); + + if (!available) { + el('routeShadowRegistryStatus').textContent = + `No shadow registry exists at ${snapshot && snapshot.registry_location ? snapshot.registry_location : 'the canonical memory path'}. ` + + 'Browser access is read-only; execution, activation, and automatic promotion remain unavailable.'; + return; + } + el('routeShadowRegistryStatus').textContent = + `${snapshot.ok ? 'Registry verification passed.' : 'Registry verification failed.'} ` + + `${campaigns.length} campaign(s), ${committed} opaque commitment(s), ${matched} matched assignment(s), ${processed} processed reveal(s), ${mismatched} mismatch(es). ` + + 'Local chain verification is not an external transparency anchor. Browser access is read-only; execution, activation, and promotion remain unavailable.'; + } + + async function refreshRouteShadowRegistry() { + const button = el('routeShadowRegistryRefresh'); + button.disabled = true; + el('routeShadowRegistryStatus').textContent = 'Reading and verifying the local shadow registry...'; + try { + const result = await api('/api/route_shadow_registry/status'); + renderRouteShadowRegistry(result.route_shadow_registry || {}); + } catch (err) { + renderRouteShadowRegistry({available:false, campaign_count:0, campaigns:[]}); + el('routeShadowRegistryStatus').textContent = `Registry status error: ${err.message}`; + } finally { + button.disabled = false; + } + } + + function renderRouteStudy(payload) { + const study = payload && payload.route_study ? payload.route_study : {}; + const charter = study.charter || {}; + const enrollment = charter.enrollment || {}; + const design = charter.probability_design || {}; + const traffic = charter.traffic_scenario || {}; + const labelScenario = traffic.observed_label_scenario || {}; + const exact = labelScenario.exact_simultaneous_target || {}; + const resources = charter.resource_forecast || {}; + const boundaries = charter.causal_boundaries || {}; + const protocol = payload && payload.route_protocol_preflight + ? payload.route_protocol_preflight : null; + const protocolCharter = protocol ? (protocol.charter || {}) : {}; + const stateful = protocolCharter.stateful_design || {}; + const targetClass = protocolCharter.target_policy_class || {}; + const protocolMeta = protocol ? (protocol.protocol || {}) : {}; + const blockerRegister = Array.isArray(protocolCharter.blocker_register) + ? protocolCharter.blocker_register : []; + const totals = resources.expected_for_planned_routes || {}; + const probabilities = design.action_probabilities || {}; + const baseline = enrollment.baseline_action || payload.baseline_agent_mode || '-'; + const floor = design.minimum_positive_exploration_probability; + const expectedByAlternate = labelScenario.expected_routes_for_target_by_alternate_action || {}; + const expectedRouteValues = Object.values(expectedByAlternate) + .map(Number).filter(Number.isFinite); + const expectedRoutes = expectedRouteValues.length ? Math.max(...expectedRouteValues) : null; + const confidenceRoutes = exact.minimum_routes_for_target_on_every_alternate_action; + const confidencePct = Number(exact.confidence_level) * 100; + const targetLabels = Number(labelScenario.target_observed_labels_per_alternate_action) || 0; + const eligible = enrollment.eligible === true; + latestRouteStudy = eligible ? JSON.parse(JSON.stringify(study)) : null; + el('routeBundleAdd').disabled = !latestRouteStudy; + el('routeStudyBaseline').textContent = policyLabLabel(baseline) || '-'; + el('routeStudyFloor').textContent = Number.isFinite(Number(floor)) + ? `${Math.round(Number(floor) * 100)}%` + : 'not enrolled'; + el('routeStudyTraffic').textContent = expectedRoutes && confidenceRoutes + ? `~${routeStudyNumber(expectedRoutes)} exp each / ${routeStudyNumber(confidenceRoutes)} @${routeStudyNumber(confidencePct)}% joint` + : 'no alternate support'; + el('routeStudyCost').textContent = routeStudyNumber(totals.cost_units, 1); + const tierOrder = ['low', 'moderate', 'high', 'frontier', 'unknown']; + const tiers = Object.values(resources.by_action || {}) + .map(row => String((row || {}).latency_tier || 'unknown')) + .filter((tier, index, values) => values.indexOf(tier) === index) + .sort((a, b) => tierOrder.indexOf(a) - tierOrder.indexOf(b)); + el('routeStudyLatency').textContent = tiers.length + ? `${policyLabLabel(tiers[0])}${tiers.length > 1 ? ' to ' + policyLabLabel(tiers[tiers.length - 1]) : ''}` + : '-'; + el('routeStudyCharter').textContent = study.design_hash ? String(study.design_hash).slice(0, 12) : '-'; + el('routeProtocolMode').textContent = protocol + ? policyLabLabel(stateful.selected_design_mode || 'not screened') : 'not enrolled'; + el('routeProtocolPolicy').textContent = protocol + ? policyLabLabel(targetClass.profile_name || 'not frozen') : '-'; + el('routeProtocolReview').textContent = protocol + ? policyLabLabel(stateful.selected_design_status || protocolMeta.state || 'blocked') + : policyLabLabel(payload.route_protocol_preflight_reason || 'unavailable'); + el('routeProtocolHash').textContent = protocol && protocol.protocol_hash + ? String(protocol.protocol_hash).slice(0, 12) : '-'; + + const distribution = el('routeStudyDistribution'); + distribution.replaceChildren(); + Object.entries(probabilities).forEach(([action, probability]) => { + const chip = document.createElement('span'); + chip.className = 'route-study-chip'; + chip.textContent = `${policyLabLabel(action)} ${Math.round(Number(probability) * 100)}%`; + distribution.appendChild(chip); + }); + const blockerList = el('routeProtocolBlockers'); + blockerList.replaceChildren(); + blockerRegister.forEach(row => { + const chip = document.createElement('span'); + chip.className = 'route-study-chip'; + chip.dataset.state = String((row || {}).status || 'unresolved'); + chip.textContent = `${policyLabLabel((row || {}).code || 'blocker')} - ${policyLabLabel((row || {}).status || 'unresolved')}`; + blockerList.appendChild(chip); + }); + const adjacent = Array.isArray(enrollment.adjacent_feasible_actions) + ? enrollment.adjacent_feasible_actions.map(policyLabLabel).join(', ') + : ''; + el('routeStudyStatus').textContent = eligible + ? `Rehearsal ready - ${adjacent || 'adjacent support'} - protocol ${policyLabLabel(protocolMeta.state || 'unavailable')} - activation, assignment, and execution remain off.` + : `Not enrolled - ${policyLabLabel(enrollment.reason || 'no feasible adjacent action')} - baseline remains deterministic.`; + const blockers = Array.isArray(boundaries.activation_blockers) + ? boundaries.activation_blockers.map(item => String(item).replaceAll('_', ' ')) + : []; + const trafficMethod = String(exact.method || 'not applicable').replaceAll('_', ' '); + const designReasons = Array.isArray(stateful.selected_design_blocking_reasons) + ? stateful.selected_design_blocking_reasons.map(item => String(item).replaceAll('_', ' ')) + : []; + el('routeStudyNote').textContent = + `Hypothetical repetition of this prompt-specific support: ${routeStudyNumber(traffic.planned_routes)} routes; ` + + `target ${routeStudyNumber(targetLabels)} observed ratings on every alternate; ` + + `assumed response ${routeStudyNumber(Number(labelScenario.assumed_feedback_rate) * 100)}%. ` + + `${trafficMethod}. This is not power, an observation model, OPE, or a promotion decision. ` + + `Stateful preflight: ${policyLabLabel(stateful.selected_design_status || 'not available')}; ` + + `${designReasons.length ? 'design review reasons: ' + designReasons.join('; ') + '. ' : ''}` + + `Declarations are not validation, route-level campaign assignment is screened out, and no seed is sealed. ` + + `Activation blockers (${blockers.length}): ${blockers.join('; ') || 'live integration unavailable'}. ` + + `No route was assigned, executed, or written.`; + } + + function routeProtocolBuildInput() { + const plannedClusters = parseInt(el('routeProtocolClusters').value, 10); + return { + study_plans: routeStudyStrata.map(plan => JSON.parse(JSON.stringify(plan))), + target_policy_profile: el('routeProtocolTarget').value, + design_mode: el('routeProtocolDesign').value, + carryover_scope: el('routeProtocolCarryover').value, + interference_scope: el('routeProtocolInterference').value, + temporal_variation: el('routeProtocolTemporal').value, + population_rule_id: 'interactive-auto-route-opt-in', + population_rule_version: '1', + cluster_key_schema_version: 'session-hash-v1', + planned_clusters: plannedClusters, + max_routes_per_cluster: 20, + analysis_every_clusters: Math.min(plannedClusters, 50), + block_length_routes: parseInt(el('routeProtocolBlock').value, 10), + washout_routes: parseInt(el('routeProtocolWashout').value, 10), + seed_commitment: null, + external_estimator_id: null, + external_reviewer_id: null, + }; + } + + function invalidateRouteReviewBundle(message) { + latestRouteReviewBundle = null; + el('routeBundleVerification').textContent = 'not built'; + el('routeBundleHash').textContent = '-'; + el('routeBundleDownload').disabled = true; + el('routeBundleStatus').textContent = message || + 'Inventory changed. Rebuild for full source-bound semantic verification.'; + } + + function renderRouteBundleInventory() { + const inventory = el('routeBundleInventory'); + inventory.replaceChildren(); + routeStudyStrata.forEach((plan, index) => { + const charter = (plan || {}).charter || {}; + const enrollment = charter.enrollment || {}; + const design = charter.probability_design || {}; + const hash = String((plan || {}).design_hash || '').slice(0, 10) || 'unhashed'; + const baseline = policyLabLabel(enrollment.baseline_action || 'unknown'); + const actions = Array.isArray(design.eligible_actions) + ? design.eligible_actions.map(policyLabLabel).join('/') : 'unknown support'; + const remove = document.createElement('button'); + remove.type = 'button'; + remove.className = 'route-study-chip'; + remove.textContent = `${index + 1}. ${hash} - ${baseline} - ${actions} - remove`; + remove.setAttribute('aria-label', `Remove support stratum ${index + 1}, ${hash}`); + remove.onclick = () => { + routeStudyStrata.splice(index, 1); + invalidateRouteReviewBundle('Support stratum removed. Rebuild before review.'); + renderRouteBundleInventory(); + }; + inventory.appendChild(remove); + }); + const count = routeStudyStrata.length; + el('routeBundleCount').textContent = `${count} ${count === 1 ? 'stratum' : 'strata'}`; + el('routeBundleBuild').disabled = count === 0; + el('routeBundleClear').disabled = count === 0; + } + + function addCurrentRouteStudyStratum() { + if (!latestRouteStudy) { + toast('err', 'Rehearse an eligible prompt-specific support stratum first'); + return; + } + const designHash = String(latestRouteStudy.design_hash || ''); + if (!designHash) { + toast('err', 'Current rehearsal has no canonical design hash'); + return; + } + if (routeStudyStrata.some(plan => String((plan || {}).design_hash || '') === designHash)) { + toast('err', 'That support stratum is already in the review inventory'); + return; + } + if (routeStudyStrata.length >= 100) { + toast('err', 'Browser review inventory is limited to 100 support strata'); + return; + } + routeStudyStrata.push(JSON.parse(JSON.stringify(latestRouteStudy))); + invalidateRouteReviewBundle('Support stratum added. Build the bundle for semantic verification.'); + renderRouteBundleInventory(); + toast('ok', 'Prompt-free support stratum added to the review inventory'); + } + + function renderRouteBundleVerification(verification) { + const checked = verification || {}; + const full = checked.verification_level === 'full_source_bound_reconstruction' && + checked.source_plan_reconstruction_performed === true; + el('routeBundleVerification').textContent = full ? 'full reconstruction' : 'verification failed'; + el('routeBundleHash').textContent = checked.bundle_hash + ? String(checked.bundle_hash).slice(0, 12) : '-'; + el('routeBundleDownload').disabled = !full || !latestRouteReviewBundle; + el('routeBundleStatus').textContent = full + ? `Verified ${checked.support_stratum_count || 0} canonical source strata by rebuilding the protocol. ` + + 'This is semantic conformance only: no signature, trusted timestamp, causal validation, assignment, or activation.' + : 'Full source-bound reconstruction did not complete; the bundle is not review-ready.'; + } + + async function buildRouteReviewBundle() { + if (!routeStudyStrata.length) { + toast('err', 'Add at least one support stratum first'); + return; + } + const button = el('routeBundleBuild'); + button.disabled = true; + el('routeBundleStatus').textContent = 'Reconstructing the campaign protocol from canonical prompt-free source plans...'; + try { + const result = await api('/api/route_study_protocol_bundle', routeProtocolBuildInput()); + latestRouteReviewBundle = result.route_protocol_review_bundle || null; + renderRouteBundleVerification(result.verification || {}); + toast('ok', 'Multi-stratum review bundle passed full source reconstruction'); + } catch (err) { + invalidateRouteReviewBundle(`Review bundle unavailable: ${err.message}`); + toast('err', err.message); + } finally { + renderRouteBundleInventory(); + } + } + + function downloadRouteReviewBundle() { + if (!latestRouteReviewBundle) return; + const rendered = JSON.stringify(latestRouteReviewBundle, null, 2) + '\n'; + const blob = new Blob([rendered], {type: 'application/json'}); + const url = URL.createObjectURL(blob); + const link = document.createElement('a'); + const hash = String(latestRouteReviewBundle.bundle_hash || 'unverified').slice(0, 12); + link.href = url; + link.download = `supermix-route-review-${hash}.json`; + document.body.appendChild(link); + link.click(); + link.remove(); + URL.revokeObjectURL(url); + } + + async function importRouteReviewFile(event) { + const input = event.target; + const file = input.files && input.files[0]; + if (!file) return; + try { + if (file.size > 2 * 1024 * 1024) { + throw new Error('Review file exceeds the 2 MiB browser limit'); + } + const parsed = JSON.parse(await file.text()); + let result; + if (parsed && parsed.schema_version === 'route-study-review-bundle-v1') { + result = await api('/api/route_study_protocol_bundle/audit', {bundle: parsed}); + latestRouteReviewBundle = parsed; + } else if (parsed && Array.isArray(parsed.study_plans)) { + result = await api('/api/route_study_protocol_bundle', parsed); + latestRouteReviewBundle = result.route_protocol_review_bundle || null; + } else { + throw new Error('Import must be a review bundle or closed prompt-free protocol build input'); + } + routeStudyStrata = JSON.parse(JSON.stringify( + (latestRouteReviewBundle || {}).source_study_plans || [] + )); + renderRouteBundleInventory(); + renderRouteBundleVerification(result.verification || {}); + toast('ok', 'Review file verified by full source-bound reconstruction'); + } catch (err) { + invalidateRouteReviewBundle(`Review import rejected: ${err.message}`); + toast('err', err.message); + } finally { + input.value = ''; + } + } + + function clearRouteReviewInventory() { + routeStudyStrata = []; + invalidateRouteReviewBundle( + 'Inventory cleared. Strata were client-side only; no ledger or memory records were written.' + ); + renderRouteBundleInventory(); + } + + async function previewRouteStudy() { + const button = el('routeStudyPreview'); + const text = el('prompt').value.trim(); + if (!text) { + renderRouteStudyUnavailable('Enter a prompt first; rehearsal needs the final prompt-specific route support.'); + return; + } + if (agentMode !== 'auto') { + renderRouteStudyUnavailable('Select Adaptive Router before rehearsing the adjacent-route study.'); + return; + } + button.disabled = true; + el('routeStudyStatus').textContent = 'Rehearsing post-filter support without assigning or running a route...'; + try { + const payload = buildRoutePayload(text); + Object.assign(payload, { + exploration_rate: Number(el('routeStudyEpsilon').value), + planned_routes: parseInt(el('routeStudyHorizon').value, 10), + scenario_confidence: 0.95, + assumed_feedback_rate: Number(el('routeStudyResponseRate').value), + target_observed_labels: parseInt(el('routeStudyTargetLabels').value, 10), + target_policy_profile: el('routeProtocolTarget').value, + protocol_design_mode: el('routeProtocolDesign').value, + carryover_scope: el('routeProtocolCarryover').value, + interference_scope: el('routeProtocolInterference').value, + temporal_variation: el('routeProtocolTemporal').value, + planned_clusters: parseInt(el('routeProtocolClusters').value, 10), + max_routes_per_cluster: 20, + analysis_every_clusters: Math.min( + parseInt(el('routeProtocolClusters').value, 10), 50 + ), + block_length_routes: parseInt(el('routeProtocolBlock').value, 10), + washout_routes: parseInt(el('routeProtocolWashout').value, 10), + }); + const result = await api('/api/route_study_plan', payload); + renderRouteStudy(result); + toast('ok', 'Adjacent-route rehearsal ready; execution stayed off'); + } catch (err) { + renderRouteStudyUnavailable(`Study rehearsal unavailable: ${err.message}`); + toast('err', err.message); + } finally { + button.disabled = false; + } + } + + async function sendRouteFeedback(rating, feedbackIntent, reason) { + if (!lastRouteFeedback) return; + const buttons = ['routeGoodBtn','routeBadBtn','routeDeeperBtn','routeCostBtn','routeSlowBtn'] + .map(id => el(id)).filter(Boolean); + buttons.forEach(btn => btn.disabled = true); + try { + const result = await api('/api/route_feedback', { + ...lastRouteFeedback, + rating, + feedback_intent: feedbackIntent, + reason: reason || feedbackIntent, + }); + const total = result.summary ? result.summary.total_feedback : ''; + setRouteFeedbackVisible(true, total ? `Route ${total}` : 'Route saved'); + if (result.summary) renderRouteHealth(result.summary); + refreshPolicyLab(); + const message = feedbackIntent === 'needs_deeper' ? 'Deeper-route preference saved' + : feedbackIntent === 'too_costly' ? 'Lower-cost preference saved' + : feedbackIntent === 'too_slow' ? 'Lower-latency preference saved' + : rating === 'up' ? 'Route preference saved' : 'Quality correction saved'; + toast('ok', message); + } catch (err) { + buttons.forEach(btn => btn.disabled = false); + toast('err', err.message); + } + } + // ── Typing indicator ───────────────────────────────────────────────── + function buildRoutePayload(text) { + const settings = { + agent_mode: agentMode, + auto_agent_budget: el('autoBudget').value, + auto_session_budget_units: parseFloat(el('sessionBudget').value) || 0, + auto_session_budget_target_routes: parseInt(el('sessionBudgetTargetRoutes').value) || 0, + loop_max_steps: parseInt(el('loopBudget').value), + memory_enabled: el('memToggle').value === 'on', + web_search_enabled: el('webToggle').value === 'on', + uploaded_image_path: currentUpload + }; + const requestedCycles = el('reasoningCycles').value; + const requestedAdaptive = el('adaptiveCompute').value; + if (requestedCycles !== 'model') settings.reasoning_cycles = requestedCycles; + if (requestedAdaptive !== 'model') settings.adaptive_compute = requestedAdaptive === 'on'; + return { + session_id: sessionId, + message: text, + model_key: el('modelSelect').value, + action_mode: 'text', + settings + }; + } + + el('routePlanBtn').onclick = async () => { + const text = el('prompt').value.trim(); + if (!text) return; + const btn = el('routePlanBtn'); + btn.disabled = true; + try { + const result = await api('/api/route_plan', buildRoutePayload(text)); + const estimate = result.route_economics_estimate || {}; + const cost = estimate.estimated_cost_units !== undefined ? `cost ${estimate.estimated_cost_units}` : 'cost -'; + const selected = result.selected_agent_mode || 'off'; + const frontier = result.route_frontier || {}; + const recommendation = frontier.recommended_agent_mode + ? `rec ${frontier.recommended_agent_mode}${frontier.selected_matches_recommendation ? '' : ' != selected'}` + : ''; + const recommendedQualityCost = frontier.recommended_estimated_quality_cost_score !== undefined + && frontier.recommended_estimated_quality_cost_score !== null + ? `rec qc ${frontier.recommended_estimated_quality_cost_score}` + : ''; + const blocker = frontier.selected_budget_blocker || frontier.recommended_budget_blocker; + const blockerText = blocker ? `block ${blocker}` : ''; + const cap = frontier.budget_cap_cost_units !== undefined && frontier.budget_cap_cost_units !== null + ? `cap ${frontier.budget_cap_cost_units}` + : ''; + const effective = frontier.effective_cap_cost_units !== undefined && frontier.effective_cap_cost_units !== null + ? `effective ${frontier.effective_cap_cost_units}` + : ''; + const remaining = frontier.remaining_cost_units !== undefined && frontier.remaining_cost_units !== null + ? `remain ${frontier.remaining_cost_units}` + : ''; + const pacing = frontier.pacing_cap_cost_units !== undefined && frontier.pacing_cap_cost_units !== null + ? `pace ${frontier.pacing_cap_cost_units}` + : ''; + const budgetPareto = Array.isArray(frontier.budget_feasible_pareto_modes) && frontier.budget_feasible_pareto_modes.length + ? `bpareto ${frontier.budget_feasible_pareto_modes.join(',')}` + : ''; + const alternatives = Array.isArray(result.route_alternatives) + ? result.route_alternatives + .slice() + .sort((a,b) => Number(a.frontier_rank || 99) - Number(b.frontier_rank || 99)) + .slice(0, 4) + .map(row => { + const value = row.estimated_cost_units !== undefined ? row.estimated_cost_units : '-'; + const rank = row.frontier_rank ? `#${row.frontier_rank}` : ''; + const qualityCost = row.estimated_quality_cost_score !== undefined && row.estimated_quality_cost_score !== null + ? `/qc${row.estimated_quality_cost_score}` + : ''; + const evidence = row.quality_source === 'adaptive_feedback' + ? '/adaptive' + : (row.quality_evidence_status && row.quality_evidence_status !== 'heuristic_prior' ? `/${row.quality_evidence_status}` : ''); + const lowerRaw = row.estimated_quality_lower_bound; + const upperRaw = row.estimated_quality_upper_bound; + const lower = Number(lowerRaw); + const upper = Number(upperRaw); + const confidence = lowerRaw !== null && lowerRaw !== undefined + && upperRaw !== null && upperRaw !== undefined + && Number.isFinite(lower) && Number.isFinite(upper) + ? `/q90:${Math.round(lower * 100)}-${Math.round(upper * 100)}` + : ''; + const riskRaw = row.risk_adjusted_quality_cost_score; + const riskQc = Number(riskRaw); + const risk = riskRaw !== null && riskRaw !== undefined + && Number.isFinite(riskQc) && row.confidence_status === 'established' + ? `/riskqc${riskQc}` + : ''; + const pareto = row.budget_feasible_pareto_frontier ? '/bpf' : (row.pareto_frontier ? '/pf' : ''); + return `${rank}${row.selected_agent_mode}${row.is_selected ? '*' : ''}:${value}${qualityCost}${evidence}${confidence}${risk}${pareto}`; + }) + .join(' ') + : ''; + const frontierBits = [ + recommendation, + recommendedQualityCost, + blockerText, + cap, + effective, + remaining, + pacing, + budgetPareto, + alternatives, + ].filter(Boolean).join(' | '); + addLoopStep(1, 'Route Plan', `${selected} | ${cost}${frontierBits ? ' | ' + frontierBits : ''}`, 'done'); + openPanelTab('mode'); + toast('ok', `Route plan: ${selected}`); + } catch (err) { + toast('err', err.message); + } finally { + btn.disabled = false; + } + }; + function addTyping() { const row = document.createElement('div'); row.className = 'msg asst'; @@ -1001,26 +2425,21 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: el('sendBtn').disabled = true; addTyping(); - if (agentMode === 'loop') { + if (agentMode === 'auto' || agentMode === 'loop' || agentMode === 'collective_loop') { el('loopSteps').innerHTML = ''; loopStep = 0; - addLoopStep(1, 'Target Initialization', 'Constructing reasoning graph...', 'active'); + lastRouteFeedback = null; + setRouteFeedbackVisible(false); + addLoopStep( + 1, + agentMode === 'auto' ? 'Adaptive Routing' : 'Target Initialization', + agentMode === 'auto' ? 'Scoring task complexity...' : 'Constructing reasoning graph...', + 'active' + ); switchPtab('mode'); } - const payload = { - session_id: sessionId, - message: text, - model_key: el('modelSelect').value, - action_mode: 'text', - settings: { - agent_mode: agentMode, - loop_max_steps: parseInt(el('loopBudget').value), - memory_enabled: el('memToggle').value === 'on', - web_search_enabled: el('webToggle').value === 'on', - uploaded_image_path: currentUpload - } - }; + const payload = buildRoutePayload(text); currentUpload = null; currentUpUrl = ''; el('uploadBar').style.display = 'none'; @@ -1030,14 +2449,38 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: const data = await api('/api/chat', payload); removeTyping(); - if (agentMode === 'loop' && data.agent_trace && data.agent_trace.loop_steps) { + if (data.agent_trace && data.agent_trace.loop_steps) { el('loopSteps').innerHTML = ''; data.agent_trace.loop_steps.forEach((s, i) => { - addLoopStep(i+1, s.goal || `Phase ${i+1}`, s.worker_excerpt || '', 'done'); + const score = scorePct(s.review_score ?? s.loop_score); + const bits = []; + if (s.worker_excerpt) bits.push(s.worker_excerpt); + if (score != null) bits.push(`Score ${score}%`); + if (s.stop_decision === 'stop') bits.push(stopReasonLabel(s.stop_reason_code)); + addLoopStep(i+1, s.goal || `Phase ${i+1}`, bits.join(' | '), 'done'); }); finaliseLoopSteps(); + } else if (data.agent_trace && data.agent_trace.auto_agent_policy) { + const p = data.agent_trace.auto_agent_policy; + el('loopSteps').innerHTML = ''; + addLoopStep(1, `Auto selected ${p.selected_agent_mode || 'off'}`, (p.reasons || []).slice(0,3).join(' | '), 'done'); + finaliseLoopSteps(); } + if (data.agent_trace && data.agent_trace.auto_agent_policy) { + const routeId = data.route_id || data.agent_trace.route_id || ''; + lastRouteFeedback = { + session_id: sessionId, + route_id: routeId, + }; + setRouteFeedbackVisible(Boolean(routeId)); + } else { + lastRouteFeedback = null; + setRouteFeedbackVisible(false); + } + refreshRouteHealth(); + refreshPolicyLab(); + addMsg('assistant', data.response || '(Inference finalized)', data.agent_trace); updateStatus(data); } catch(err) { @@ -1055,6 +2498,10 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: await api('/api/clear', { session_id: sessionId }); el('thread').innerHTML = ''; el('loopSteps').innerHTML = ''; + lastRouteFeedback = null; + setRouteFeedbackVisible(false); + renderRouteHealth({}); + renderPolicyLab({}); addMsg('assistant', 'Session memory cleared. Omni V46 is ready for the next message.'); toast('ok', 'System memory purged'); } catch(e) { toast('err', e.message); } @@ -1097,8 +2544,11 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: const label = st.active_model_label || selectedLabel || 'Auto Router'; const score = selected.common_overall_exact != null ? selected.common_overall_exact : selected.recipe_eval_accuracy; const scoreText = score != null ? `${(Number(score) * 100).toFixed(1)}%` : 'pending'; + const policy = data && data.agent_trace ? data.agent_trace.auto_agent_policy : null; + const budgetText = policy && policy.budget_profile ? `/${policy.budget_profile}` : ''; + const modeText = policy ? `${agentMode}${budgetText} -> ${policy.selected_agent_mode || 'off'}` : agentMode; el('panelStatus').textContent = - `model: ${label || '-'}\ndevice: ${st.device || '-'}\nbenchmark: ${scoreText}\nmode: ${agentMode}`; + `model: ${label || '-'}\ndevice: ${st.device || '-'}\nbenchmark: ${scoreText}\nmode: ${modeText}`; el('activePill').textContent = String(label).toLowerCase().includes('v46') ? 'V46 Champion' : (label || 'Auto'); const lowered = label.toLowerCase(); const pillClass = lowered.includes('v46') ? ' v46' : (lowered.includes('v48') ? ' v48' : (lowered.includes('v47') ? ' v47' : '')); @@ -1113,6 +2563,31 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: } catch(_) {} } + el('routeGoodBtn').onclick = () => sendRouteFeedback('up', 'good', 'satisfied'); + el('routeBadBtn').onclick = () => sendRouteFeedback('down', 'bad_quality', 'bad quality'); + el('routeDeeperBtn').onclick = () => sendRouteFeedback('down', 'needs_deeper', 'needs deeper reasoning'); + el('routeCostBtn').onclick = () => sendRouteFeedback('down', 'too_costly', 'too costly'); + el('routeSlowBtn').onclick = () => sendRouteFeedback('down', 'too_slow', 'too slow'); + el('policyLabRefresh').onclick = () => refreshPolicyLab(); + el('policyLabProfile').onchange = () => refreshPolicyLab(); + el('routeStudyPreview').onclick = () => previewRouteStudy(); + el('routeShadowRegistryRefresh').onclick = () => refreshRouteShadowRegistry(); + el('routeBundleAdd').onclick = () => addCurrentRouteStudyStratum(); + el('routeBundleBuild').onclick = () => buildRouteReviewBundle(); + el('routeBundleDownload').onclick = () => downloadRouteReviewBundle(); + el('routeBundleImport').onchange = event => importRouteReviewFile(event); + el('routeBundleClear').onclick = () => clearRouteReviewInventory(); + [ + 'routeProtocolTarget', 'routeProtocolDesign', 'routeProtocolCarryover', + 'routeProtocolInterference', 'routeProtocolTemporal', 'routeProtocolClusters', + 'routeProtocolBlock', 'routeProtocolWashout', + ].forEach(id => el(id).addEventListener('change', () => { + if (latestRouteReviewBundle) { + invalidateRouteReviewBundle('Protocol declarations changed. Rebuild before review.'); + } + })); + renderRouteBundleInventory(); + // ── Benchmark tab ───────────────────────────────────────────────────── async function loadBenchData() { try { @@ -1157,8 +2632,77 @@ def _benchmark_rows_for_ui(data: Dict[str, Any]) -> List[Dict[str, Any]]: // ── Init ───────────────────────────────────────────────────────────── initModels(); + refreshRouteHealth(); + refreshPolicyLab(); + refreshRouteShadowRegistry(); })(); + + + + """ @@ -1196,6 +2740,168 @@ def api_chat(): logging.exception("Chat request failed") return jsonify({"ok": False, "error": str(exc)}), 500 +@app.route("/api/route_plan", methods=["POST"]) +def api_route_plan(): + data = request.json or {} + try: + result = manager.preview_route_plan( + session_id=data.get("session_id", "default"), + prompt=data.get("message", ""), + model_key=data.get("model_key", "auto"), + action_mode=data.get("action_mode", "text"), + settings=data.get("settings", {}), + ) + return jsonify(result) + except (KeyError, ValueError, RuntimeError) as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route plan request failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_study_plan", methods=["POST"]) +def api_route_study_plan(): + data = request.json or {} + try: + result = manager.preview_route_study( + session_id=data.get("session_id", "default"), + prompt=data.get("message", ""), + model_key=data.get("model_key", "auto"), + action_mode=data.get("action_mode", "text"), + settings=data.get("settings", {}), + exploration_rate=data.get("exploration_rate", 0.10), + planned_routes=data.get("planned_routes", 2_000), + scenario_confidence=data.get("scenario_confidence", 0.95), + assumed_feedback_rate=data.get("assumed_feedback_rate", 0.30), + target_observed_labels=data.get("target_observed_labels", 20), + target_policy_profile=data.get("target_policy_profile", "balanced"), + protocol_design_mode=data.get( + "protocol_design_mode", "sticky_session_cluster" + ), + carryover_scope=data.get("carryover_scope", "unknown"), + interference_scope=data.get("interference_scope", "unknown"), + temporal_variation=data.get("temporal_variation", "unknown"), + planned_clusters=data.get("planned_clusters", 200), + max_routes_per_cluster=data.get("max_routes_per_cluster", 20), + analysis_every_clusters=data.get("analysis_every_clusters", 50), + block_length_routes=data.get("block_length_routes", 20), + washout_routes=data.get("washout_routes", 0), + ) + return jsonify(result) + except (KeyError, ValueError, RuntimeError) as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route study rehearsal failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_study_protocol_bundle", methods=["POST"]) +def api_route_study_protocol_bundle(): + try: + data = _read_strict_route_review_json() + if not isinstance(data, dict): + raise ValueError("route protocol review build input must be a JSON object") + _validate_route_review_request_size(data) + _validate_route_review_strata(data.get("study_plans")) + return jsonify(manager.build_route_protocol_review_bundle(data)) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route protocol review bundle build failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_study_protocol_bundle/audit", methods=["POST"]) +def api_route_study_protocol_bundle_audit(): + try: + data = _read_strict_route_review_json() + if not isinstance(data, dict) or set(data) != {"bundle"}: + raise ValueError("route protocol review audit accepts only a bundle object") + _validate_route_review_request_size(data) + bundle = data.get("bundle") + if not isinstance(bundle, dict): + raise ValueError("route protocol review bundle must be a JSON object") + _validate_route_review_strata(bundle.get("source_study_plans")) + return jsonify(manager.audit_route_protocol_review_bundle(bundle)) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route protocol review bundle audit failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_feedback", methods=["POST"]) +def api_route_feedback(): + data = request.json or {} + session_id = str(data.get("session_id") or "default") + try: + if not str(data.get("route_id") or "").strip(): + raise ValueError("route_id is required for route feedback") + result = manager.record_route_feedback( + session_id=session_id, + feedback={ + "route_id": data.get("route_id") or "", + "rating": data.get("rating") or "", + "feedback_intent": data.get("feedback_intent") or data.get("intent") or "", + "feedback_tags": data.get("feedback_tags") or [], + "reason": data.get("reason") or "", + }, + ) + return jsonify(result) + except ValueError as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route feedback request failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + +@app.route("/api/route_health", methods=["POST"]) +def api_route_health(): + data = request.json or {} + session_id = str(data.get("session_id") or "default") + try: + return jsonify({"ok": True, "route_health": manager.route_health_snapshot(session_id)}) + except Exception as exc: + logging.exception("Route health request failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_policy_lab", methods=["POST"]) +def api_route_policy_lab(): + data = request.json or {} + session_id = str(data.get("session_id") or "default") + profile = str(data.get("profile") or "balanced") + try: + return jsonify( + { + "ok": True, + "policy_lab": manager.route_policy_lab_snapshot(session_id, profile=profile), + } + ) + except ValueError as exc: + return jsonify({"ok": False, "error": str(exc)}), 400 + except Exception as exc: + logging.exception("Route policy lab request failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/route_shadow_registry/status", methods=["GET"]) +def api_route_shadow_registry_status(): + try: + response = jsonify( + { + "ok": True, + "route_shadow_registry": manager.route_shadow_registry_snapshot(), + } + ) + response.headers["Cache-Control"] = "no-store" + return response + except Exception as exc: + logging.exception("Route shadow registry status failed") + response = jsonify({"ok": False, "error": str(exc)}) + response.headers["Cache-Control"] = "no-store" + return response, 500 + + @app.route("/api/clear", methods=["POST"]) def api_clear(): data = request.json or {} @@ -1246,16 +2952,98 @@ def serve_upload(session_slug, filename): return send_from_directory(manager.uploads_dir / safe_slug, filename) +# ─── 3D Model View & Downloads ─────────────────────────────────────────────── + +@app.route("/api/three_d_model_view") +def api_three_d_model_view(): + try: + model_data = manager.three_d_model_view() + model_data["download_zip_url"] = "/download/three_d_model_zip" + model_data["download_summary_url"] = "/download/three_d_model_summary" + return jsonify({"ok": True, "model": model_data}) + except Exception as exc: + logging.exception("3D model view failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/download/three_d_model_zip") +def download_three_d_model_zip(): + try: + model_data = manager.three_d_model_view() + file_path = Path(model_data["zip_path"]) + return send_from_directory(str(file_path.parent), file_path.name, as_attachment=True) + except Exception as exc: + logging.exception("3D zip download failed") + return str(exc), 500 + + +@app.route("/download/three_d_model_summary") +def download_three_d_model_summary(): + try: + model_data = manager.three_d_model_view() + file_path = Path(model_data["summary_path"]) + return send_from_directory(str(file_path.parent), file_path.name, as_attachment=True) + except Exception as exc: + logging.exception("3D summary download failed") + return str(exc), 500 + + +# ─── Model Store API ───────────────────────────────────────────────────────── + +@app.route("/api/model_store") +def api_model_store(): + try: + force_refresh = request.args.get("force_refresh", "false").lower() == "true" + catalog = manager.model_store_catalog(force_refresh=force_refresh) + return jsonify({"ok": True, **catalog}) + except Exception as exc: + logging.exception("Model store catalog failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/model_store/jobs") +def api_model_store_jobs(): + try: + jobs_data = manager.model_store_jobs() + return jsonify({"ok": True, **jobs_data}) + except Exception as exc: + logging.exception("Model store jobs failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + +@app.route("/api/model_store/install", methods=["POST"]) +def api_model_store_install(): + try: + data = request.json or {} + file_name = data.get("file_name") + if not file_name: + return jsonify({"ok": False, "error": "file_name is required"}), 400 + job = manager.install_model_store_artifact(file_name) + return jsonify({"ok": True, "job": job}) + except Exception as exc: + logging.exception("Model store install failed") + return jsonify({"ok": False, "error": str(exc)}), 500 + + # ─── Entrypoint ─────────────────────────────────────────────────────────────── -def main(): - global manager +def build_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Supermix Studio X - V46 20-Suite Champion") parser.add_argument("--port", type=int, default=5000, help="Port to listen on") parser.add_argument("--models", type=str, default=str(DEFAULT_MODELS_DIR), help="Path to local models directory") - parser.add_argument("--host", type=str, default="0.0.0.0") - args = parser.parse_args() + parser.add_argument( + "--host", + type=str, + default="127.0.0.1", + help="Interface to bind (default: loopback only; opt in explicitly for remote access)", + ) + return parser + + +def main(): + global manager + args = build_arg_parser().parse_args() manager = UnifiedModelManager( records=discover_model_records(Path(args.models)), diff --git a/test_multimodel_memory_tools.py b/test_multimodel_memory_tools.py index 9ba7e3aa4..089f0317e 100644 --- a/test_multimodel_memory_tools.py +++ b/test_multimodel_memory_tools.py @@ -1,3 +1,4 @@ +import json from pathlib import Path from source.multimodel_memory import ConversationMemoryStore @@ -34,6 +35,374 @@ def test_memory_store_extracts_facts_and_examples(tmp_path: Path) -> None: assert "Relevant prior conversation examples" in bundle["context_block"] +def test_memory_store_session_filenames_isolate_colliding_legacy_slugs(tmp_path: Path) -> None: + root = tmp_path / "memory" + store = ConversationMemoryStore(root) + + store.update( + session_id="route/A", + user_text="Remember that this belongs to slash.", + assistant_text="Stored slash memory.", + model_key="v33_final", + route_reason="collision regression", + ) + store.update( + session_id="route?A", + user_text="Remember that this belongs to question mark.", + assistant_text="Stored question-mark memory.", + model_key="v33_final", + route_reason="collision regression", + ) + + slash_payload = store.load_session("route/A") + question_payload = store.load_session("route?A") + session_files = sorted(root.glob("route-a-*.json")) + + assert len(session_files) == 2 + assert session_files[0].name != session_files[1].name + assert slash_payload["session_id"] == "route/A" + assert question_payload["session_id"] == "route?A" + assert "belongs to slash" in json.dumps(slash_payload) + assert "question mark" not in json.dumps(slash_payload) + assert "belongs to question mark" in json.dumps(question_payload) + assert "belongs to slash" not in json.dumps(question_payload) + + +def test_memory_store_migrates_only_matching_legacy_session_file(tmp_path: Path) -> None: + root = tmp_path / "memory" + root.mkdir() + legacy_path = root / "route-a.json" + legacy_path.write_text( + json.dumps( + { + "session_id": "route/A", + "created_at": 10.0, + "updated_at": 20.0, + "memories": [{"kind": "fact", "text": "legacy slash memory"}], + "turns": [], + } + ), + encoding="utf-8", + ) + store = ConversationMemoryStore(root) + + payload = store.load_session("route/A") + migrated_files = list(root.glob("route-a-*.json")) + + assert payload["session_id"] == "route/A" + assert payload["memories"][0]["text"] == "legacy slash memory" + assert len(migrated_files) == 1 + assert json.loads(migrated_files[0].read_text(encoding="utf-8"))["session_id"] == "route/A" + assert not legacy_path.exists() + + +def test_memory_store_rejects_mismatched_legacy_session_identity(tmp_path: Path) -> None: + root = tmp_path / "memory" + root.mkdir() + legacy_path = root / "route-a.json" + legacy_payload = { + "session_id": "route/A", + "memories": [{"kind": "fact", "text": "must not cross sessions"}], + "turns": [{"user": "private slash turn"}], + } + legacy_path.write_text(json.dumps(legacy_payload), encoding="utf-8") + store = ConversationMemoryStore(root) + + payload = store.load_session("route?A") + + assert payload["session_id"] == "route?A" + assert payload["memories"] == [] + assert payload["turns"] == [] + assert legacy_path.exists() + assert json.loads(legacy_path.read_text(encoding="utf-8")) == legacy_payload + assert list(root.glob("route-a-*.json")) == [] + + +def test_memory_store_rejects_mismatched_hashed_session_identity(tmp_path: Path) -> None: + root = tmp_path / "memory" + store = ConversationMemoryStore(root) + session_path = store._path_for("route?A") + stored_payload = { + "session_id": "route/A", + "memories": [{"kind": "fact", "text": "must not trust filename alone"}], + "turns": [{"user": "private slash turn"}], + } + session_path.write_text(json.dumps(stored_payload), encoding="utf-8") + + payload = store.load_session("route?A") + + assert payload["session_id"] == "route?A" + assert payload["memories"] == [] + assert payload["turns"] == [] + assert json.loads(session_path.read_text(encoding="utf-8")) == stored_payload + + +def test_memory_store_route_feedback_round_trips_without_prompt_injection(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + session_id = "route-feedback-session" + + result = store.add_feedback( + session_id=session_id, + feedback={ + "route_id": "route-1", + "prompt": "Research and implement the runtime router.", + "response": "Internal route feedback should stay out of prompts.", + "selected_agent_mode": "collective_loop", + "rating": "down", + "reason": "too expensive", + "auto_agent_policy": { + "score": 5, + "reason": "high_complexity_with_collective_capacity", + "reasons": ["external_evidence", "workflow_depth"], + }, + "route_economics": { + "estimate": { + "selected_agent_mode": "collective_loop", + "budget_profile": "balanced", + "estimated_model_calls": 12, + "estimated_cost_units": 12.25, + "latency_tier": "frontier", + }, + "actual": { + "elapsed_ms": 251.5, + "model_calls": 9, + "tool_calls": 1, + "cost_units": 9.25, + "latency_tier": "high", + }, + }, + "model_key": "omni_collective_v8", + "route_reason": "Auto orchestration selected collective_loop.", + }, + ) + + assert result["feedback"]["rating"] == "down" + assert result["feedback"]["score_delta"] == -1 + assert result["feedback"]["route_economics"]["estimate"]["estimated_model_calls"] == 12 + assert result["feedback"]["route_economics"]["actual"]["cost_units"] == 9.25 + assert result["summary"]["mode_scores"]["collective_loop"]["net"] == -1 + assert result["summary"]["mode_scores"]["collective_loop"]["economics"]["sample_count"] == 1 + assert result["summary"]["mode_scores"]["collective_loop"]["economics"]["avg_cost_units"] == 9.25 + assert result["summary"]["mode_scores"]["collective_loop"]["adaptive"]["quality_score"] == 0.0 + assert result["summary"]["adaptive"]["quality_cost_score"] is not None + assert result["summary"]["mode_scores"]["loop"]["economics"]["sample_count"] == 0 + assert result["summary"]["economics"]["sample_count"] == 1 + assert result["summary"]["economics"]["avg_cost_units"] == 9.25 + assert result["summary"]["economics"]["avg_elapsed_ms"] == 251.5 + snapshot = store.session_snapshot(session_id) + assert snapshot["route_feedback_count"] == 1 + assert snapshot["route_feedback"]["economics"]["avg_model_calls"] == 9.0 + + bundle = store.build_context(session_id, "Research the runtime router.") + assert "route feedback" not in bundle["context_block"].lower() + assert "too expensive" not in bundle["context_block"].lower() + assert "cost_units" not in bundle["context_block"].lower() + + +def test_route_feedback_joins_server_usage_and_revisions_by_route_id(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + session_id = "trusted-route-session" + store.add_route_usage( + session_id=session_id, + route_id="server-route-1", + prompt="Authoritative prompt", + selected_agent_mode="collective", + route_economics={"actual": {"elapsed_ms": 250.0, "model_calls": 2, "cost_units": 2.5}}, + auto_agent_policy={ + "policy_id": "auto-route-v2", + "policy_version": "2.0.0", + "feature_schema_version": "route-context-v1", + "decision_type": "deterministic", + "action_mode": "text", + "score": 3, + "selected_agent_mode": "collective", + "allowed_agent_modes": ["off", "collective", "loop"], + "eligible_actions": ["off", "collective", "loop"], + "action_probabilities": {"off": 0.0, "collective": 1.0, "loop": 0.0}, + "post_filter_action_probabilities": {"off": 0.0, "collective": 1.0, "loop": 0.0}, + "probability_stage": "post_filter", + "logging_propensity": 1.0, + "decision_context": { + "action_mode": "text", + "budget_profile": "balanced", + "score": 3, + "allowed_agent_modes": ["off", "collective", "loop"], + }, + }, + route_reason="server reason", + model_key="server-model", + ) + + first = store.add_feedback( + session_id=session_id, + feedback={ + "route_id": "server-route-1", + "prompt": "spoofed prompt", + "selected_agent_mode": "collective_loop", + "rating": "up", + "route_economics": {"actual": {"elapsed_ms": 9999, "cost_units": 99}}, + "model_key": "spoofed-model", + }, + ) + row = first["feedback"] + assert row["prompt"] == "Authoritative prompt" + assert row["selected_agent_mode"] == "collective" + assert row["route_economics"]["actual"]["cost_units"] == 2.5 + assert row["model_key"] == "server-model" + assert row["evidence_source"] == "server_route_join" + assert row["auto_agent_policy"]["policy_version"] == "2.0.0" + assert row["auto_agent_policy"]["decision_context"]["score"] == 3 + + revised = store.add_feedback( + session_id=session_id, + feedback={ + "route_id": "server-route-1", + "rating": "down", + "feedback_intent": "too_slow", + }, + ) + assert revised["feedback"]["feedback_revision"] == 2 + assert revised["feedback"]["feedback_intent"] == "too_slow" + assert revised["summary"]["total_feedback"] == 1 + + +def test_route_feedback_summary_geometrically_discounts_stale_positives(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + session_id = "route-adaptive-session" + prompt = "Implement a runtime integration." + + for idx in range(6): + store.add_feedback( + session_id=session_id, + feedback={ + "route_id": f"old-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "route_economics": {"actual": {"elapsed_ms": 100.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + for idx in range(2): + store.add_feedback( + session_id=session_id, + feedback={ + "route_id": f"new-bad-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "route_economics": {"actual": {"elapsed_ms": 120.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + + summary = store.route_feedback_summary(session_id, prompt) + loop_score = summary["mode_scores"]["loop"] + + assert loop_score["net"] == 4 + assert loop_score["adaptive"]["weighted_net"] < 0 + assert loop_score["adaptive"]["quality_score"] < 0.5 + assert loop_score["adaptive"]["regression_signal"] is True + assert summary["adaptive"]["regression_signal"] is True + + +def test_route_feedback_intents_keep_preference_pressure_out_of_quality(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + prompt = "Implement and verify the runtime router." + + cost_result = store.add_feedback( + session_id="cost-intent-session", + feedback={ + "route_id": "route-cost", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "feedback_intent": "too_costly", + "route_economics": {"actual": {"elapsed_ms": 8000, "model_calls": 7, "cost_units": 7.0}}, + }, + ) + cost_feedback = cost_result["feedback"] + cost_adaptive = cost_result["summary"]["mode_scores"]["loop"]["adaptive"] + assert cost_feedback["score_delta"] == 0 + assert cost_feedback["feedback_axes"]["quality"] is None + assert cost_adaptive["quality_sample_count"] == 0 + assert cost_adaptive["quality_score"] is None + assert cost_adaptive["preference_direction"] == "shallower" + assert cost_adaptive["weighted_cost_pressure"] == 1.0 + + depth_result = store.add_feedback( + session_id="depth-intent-session", + feedback={ + "route_id": "route-depth", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "down", + "feedback_intent": "needs_deeper", + }, + ) + depth_adaptive = depth_result["summary"]["mode_scores"]["collective"]["adaptive"] + assert depth_result["feedback"]["score_delta"] == 0 + assert depth_adaptive["quality_sample_count"] == 0 + assert depth_adaptive["preference_direction"] == "deeper" + + +def test_route_feedback_quality_reports_effective_sample_confidence_bounds(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + prompt = "Calibrate the route controller." + result = None + for idx in range(6): + result = store.add_feedback( + session_id="confidence-session", + feedback={ + "route_id": f"route-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "feedback_intent": "good", + "route_economics": {"actual": {"elapsed_ms": 500, "model_calls": 2, "cost_units": 2.0}}, + }, + ) + + adaptive = result["summary"]["mode_scores"]["collective"]["adaptive"] + assert adaptive["quality_sample_count"] == 6 + assert 3.0 <= adaptive["effective_sample_size"] < 6.0 + assert 0.0 <= adaptive["quality_lower_bound"] <= adaptive["quality_score"] + assert adaptive["quality_score"] <= adaptive["quality_upper_bound"] <= 1.0 + assert adaptive["quality_cost_lower_bound"] < adaptive["quality_cost_score"] + assert adaptive["confidence_status"] == "established" + assert adaptive["interval_kind"] == "recency_weighted_wilson_heuristic" + assert adaptive["coverage_claim"] == "heuristic_associational_only" + assert adaptive["effective_sample_size_ceiling"] == 4.0 + + +def test_memory_store_route_usage_summary_counts_full_retained_ledger(tmp_path: Path) -> None: + store = ConversationMemoryStore(tmp_path / "memory") + session_id = "route-usage-session" + payload = store.load_session(session_id) + payload["route_usage"] = [ + { + "route_id": f"route-{idx}", + "prompt": "Track automatic route usage.", + "selected_agent_mode": "off", + "route_economics": { + "actual": { + "elapsed_ms": 10.0, + "model_calls": 1, + "cost_units": 1.0, + } + }, + } + for idx in range(130) + ] + store.save_session(session_id, payload) + + summary = store.route_usage_summary(session_id) + + assert summary["total_routes"] == 130 + assert summary["recent_routes"] == 130 + assert summary["economics"]["sample_count"] == 130 + assert summary["economics"]["total_cost_units"] == 130.0 + assert summary["mode_economics"]["off"]["total_cost_units"] == 130.0 + + def test_tool_call_parsing_and_stripping() -> None: text = "TOOL:web_search: latest OpenAI model docs\nTOOL:open_cmd: C:\\work\nThen summarize the result." assert parse_tool_calls(text) == ["latest OpenAI model docs"] diff --git a/test_multimodel_runtime.py b/test_multimodel_runtime.py index 329b3488e..f1e52ee18 100644 --- a/test_multimodel_runtime.py +++ b/test_multimodel_runtime.py @@ -1,12 +1,16 @@ +import json import sys from pathlib import Path +import pytest + from source.multimodel_catalog import ModelRecord SOURCE_DIR = Path(__file__).resolve().parent / "source" if str(SOURCE_DIR) not in sys.path: sys.path.insert(0, str(SOURCE_DIR)) +import source.multimodel_runtime as runtime_module from source.multimodel_runtime import ChatResult, UnifiedModelManager @@ -44,6 +48,105 @@ def unload(self) -> None: return None +def test_champion_backend_forwards_runtime_compute_controls_and_telemetry() -> None: + record = _record("v51-runtime", "champion_chat", ("chat",), 0.9) + captured = {} + + class _Engine: + def chat(self, **kwargs): + captured.update(kwargs) + return { + "response": "computed response", + "timing_ms": {"total": 12.5}, + "compute": { + "applied": True, + "cycles_used": 5.0, + "exit_reason": "prediction_stable", + }, + } + + backend = runtime_module.ChampionChatBackend.__new__(runtime_module.ChampionChatBackend) + backend.record = record + backend.engine = _Engine() + result = backend.chat( + "compute-session", + "Solve this carefully.", + { + "reasoning_cycles": 5, + "adaptive_compute": True, + "adaptive_exit_tol": 0.04, + "adaptive_exit_entropy": 0.25, + "prediction_stability_patience": 3, + "prediction_stability_tol": 0.01, + }, + ) + + assert captured["reasoning_cycles"] == 5 + assert captured["adaptive_compute"] is True + assert captured["adaptive_exit_tol"] == 0.04 + assert captured["adaptive_exit_entropy"] == 0.25 + assert captured["prediction_stability_patience"] == 3 + assert captured["prediction_stability_tol"] == 0.01 + assert result.compute == { + "applied": True, + "cycles_used": 5.0, + "exit_reason": "prediction_stable", + } + assert result.to_dict()["compute"]["exit_reason"] == "prediction_stable" + + +def _auto_route_records() -> tuple[ModelRecord, ...]: + return ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + + +def test_manager_exposes_only_read_only_shadow_registry_status(tmp_path: Path) -> None: + from source.route_policy_protocol import build_route_study_review_bundle_from_input + from source.route_policy_protocol_cli import _example_bundle_input + from source.route_policy_shadow_registry import RouteShadowAssignmentRegistry + + manager = UnifiedModelManager( + records=(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + absent = manager.route_shadow_registry_snapshot() + assert absent["available"] is False + assert absent["read_only"] is True + assert absent["execution_enabled"] is False + assert not manager.route_shadow_registry_path.exists() + + registry = RouteShadowAssignmentRegistry(manager.route_shadow_registry_path) + bundle = build_route_study_review_bundle_from_input(_example_bundle_input()) + sealed = registry.seal_campaign(bundle, bytes(range(32))) + campaign_id = sealed["public_package"]["campaign_seal"]["seal"]["campaign_id"] + registry.append_assignment_commitment( + campaign_id=campaign_id, + seed_capsule=sealed["private_seed_capsule"], + cluster_identifier="runtime-status-private-cluster", + ) + before_database = manager.route_shadow_registry_path.read_bytes() + + status = manager.route_shadow_registry_snapshot() + assert status["available"] is True + assert status["status"] == "verified" + assert status["read_only"] is True + assert status["registry_location"] == "memory/route-policy-shadow-registry.sqlite3" + assert "registry_path" not in status + assert status["campaign_count"] == 1 + assert status["campaigns"][0]["artifact_audit_ok"] is True + assert status["campaigns"][0]["commitment_count"] == 1 + assert status["execution_enabled"] is False + assert status["activation_available"] is False + assert status["automatic_promotion_allowed"] is False + assert manager.route_shadow_registry_path.read_bytes() == before_database + rendered = json.dumps(status, sort_keys=True) + assert sealed["private_seed_capsule"]["seed_material_base64url"] not in rendered + assert "runtime-status-private-cluster" not in rendered + + def test_collective_panel_includes_omni_collective_v2_v3_v4_v5_v6_v7_v8_v8_preview_v40_and_domain_specialists(tmp_path: Path) -> None: records = ( _record("v33_final", "champion_chat", ("chat",), 0.18), @@ -189,119 +292,503 @@ def fake_build_backend(record: ModelRecord): assert "fell back" in payload["route_reason"].lower() -def test_model_store_catalog_marks_installed_and_selectable_records(tmp_path: Path) -> None: - models_dir = tmp_path / "models" - models_dir.mkdir() - installed = models_dir / "dcgan_v2_in_progress.zip" - installed.write_bytes(b"zip") +def test_auto_agent_mode_routes_complex_prompt_to_collective_loop(tmp_path: Path, monkeypatch) -> None: records = ( - _record("dcgan_v2_in_progress", "dcgan_image", ("image",), None), + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), ) manager = UnifiedModelManager( records=records, extraction_root=tmp_path / "extract", generated_dir=tmp_path / "generated", - models_dir=models_dir, - common_summary_path=tmp_path / "missing_summary.json", ) + captured = {} - manager._fetch_model_store_manifest_locked = lambda force_refresh=False: { # type: ignore[method-assign] - "models": [ - {"file_name": "dcgan_v2_in_progress.zip", "size_bytes": 3, "size_mb": 0.0, "family": "gan"}, - {"file_name": "supermix_omni_collective_v8_preview_20260407_001155.zip", "size_bytes": 10, "size_mb": 0.0, "family": "fusion"}, - ] + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + captured["pre_execution_decision"] = manager.route_policy_ledger.list_decisions( + session_id="auto-complex-session" + )[0] + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response="auto routed through collective loop", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + + payload = manager.handle_prompt( + session_id="auto-complex-session", + prompt=( + "Research the latest evidence, design and implement a multi-step runtime integration, " + "debug regressions, verify with tests, and explain benchmark tradeoffs." + ), + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + pre_execution_decision = captured["pre_execution_decision"] + assert set(pre_execution_decision["outcome_contracts"]) == { + "route_success", + "user_quality_rating", + "cost", + "latency", } + assert pre_execution_decision["outcome_contracts_precommitted_at_begin"] is True + assert pre_execution_decision["outcome_contracts_defaulted_at_begin"] is False + assert pre_execution_decision["outcome_events"] == [] + for outcome_name, contract in pre_execution_decision["outcome_contracts"].items(): + assert contract["schema_version"] == "route-outcome-contract-v1" + assert contract["outcome_name"] == outcome_name + assert contract["precommitted"] is True + assert contract["commitment_source"] == "caller" + assert len(contract["contract_hash"]) == 64 + assert captured["collective_mode"] is True + assert captured["settings"]["agent_mode"] == "collective_loop" + assert captured["settings"]["reasoning_cycles"] == policy["reasoning_cycles"] + assert captured["settings"]["adaptive_compute"] is True + assert policy["selected_agent_mode"] == "collective_loop" + assert policy["policy_id"] == "auto-route-v2" + assert policy["policy_version"] == "2.0.0" + assert policy["feature_schema_version"] == "route-context-v1" + assert policy["decision_type"] == "deterministic" + assert policy["probability_stage"] == "post_filter" + assert policy["logging_propensity"] == 1.0 + assert policy["post_filter_action_probabilities"]["collective_loop"] == 1.0 + assert sum(policy["post_filter_action_probabilities"].values()) == 1.0 + assert policy["counterfactual_support"] == "none_deterministic_logging" + assert policy["logging_support"]["schema_version"] == "route-support-v1" + assert policy["logging_support"]["decision_type"] == "deterministic" + assert policy["logging_support"]["sampler"]["exploration_rate"] == 0.0 + assert policy["candidate_set_hash"] == policy["logging_support"]["candidate_set_hash"] + assert policy["distribution_hash"] == policy["logging_support"]["distribution_hash"] + assert payload["route_id"] == payload["agent_trace"]["route_id"] + assert len(payload["route_id"]) == 32 + assert policy["runtime_compute_request"] == { + "reasoning_cycles": policy["reasoning_cycles"], + "adaptive_compute": True, + "source": "auto_route_policy", + } + assert policy["collective_available"] is True + assert policy["route_economics_estimate"]["selected_agent_mode"] == "collective_loop" + assert policy["route_economics_estimate"]["estimated_model_calls"] >= 8 + assert payload["agent_trace"]["route_economics"]["estimate"]["selected_agent_mode"] == "collective_loop" + assert payload["agent_trace"]["route_economics"]["actual"]["elapsed_ms"] >= 0 + assert payload["agent_trace"]["requested_agent_mode"] == "auto" + assert payload["agent_trace"]["resolved_agent_mode"] == "collective_loop" + assert "Auto orchestration selected collective_loop" in payload["route_reason"] - payload = manager.model_store_catalog(force_refresh=True) - by_name = {row["file_name"]: row for row in payload["models"]} - assert by_name["dcgan_v2_in_progress.zip"]["installed"] is True - assert by_name["dcgan_v2_in_progress.zip"]["selectable"] is True - assert by_name["supermix_omni_collective_v8_preview_20260407_001155.zip"]["known"] is True - assert by_name["supermix_omni_collective_v8_preview_20260407_001155.zip"]["installed"] is False + policy_lab = manager.route_policy_lab_snapshot("auto-complex-session", profile="efficiency") + assert policy_lab["profile"]["name"] == "efficiency" + assert policy_lab["support"]["usage"]["unique_route_ids"] == 1 + assert policy_lab["evidence_source"] == "durable_sqlite_ledger" + assert policy_lab["evaluation_readiness"]["schema_version"] == "route-readiness-v2" + assert policy_lab["evaluation_readiness"]["logging_integrity"]["valid_routes"] == 0 + assert policy_lab["evaluation_readiness"]["target_overlap"]["effective_sample_size"] == 0.0 + assert policy_lab["evaluation_readiness"]["ready_for_external_ope"] is False + assert "no_valid_randomized_overlap" in policy_lab["promotion_gate"]["blocking_reason_codes"] + assert policy_lab["promotion_gate"]["deployment"] == "shadow_only" + assert policy_lab["promotion_gate"]["automatic_promotion_allowed"] is False + assert policy_lab["evaluation_readiness"]["policy_value_estimated"] is False + outcome_maturity = policy_lab["outcome_contract_maturity"] + assert outcome_maturity["schema_version"] == "route-outcome-maturity-v1" + assert outcome_maturity["descriptive_only"] is True + assert outcome_maturity["policy_value_estimate"] is None + assert outcome_maturity["causal_identification"] == "not_performed" + assert outcome_maturity["missingness_identification"] == "not_performed" + assert any( + "diagnostic-only telemetry" in warning + and "not a policy-value estimator" in warning + for warning in policy_lab["warnings"] + ) + assert policy_lab["durable_ledger"]["counts"] == { + "started": 1, + "completed": 1, + "failed": 0, + "inflight": 0, + } + decision = manager.route_policy_ledger.get_decision(payload["route_id"]) + assert decision["status"] == "completed" + assert decision["chosen_mode"] == "collective_loop" + assert decision["executed_mode"] == "collective_loop" + assert decision["decision_context"]["requested_agent_mode"] == "auto" + assert decision["decision_type"] == "deterministic" + assert decision["probability_stage"] == "post_filter" + assert decision["candidate_set_hash"] == policy["candidate_set_hash"] + assert decision["distribution_hash"] == policy["distribution_hash"] + assert "actual" not in decision["actual_economics"] + assert decision["actual_economics"]["cost_units"] == pytest.approx( + payload["agent_trace"]["route_economics"]["actual"]["cost_units"] + ) + assert decision["actual_economics"]["elapsed_ms"] == pytest.approx( + payload["agent_trace"]["route_economics"]["actual"]["elapsed_ms"] + ) + terminal_outcome_names = [event["outcome_name"] for event in decision["outcome_events"]] + assert terminal_outcome_names.count("route_success") == 1 + assert terminal_outcome_names.count("cost") == 1 + assert terminal_outcome_names.count("latency") == 1 + assert "user_quality_rating" not in terminal_outcome_names + terminal_events = {event["outcome_name"]: event for event in decision["outcome_events"]} + assert terminal_events["route_success"]["value"] is True + assert terminal_events["route_success"]["observation_status"] == "observed" + assert terminal_events["route_success"]["event_source"] == "route_completion" + assert terminal_events["cost"]["observation_status"] == "observed" + assert terminal_events["cost"]["value"] == pytest.approx( + decision["actual_economics"]["cost_units"] + ) + assert terminal_events["latency"]["observation_status"] == "observed" + assert terminal_events["latency"]["value"] == pytest.approx( + decision["actual_economics"]["elapsed_ms"] + ) + assert payload["agent_trace"]["route_ledger"]["status"] == "completed" + assert payload["agent_trace"]["route_ledger"]["candidate_set_hash"] == policy["candidate_set_hash"] + first_feedback = manager.record_route_feedback( + session_id="auto-complex-session", + feedback={"route_id": payload["route_id"], "rating": "up", "feedback_intent": "good"}, + ) + assert first_feedback["durable_feedback"] == { + "status": "committed", + "revision": 1, + "idempotent": False, + "eligible_for_readiness": True, + } + observed_lab = manager.route_policy_lab_snapshot("auto-complex-session", profile="efficiency") + assert observed_lab["matched_observed"]["avg_cost_units"] == pytest.approx( + decision["actual_economics"]["cost_units"] + ) + assert observed_lab["matched_observed"]["avg_elapsed_ms"] == pytest.approx( + decision["actual_economics"]["elapsed_ms"] + ) + identical_retry = manager.record_route_feedback( + session_id="auto-complex-session", + feedback={"route_id": payload["route_id"], "rating": "up", "feedback_intent": "good"}, + ) + assert identical_retry["durable_feedback"]["revision"] == 1 + assert identical_retry["durable_feedback"]["idempotent"] is True + assert identical_retry["feedback"]["feedback_revision"] == 1 + decision_with_feedback = manager.route_policy_ledger.get_decision(payload["route_id"]) + assert decision_with_feedback["feedback_status"] == "known" + assert decision_with_feedback["feedback_revision_count"] == 1 + assert decision_with_feedback["latest_feedback"]["feedback"]["observation_status"] == "observed" + feedback_outcome_names = [ + event["outcome_name"] for event in decision_with_feedback["outcome_events"] + ] + assert feedback_outcome_names.count("user_quality_rating") == 1 + quality_event = next( + event + for event in decision_with_feedback["outcome_events"] + if event["outcome_name"] == "user_quality_rating" + ) + assert quality_event["value"] == 1 + assert quality_event["observation_status"] == "observed" + assert quality_event["event_source"] == "feedback_revision" + non_quality_feedback = manager.record_route_feedback( + session_id="auto-complex-session", + feedback={ + "route_id": payload["route_id"], + "rating": "down", + "feedback_intent": "too_slow", + }, + ) + assert non_quality_feedback["durable_feedback"]["revision"] == 2 + decision_with_non_quality_feedback = manager.route_policy_ledger.get_decision(payload["route_id"]) + assert decision_with_non_quality_feedback["feedback_revision_count"] == 2 + quality_events = [ + event + for event in decision_with_non_quality_feedback["outcome_events"] + if event["outcome_name"] == "user_quality_rating" + ] + assert len(quality_events) == 2 + assert [event["observation_status"] for event in quality_events] == [ + "observed", + "not_observed", + ] + assert quality_events[1]["value"] is None + assert quality_events[1]["event_source"] == "feedback_revision" + with pytest.raises(ValueError, match="does not belong to this session"): + manager.record_route_feedback( + session_id="different-session", + feedback={"route_id": payload["route_id"], "rating": "up", "feedback_intent": "good"}, + ) -def test_loop_agent_runs_until_reviewer_marks_complete(tmp_path: Path, monkeypatch) -> None: +def test_durable_feedback_survives_compatibility_mirror_failure_and_reconciles( + tmp_path: Path, + monkeypatch, +) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + session_id = "mirror-reconciliation-session" + route_id = "11111111-1111-4111-8111-111111111111" + manager.route_policy_ledger.begin_decision( + session_id=session_id, + policy_name="explicit-route-v1", + policy_version="1.0.0", + policy_schema_version="route-context-v1", + decision_context={"action_mode": "text", "selected_agent_mode": "off"}, + eligible_modes=["off"], + chosen_mode="off", + action_probabilities={"off": 1.0}, + estimated_economics={"estimated_cost_units": 1.0}, + route_id=route_id, + ) + manager.route_policy_ledger.complete_decision( + route_id, + success=True, + executed_mode="off", + actual_economics={"cost_units": 1.0}, + ) + + original_commit = manager.memory_store.commit_feedback + + def fail_mirror(**_kwargs): + raise OSError("simulated compatibility mirror failure") + + monkeypatch.setattr(manager.memory_store, "commit_feedback", fail_mirror) + accepted = manager.record_route_feedback( + session_id=session_id, + feedback={"route_id": route_id, "rating": "up", "feedback_intent": "good"}, + ) + + assert accepted["ok"] is True + assert accepted["compatibility_mirror"]["status"] == "pending_reconciliation" + assert accepted["durable_feedback"]["revision"] == 1 + durable = manager.route_policy_ledger.get_decision(route_id) + assert durable["feedback_revision_count"] == 1 + durable_quality_events = [ + event for event in durable["outcome_events"] + if event["outcome_name"] == "user_quality_rating" + ] + assert len(durable_quality_events) == 1 + assert durable_quality_events[0]["value"] == 1 + assert durable_quality_events[0]["event_source"] == "feedback_revision" + assert manager.memory_store.load_session(session_id)["route_feedback"] == [] + + monkeypatch.setattr(manager.memory_store, "commit_feedback", original_commit) + reconciled = manager.record_route_feedback( + session_id=session_id, + feedback={"route_id": route_id, "rating": "up", "feedback_intent": "good"}, + ) + assert reconciled["durable_feedback"]["idempotent"] is True + assert reconciled["compatibility_mirror"]["status"] == "committed" + assert reconciled["feedback"]["feedback_revision"] == 1 + assert manager.route_policy_ledger.get_decision(route_id)["feedback_revision_count"] == 1 + assert [ + event["outcome_name"] + for event in manager.route_policy_ledger.get_decision(route_id)["outcome_events"] + ].count("user_quality_rating") == 1 + + +def test_legacy_json_policy_lab_fallback_is_descriptive_only(tmp_path: Path) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + session_id = "legacy-policy-lab-session" + manager.memory_store.add_route_usage( + session_id=session_id, + route_id="legacy-route-1", + prompt="Legacy compatibility evidence", + selected_agent_mode="off", + route_economics={"actual": {"cost_units": 1.0, "elapsed_ms": 20.0}}, + auto_agent_policy={ + "policy_id": "auto-route-v2", + "policy_version": "2.0.0", + "feature_schema_version": "route-context-v1", + "decision_type": "deterministic", + "probability_stage": "post_filter", + "score": 0, + "selected_agent_mode": "off", + "allowed_agent_modes": ["off"], + "eligible_actions": ["off"], + "post_filter_action_probabilities": {"off": 1.0}, + }, + ) + + report = manager.route_policy_lab_snapshot(session_id) + lifecycle = report["evaluation_readiness"]["lifecycle_integrity"] + assert report["evidence_source"] == "legacy_json_compatibility" + assert report["compatibility_view"]["used_for_analysis"] is True + assert report["compatibility_view"]["used_for_readiness"] is False + assert report["compatibility_view"]["eligible_for_readiness"] is False + assert lifecycle["durable_lifecycle_present"] is False + assert lifecycle["reconciled"] is False + assert "durable_lifecycle_required" in report["promotion_gate"]["blocking_reason_codes"] + assert report["evaluation_readiness"]["ready_for_external_ope"] is False + + +def test_route_ledger_records_runtime_failure_before_reraising(tmp_path: Path, monkeypatch) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + observed = {} + + def fail_loop_agent(**_kwargs): + observed["during"] = manager.route_policy_ledger.report(session_id="failed-route-session") + observed["policy_lab_during"] = manager.route_policy_lab_snapshot( + "failed-route-session", + profile="balanced", + ) + raise RuntimeError("simulated route backend failure") + + monkeypatch.setattr(manager, "_run_loop_agent_text", fail_loop_agent) + + with pytest.raises(RuntimeError, match="simulated route backend failure"): + manager.handle_prompt( + session_id="failed-route-session", + prompt=( + "Research the latest evidence, implement the whole project, verify every regression, " + "and complete a multi-step production audit." + ), + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert observed["during"]["counts"] == { + "started": 1, + "completed": 0, + "failed": 0, + "inflight": 1, + } + during_lab = observed["policy_lab_during"] + during_lifecycle = during_lab["evaluation_readiness"]["lifecycle_integrity"] + assert during_lab["evidence_source"] == "durable_sqlite_ledger" + assert during_lifecycle == { + "durable_lifecycle_present": True, + "durable_started_routes": 1, + "durable_terminal_routes": 0, + "durable_inflight_routes": 1, + "replay_usage_routes": 1, + "replay_decision_coverage_rate": 1.0, + "reconciled": False, + } + assert "lifecycle_not_reconciled" in during_lab["promotion_gate"]["blocking_reason_codes"] + report = manager.route_policy_ledger.report(session_id="failed-route-session") + assert report["counts"] == {"started": 1, "completed": 0, "failed": 1, "inflight": 0} + failed_lab = manager.route_policy_lab_snapshot("failed-route-session", profile="balanced") + failed_lifecycle = failed_lab["evaluation_readiness"]["lifecycle_integrity"] + assert failed_lab["durable_ledger"]["counts"] == { + "started": 1, + "completed": 0, + "failed": 1, + "inflight": 0, + } + assert failed_lifecycle == { + "durable_lifecycle_present": True, + "durable_started_routes": 1, + "durable_terminal_routes": 1, + "durable_inflight_routes": 0, + "replay_usage_routes": 1, + "replay_decision_coverage_rate": 1.0, + "reconciled": True, + } + assert "lifecycle_not_reconciled" not in failed_lab["promotion_gate"]["blocking_reason_codes"] + decision = manager.route_policy_ledger.list_decisions(session_id="failed-route-session")[0] + assert decision["status"] == "failed" + assert decision["success"] is False + assert decision["error_category"] == "runtime_error" + assert decision["error_message"] == "simulated route backend failure" + assert set(decision["actual_economics"]) == {"elapsed_ms"} + assert 0.0 <= decision["actual_economics"]["elapsed_ms"] < 5_000.0 + failure_outcome_names = [event["outcome_name"] for event in decision["outcome_events"]] + assert failure_outcome_names.count("route_success") == 1 + assert failure_outcome_names.count("cost") == 1 + assert failure_outcome_names.count("latency") == 1 + assert "user_quality_rating" not in failure_outcome_names + failure_events = {event["outcome_name"]: event for event in decision["outcome_events"]} + assert failure_events["route_success"]["value"] is False + assert failure_events["route_success"]["event_source"] == "route_completion" + assert failure_events["cost"]["observation_status"] == "not_observed" + assert failure_events["cost"]["value"] is None + assert failure_events["latency"]["value"] == pytest.approx( + decision["actual_economics"]["elapsed_ms"] + ) + assert manager.memory_store.load_session("failed-route-session").get("route_usage") == [] + with pytest.raises(ValueError, match="successfully completed route"): + manager.record_route_feedback( + session_id="failed-route-session", + feedback={"route_id": decision["route_id"], "rating": "down", "feedback_intent": "bad_quality"}, + ) + + +def test_auto_agent_budget_fast_caps_complex_prompt_to_collective(tmp_path: Path, monkeypatch) -> None: records = ( _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), ) manager = UnifiedModelManager( records=records, extraction_root=tmp_path / "extract", generated_dir=tmp_path / "generated", ) - counters = {"planner": 0, "worker": 0, "review": 0} + captured = {} - def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): - if "planner sub-agent" in prompt: - counters["planner"] += 1 - return ChatResult( - kind="text", - model_key=record.key, - model_label=record.label, - route_reason=route_reason, - response=( - "DONE: no\n" - "STEP_GOAL: produce the strongest final answer\n" - "SUCCESS_SIGNAL: the user request is fully handled\n" - "WORKING_NOTES: tighten the answer and check completeness" - ), - prompt_used=prompt, - ), [] - if "worker sub-agent" in prompt: - counters["worker"] += 1 - return ChatResult( - kind="text", - model_key=record.key, - model_label=record.label, - route_reason=route_reason, - response=( - f"DONE: {'yes' if counters['worker'] >= 2 else 'no'}\n" - f"OUTPUT: draft answer pass {counters['worker']}\n" - "NEXT_FOCUS: close the remaining gaps" - ), - prompt_used=prompt, - ), [] - counters["review"] += 1 + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] return ChatResult( kind="text", - model_key=record.key, - model_label=record.label, - route_reason=route_reason, - response=( - f"COMPLETE: {'yes' if counters['review'] >= 2 else 'no'}\n" - f"FINAL_RESPONSE: final answer pass {counters['review']}\n" - "REASON: the loop has converged on a complete response\n" - "NEXT_STEP: none" - ), - prompt_used=prompt, - ), [] + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response="auto routed through fast collective", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) - monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) payload = manager.handle_prompt( - session_id="loop-session", - prompt="Finish this task autonomously.", + session_id="auto-fast-session", + prompt=( + "Research the latest evidence, design and implement a multi-step runtime integration, " + "debug regressions, verify with tests, and explain benchmark tradeoffs." + ), model_key="omni_collective_v8", action_mode="text", settings={ - "agent_mode": "loop", - "loop_max_steps": 4, + "agent_mode": "auto", + "auto_agent_budget": "fast", "memory_enabled": False, "web_search_enabled": False, "cmd_open_enabled": False, }, ) - assert payload["model_key"] == "omni_collective_v8" - assert payload["agent_trace"]["agent_mode"] == "loop_agent" - assert payload["agent_trace"]["loop_completed"] is True - assert payload["agent_trace"]["loop_budget"] == 4 - assert len(payload["agent_trace"]["loop_steps"]) == 2 - assert "final answer pass 2" in payload["response"] - assert counters == {"planner": 2, "worker": 2, "review": 2} + policy = payload["agent_trace"]["auto_agent_policy"] + assert captured["settings"]["agent_mode"] == "collective" + assert policy["budget_profile"] == "fast" + assert policy["selected_agent_mode"] == "collective" + assert "loop" not in policy["allowed_agent_modes"] + assert policy["score"] < policy["score_before_budget"] + assert "fast budget profile" in payload["route_reason"] -def test_collective_loop_agent_uses_collective_worker(tmp_path: Path, monkeypatch) -> None: +def test_auto_agent_budget_max_promotes_moderate_prompt_to_collective_loop(tmp_path: Path, monkeypatch) -> None: records = ( _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), @@ -311,112 +798,3494 @@ def test_collective_loop_agent_uses_collective_worker(tmp_path: Path, monkeypatc extraction_root=tmp_path / "extract", generated_dir=tmp_path / "generated", ) - collective_calls = {"count": 0} - - def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): - if "planner sub-agent" in prompt: - response = ( - "DONE: no\n" - "STEP_GOAL: consult the panel and produce the final answer\n" - "SUCCESS_SIGNAL: the task is fully addressed\n" - "WORKING_NOTES: use the panel once and stop if complete" - ) - else: - response = ( - "COMPLETE: yes\n" - "FINAL_RESPONSE: collective loop final\n" - "REASON: the task is complete\n" - "NEXT_STEP: none" - ) - return ChatResult( - kind="text", - model_key=record.key, - model_label=record.label, - route_reason=route_reason, - response=response, - prompt_used=prompt, - ), [] + captured = {} - def fake_run_agent_text(*, session_id, prompt, chosen_record, settings, route_reason, action_mode, memory_bundle): - collective_calls["count"] += 1 + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] return ChatResult( kind="text", model_key=chosen_record.key, model_label=chosen_record.label, - route_reason=route_reason, - response="DONE: yes\nOUTPUT: panel-produced answer\nNEXT_FOCUS: none", - prompt_used=prompt, - agent_trace={ - "agent_mode": "collective_panel", - "consulted_models": ["omni_collective_v8", "v40_benchmax"], - "consultation_rows": [{"model_key": "v40_benchmax", "model_label": "v40_benchmax", "response": "panel note"}], - "tool_events": [], - }, + route_reason=kwargs["route_reason"], + response="auto routed through max collective loop", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_loop_agent", "loop_steps": []}, ) - monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) - monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) payload = manager.handle_prompt( - session_id="collective-loop-session", - prompt="Solve this with the collective loop.", + session_id="auto-max-session", + prompt="Research latest runtime integration and verify tests.", model_key="omni_collective_v8", action_mode="text", settings={ - "agent_mode": "collective_loop", - "loop_max_steps": 3, + "agent_mode": "auto", + "auto_agent_budget": "max", "memory_enabled": False, "web_search_enabled": False, "cmd_open_enabled": False, }, ) - assert payload["agent_trace"]["agent_mode"] == "collective_loop_agent" - assert payload["agent_trace"]["loop_worker_mode"] == "collective" - assert payload["agent_trace"]["loop_completed"] is True - assert collective_calls["count"] == 1 - assert payload["agent_trace"]["consulted_models"] == ["omni_collective_v8", "v40_benchmax"] + policy = payload["agent_trace"]["auto_agent_policy"] + assert captured["collective_mode"] is True + assert captured["settings"]["agent_mode"] == "collective_loop" + assert policy["budget_profile"] == "max" + assert policy["selected_agent_mode"] == "collective_loop" + assert policy["score"] > policy["score_before_budget"] + assert "max budget profile" in payload["route_reason"] -def test_collective_agent_skips_broken_consultant_and_still_returns(tmp_path: Path, monkeypatch) -> None: +def test_auto_agent_mode_keeps_simple_prompt_single_pass(tmp_path: Path, monkeypatch) -> None: records = ( _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), - _record("qwen_v28", "qwen_adapter", ("chat",), 0.42), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), ) manager = UnifiedModelManager( records=records, extraction_root=tmp_path / "extract", generated_dir=tmp_path / "generated", ) - - def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): - if record.key == "qwen_v28": - raise FileNotFoundError("missing adapter_model.safetensors") - return ChatResult( - kind="text", - model_key=record.key, - model_label=record.label, - route_reason=route_reason, - response="safe panel answer", - prompt_used=prompt, - ), [] - - monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) payload = manager.handle_prompt( - session_id="collective-session", - prompt="Answer despite one broken consultant.", + session_id="auto-simple-session", + prompt="Say hello.", model_key="omni_collective_v8", action_mode="text", settings={ - "agent_mode": "collective", + "agent_mode": "auto", "memory_enabled": False, "web_search_enabled": False, "cmd_open_enabled": False, }, ) - assert payload["response"] == "safe panel answer" + policy = payload["agent_trace"]["auto_agent_policy"] + assert payload["agent_trace"]["agent_mode"] == "off" + assert payload["agent_trace"]["requested_agent_mode"] == "auto" + assert payload["agent_trace"]["resolved_agent_mode"] == "off" + assert policy["selected_agent_mode"] == "off" + assert policy["reason"] == "low_complexity_single_pass" + assert policy["route_confidence"]["level"] == "high" + assert policy["uncertainty_adjustment"] is None + economics = payload["agent_trace"]["route_economics"] + assert economics["estimate"]["selected_agent_mode"] == "off" + assert economics["estimate"]["estimated_model_calls"] == 1 + assert economics["actual"]["model_calls"] == 1 + assert economics["actual"]["elapsed_ms"] >= 0 + assert payload["timing"]["route_elapsed_ms"] == economics["actual"]["elapsed_ms"] + assert policy["route_economics_actual"]["model_calls"] == 1 + assert payload["response"] == "reply via omni_collective_v8" + + +def test_auto_agent_uncertainty_margin_escalates_near_threshold(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 1, "cycles": 1, "reasons": ["borderline"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + + payload = manager.handle_prompt( + session_id="uncertainty-margin-session", + prompt="Audit this answer for correctness risk.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["uncertainty_adjustment"] + assert captured["settings"]["agent_mode"] == "collective" + assert policy["selected_agent_mode"] == "collective" + assert policy["reason"] == "borderline_score_with_uncertainty_signals" + assert adjustment["direction"] == "upgrade" + assert adjustment["from"] == "off" + assert adjustment["to"] == "collective" + assert adjustment["score_to_next_mode"] == 1 + assert {"audit", "correctness", "risk"}.issubset(set(adjustment["uncertainty_signals"])) + assert policy["route_confidence"]["adjusted"] is True + assert policy["route_confidence"]["score_to_next_mode"] == 1 + assert policy["route_confidence"]["selected_agent_mode"] == "collective" + assert "borderline_score_with_uncertainty_signals" in payload["route_reason"] + + +def test_auto_agent_uncertainty_margin_respects_fast_budget(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 2, "cycles": 2, "reasons": ["borderline"]}, + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="uncertainty-fast-budget-session", + prompt="Audit this answer for correctness risk.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": "fast", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert payload["agent_trace"]["resolved_agent_mode"] == "off" + assert policy["budget_profile"] == "fast" + assert policy["selected_agent_mode"] == "off" + assert policy["uncertainty_adjustment"] is None + assert policy["route_confidence"]["adjusted"] is False + assert policy["route_confidence"]["score_to_next_mode"] == 1 + assert policy["score"] < policy["score_before_budget"] + + +def test_auto_agent_session_budget_paces_route_to_remaining_budget(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="session-budget-pacing", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + + payload = manager.handle_prompt( + session_id="session-budget-pacing", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 8, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert captured["settings"]["agent_mode"] == "collective" + assert policy["selected_agent_mode"] == "collective" + assert payload["agent_trace"]["resolved_agent_mode"] == "collective" + assert adjustment["direction"] == "downgrade" + assert adjustment["from"] == "collective_loop" + assert adjustment["to"] == "collective" + assert adjustment["remaining_cost_units"] == 3.0 + assert policy["session_budget"]["limit_cost_units"] == 8.0 + assert policy["session_budget"]["used_cost_units"] == 5.0 + assert policy["session_budget"]["estimated_cost_units"] <= policy["session_budget"]["remaining_cost_units"] + assert policy["route_confidence"]["session_budget_adjusted"] is True + assert policy["route_economics_estimate"]["selected_agent_mode"] == "collective" + assert "Session budget pacing downgraded collective_loop to collective" in payload["route_reason"] + + health = manager.route_health_snapshot("session-budget-pacing") + assert health["route_usage"]["total_routes"] == 2 + assert health["route_usage"]["economics"]["total_cost_units"] == 8.0 + + +def test_auto_agent_session_budget_target_routes_paces_early_route(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + + payload = manager.handle_prompt( + session_id="session-budget-target-routes", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 12, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert captured["settings"]["agent_mode"] == "collective" + assert policy["selected_agent_mode"] == "collective" + assert adjustment["reason"] == "session_route_budget_target_pacing" + assert adjustment["from"] == "collective_loop" + assert adjustment["to"] == "collective" + assert adjustment["target_route_count"] == 4 + assert adjustment["target_remaining_routes"] == 4 + assert adjustment["pacing_cap_cost_units"] == 3.0 + assert policy["session_budget"]["target_route_count"] == 4 + assert policy["session_budget"]["pacing_cap_cost_units"] == 3.0 + assert policy["session_budget"]["effective_cap_cost_units"] == 3.0 + assert policy["session_budget"]["would_exceed_pacing_cap"] is False + assert policy["route_economics_estimate"]["estimated_cost_units"] <= 3.0 + assert "Session budget target pacing downgraded collective_loop to collective" in payload["route_reason"] + + +def test_auto_agent_session_budget_target_routes_preserves_remaining_budget_precedence(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="session-budget-target-hard-remaining", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="session-budget-target-hard-remaining", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 5.5, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert policy["selected_agent_mode"] == "off" + assert adjustment["reason"] == "session_route_budget_would_exceed_remaining" + assert adjustment["to"] == "off" + assert adjustment["remaining_cost_units"] == 0.5 + assert policy["session_budget"]["would_exceed_remaining"] is True + assert policy["session_budget"]["would_exceed_pacing_cap"] is True + assert "Session budget pacing downgraded collective_loop to off" in payload["route_reason"] + + +def test_auto_agent_session_budget_target_routes_reports_single_pass_pacing_floor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="session-budget-target-single-pass-floor", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 1.5, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert policy["selected_agent_mode"] == "off" + assert adjustment["reason"] == "session_route_budget_target_pacing" + assert adjustment["to"] == "off" + assert adjustment["pacing_cap_cost_units"] == 0.375 + assert policy["session_budget"]["would_exceed_remaining"] is False + assert policy["session_budget"]["would_exceed_pacing_cap"] is True + assert policy["route_economics_estimate"]["estimated_cost_units"] == 1.0 + assert "Session budget target pacing downgraded collective_loop to off" in payload["route_reason"] + + +def test_route_plan_preview_does_not_run_inference_or_write_memory(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + def fail(*args, **kwargs): + raise AssertionError("preview must not run inference or write memory") + + monkeypatch.setattr(manager, "ensure_backend", fail) + monkeypatch.setattr(manager, "_run_text_model", fail) + monkeypatch.setattr(manager, "_run_agent_text", fail) + monkeypatch.setattr(manager, "_run_loop_agent_text", fail) + monkeypatch.setattr(manager.memory_store, "add_route_usage", fail) + monkeypatch.setattr(manager.memory_store, "update", fail) + + plan = manager.preview_route_plan( + session_id="route-plan-dry-run", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert plan["dry_run"] is True + assert plan["execution_plan"]["will_run_inference"] is False + assert plan["execution_plan"]["will_write_memory"] is False + assert plan["auto_agent_policy"]["selected_agent_mode"] == plan["selected_agent_mode"] + assert manager.route_health_snapshot("route-plan-dry-run")["route_usage"]["total_routes"] == 0 + assert manager.route_policy_ledger.report(session_id="route-plan-dry-run")["counts"]["started"] == 0 + + +def test_route_study_preview_uses_final_budget_support_without_assigning_or_writing( + tmp_path: Path, + monkeypatch, +) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + def fail(*args, **kwargs): + raise AssertionError("study rehearsal must not infer, assign, or write") + + monkeypatch.setattr(manager, "ensure_backend", fail) + monkeypatch.setattr(manager, "_run_text_model", fail) + monkeypatch.setattr(manager, "_run_agent_text", fail) + monkeypatch.setattr(manager, "_run_loop_agent_text", fail) + monkeypatch.setattr(manager.memory_store, "add_route_usage", fail) + monkeypatch.setattr(manager.memory_store, "update", fail) + + prompt = "Research, implement, and verify a private production integration." + session_id = "route-study-private-session" + result = manager.preview_route_study( + session_id=session_id, + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 3, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + exploration_rate=0.10, + planned_routes=2000, + assumed_feedback_rate=0.30, + target_observed_labels=20, + ) + + assert result["dry_run"] is True + assert result["baseline_agent_mode"] == "collective" + assert result["execution_plan"] == { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + } + study = result["route_study"] + assert study["study"]["study_id"] == "auto-route-adjacent-explorer-v1" + assert study["charter"]["source_contract"] == result["deterministic_support"] + assert study["charter"]["enrollment"]["adjacent_feasible_actions"] == ["off"] + assert study["charter"]["probability_design"]["action_probabilities"] == { + "off": 0.1, + "collective": 0.9, + } + assert study["charter"]["probability_design"]["assignment_performed"] is False + label_forecast = study["charter"]["traffic_scenario"]["observed_label_scenario"] + assert label_forecast["target_scope"] == ( + "at_least_target_observed_labels_on_every_alternate_action" + ) + assert label_forecast["exact_simultaneous_target"]["method"] == ( + "exact_binomial_tail_inversion_single_alternate" + ) + assert study["charter"]["causal_boundaries"]["execution_enabled"] is False + assert study["charter"]["causal_boundaries"]["automatic_promotion_allowed"] is False + assert study["charter"]["causal_boundaries"]["activation_blockers"] + protocol = result["route_protocol_preflight"] + assert result["route_protocol_preflight_reason"] == "draft_for_independent_review" + assert protocol["protocol"]["label"] == "Stateful Route Experiment Preflight v1" + assert protocol["protocol"]["activation_available"] is False + assert protocol["charter"]["source_studies"]["support_strata"][0][ + "study_design_hash" + ] == study["design_hash"] + assert protocol["charter"]["stateful_design"]["selected_design_mode"] == ( + "sticky_session_cluster" + ) + assert protocol["charter"]["stateful_design"]["selected_design_status"] == ( + "declaration_incomplete" + ) + assert protocol["charter"]["randomness"]["assignment_performed"] is False + assert protocol["charter"]["causal_boundaries"]["activation_blockers"] + encoded = json.dumps(study, sort_keys=True) + assert prompt not in encoded + assert session_id not in encoded + encoded_protocol = json.dumps(protocol, sort_keys=True) + assert prompt not in encoded_protocol + assert session_id not in encoded_protocol + assert manager.route_health_snapshot(session_id)["route_usage"]["total_routes"] == 0 + assert manager.route_policy_ledger.report(session_id=session_id)["counts"]["started"] == 0 + + +def test_route_study_preview_rejects_manual_agent_mode(tmp_path: Path) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + with pytest.raises(ValueError, match="requires Auto Router"): + manager.preview_route_study( + session_id="manual-study", + prompt="Preview this route.", + model_key="omni_collective_v8", + action_mode="text", + settings={"agent_mode": "off"}, + ) + + +def test_route_study_protocol_preflight_binds_stateful_design_declarations( + tmp_path: Path, + monkeypatch, +) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + result = manager.preview_route_study( + session_id="stateful-preflight-session", + prompt="Research a difficult private integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 3, + "memory_enabled": False, + }, + target_policy_profile="quality_first", + protocol_design_mode="clustered_switchback", + carryover_scope="within_session", + interference_scope="none_declared", + temporal_variation="stable_declared", + planned_clusters=480, + max_routes_per_cluster=12, + analysis_every_clusters=60, + block_length_routes=20, + washout_routes=4, + ) + + protocol = result["route_protocol_preflight"] + assert protocol["charter"]["target_policy_class"]["profile_name"] == "quality_first" + assert protocol["charter"]["population"]["planned_clusters"] == 480 + assert protocol["charter"]["stopping_and_resources"]["planned_route_ceiling"] == 5760 + screen = protocol["charter"]["stateful_design"] + assert screen["selected_design_mode"] == "clustered_switchback" + assert screen["selected_design_status"] == "assumptions_declared_unvalidated" + assert screen["assignment_unit"] == "session_hash_x_time_block" + assert screen["block_length_routes"] == 20 + assert screen["washout_routes"] == 4 + assert protocol["charter"]["causal_boundaries"]["activation_available"] is False + + +def test_multi_stratum_review_bundle_is_pure_source_bound_and_fully_reconstructed( + tmp_path: Path, + monkeypatch, +) -> None: + from source.route_policy_protocol_cli import _example_bundle_input + + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + before_ledger = manager.route_policy_ledger.report() + + def fail(*_args, **_kwargs): + raise AssertionError("review-bundle assembly must not enter route or inference paths") + + monkeypatch.setattr(manager, "handle_prompt", fail) + monkeypatch.setattr(manager, "preview_route_plan", fail) + result = manager.build_route_protocol_review_bundle(_example_bundle_input()) + + bundle = result["route_protocol_review_bundle"] + verification = result["verification"] + assert bundle["schema_version"] == "route-study-review-bundle-v1" + assert verification["verification_level"] == "full_source_bound_reconstruction" + assert verification["support_stratum_count"] == 2 + assert verification["source_plan_reconstruction_performed"] is True + assert result["execution_plan"] == { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + } + assert manager.route_policy_ledger.report() == before_ledger + assert manager.audit_route_protocol_review_bundle(bundle)["verification"]["ok"] is True + + with pytest.raises(ValueError, match="non-prompt-free fields: prompt"): + manager.build_route_protocol_review_bundle( + {**_example_bundle_input(), "prompt": "must never enter the bundle"} + ) + + +def test_route_plan_preview_and_execution_share_budget_filtered_support_hashes( + tmp_path: Path, + monkeypatch, +) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="route-support-preview-execution", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + + settings = { + "agent_mode": "auto", + "auto_session_budget_units": 8, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + } + prompt = "Research, implement, and verify a complex production integration." + preview = manager.preview_route_plan( + session_id="route-support-preview-execution", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings=settings, + ) + + def fake_run_agent_text(**kwargs): + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response="budget-filtered collective response", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + executed = manager.handle_prompt( + session_id="route-support-preview-execution", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings=settings, + ) + + preview_policy = preview["auto_agent_policy"] + executed_policy = executed["agent_trace"]["auto_agent_policy"] + preview_support = preview_policy["logging_support"] + executed_support = executed_policy["logging_support"] + preview_exclusions = { + row["action"]: row["reasons"] for row in preview_support["exclusions"] + } + + assert preview["selected_agent_mode"] == "collective" + assert executed["agent_trace"]["resolved_agent_mode"] == "collective" + assert [row["action"] for row in preview_support["candidates"]] == ["off", "collective"] + assert preview_exclusions == { + "loop": ["session_budget_post_filter"], + "collective_loop": ["session_budget_post_filter"], + } + assert executed_support["candidates"] == preview_support["candidates"] + assert executed_support["exclusions"] == preview_support["exclusions"] + assert executed_policy["candidate_set_hash"] == preview_policy["candidate_set_hash"] + assert executed_policy["distribution_hash"] == preview_policy["distribution_hash"] + + ledger_row = manager.route_policy_ledger.get_decision(executed["route_id"]) + assert ledger_row["candidate_set_hash"] == preview_policy["candidate_set_hash"] + assert ledger_row["distribution_hash"] == preview_policy["distribution_hash"] + + +def test_route_plan_preview_is_idempotent_for_session_budget_usage(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="route-plan-budget-idempotent", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + + settings = { + "agent_mode": "auto", + "auto_session_budget_units": 8, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + } + first = manager.preview_route_plan( + session_id="route-plan-budget-idempotent", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings=settings, + ) + second = manager.preview_route_plan( + session_id="route-plan-budget-idempotent", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings=settings, + ) + + usage = manager.memory_store.route_usage_summary("route-plan-budget-idempotent") + assert usage["total_routes"] == 1 + assert usage["economics"]["total_cost_units"] == 5.0 + assert first["auto_agent_policy"]["session_budget"]["used_cost_units"] == 5.0 + assert second["auto_agent_policy"]["session_budget"]["used_cost_units"] == 5.0 + assert first["route_economics_estimate"]["selected_agent_mode"] == first["selected_agent_mode"] + assert second["route_economics_estimate"]["selected_agent_mode"] == second["selected_agent_mode"] + + +def test_route_plan_preview_sanitizes_bad_numeric_settings(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id="route-plan-bad-settings", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "loop_max_steps": "bad", + "web_search_enabled": True, + "web_search_budget": -5, + "web_search_results": "bad", + "auto_session_budget_units": "bad", + "auto_session_budget_target_routes": "bad", + "memory_enabled": False, + "cmd_open_enabled": False, + }, + ) + + estimate = plan["route_economics_estimate"] + assert plan["ok"] is True + assert estimate["estimated_tool_calls"] == 0 + assert estimate["planned_loop_steps"] == 4 + assert estimate["estimated_cost_units"] >= 0.0 + assert plan["auto_agent_policy"]["session_budget_adjustment"] is None + for row in plan["route_alternatives"]: + row_estimate = row["estimate"] + assert row_estimate["estimated_tool_calls"] >= 0 + assert 0 <= row_estimate["planned_loop_steps"] <= runtime_module.LOOP_AGENT_HARD_MAX_STEPS + assert row_estimate["estimated_cost_units"] >= 0.0 + + +def test_route_plan_preview_reports_allowed_route_alternatives(tmp_path: Path, monkeypatch) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id="route-plan-alternatives", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": "balanced", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + alternatives = plan["route_alternatives"] + modes = [row["selected_agent_mode"] for row in alternatives] + selected = [row for row in alternatives if row["is_selected"]] + assert modes == ["off", "collective", "loop", "collective_loop"] + assert selected == [next(row for row in alternatives if row["selected_agent_mode"] == plan["selected_agent_mode"])] + assert all(row["estimated_cost_units"] == row["estimate"]["estimated_cost_units"] for row in alternatives) + assert alternatives[0]["estimated_cost_units"] < alternatives[-1]["estimated_cost_units"] + assert all(row["frontier_rank"] >= 1 for row in alternatives) + assert all(0.0 <= row["estimated_quality_score"] <= 1.0 for row in alternatives) + assert all(row["estimated_quality_cost_score"] is not None for row in alternatives) + assert all(row["quality_cost_source"] == "heuristic_cost_adjusted" for row in alternatives) + assert all(row["quality_evidence_status"] == "heuristic_prior" for row in alternatives) + assert all(row["budget_fit"] is True for row in alternatives) + assert all(row["budget_feasible_pareto_frontier"] == row["pareto_frontier"] for row in alternatives) + assert plan["route_frontier"]["budget_feasible_pareto_modes"] == plan["route_frontier"]["pareto_modes"] + assert plan["route_frontier"]["recommended_agent_mode"] == plan["selected_agent_mode"] + assert plan["route_frontier"]["selected_matches_recommendation"] is True + assert plan["route_frontier"]["budget_fit_count"] == len(alternatives) + assert plan["route_frontier"]["budget_blockers"]["none"] == len(alternatives) + assert plan["route_frontier"]["selected_budget_blocker"] is None + assert plan["route_frontier"]["ranked_modes"][0]["frontier_rank"] == 1 + + +def test_route_plan_preview_alternatives_reflect_budget_adjusted_selection(tmp_path: Path, monkeypatch) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="route-plan-budget-frontier", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + + plan = manager.preview_route_plan( + session_id="route-plan-budget-frontier", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 8, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + alternatives = plan["route_alternatives"] + selected_rows = [row for row in alternatives if row["is_selected"]] + expensive_row = next(row for row in alternatives if row["selected_agent_mode"] == "collective_loop") + assert plan["selected_agent_mode"] == "collective" + assert selected_rows == [next(row for row in alternatives if row["selected_agent_mode"] == "collective")] + assert selected_rows[0]["estimate"] == plan["route_economics_estimate"] + assert expensive_row["estimated_cost_units"] > selected_rows[0]["estimated_cost_units"] + assert selected_rows[0]["budget_fit"] is True + assert expensive_row["budget_fit"] is False + assert selected_rows[0]["fits_remaining_budget"] is True + assert selected_rows[0]["fits_pacing_cap"] is None + assert selected_rows[0]["budget_blocker"] is None + assert expensive_row["fits_remaining_budget"] is False + assert expensive_row["budget_blocker"] == "remaining_budget" + assert plan["route_frontier"]["budget_cap_cost_units"] == 3.0 + assert plan["route_frontier"]["remaining_cost_units"] == 3.0 + assert plan["route_frontier"]["pacing_cap_cost_units"] is None + assert plan["route_frontier"]["effective_cap_cost_units"] == 3.0 + assert plan["route_frontier"]["recommended_agent_mode"] == "collective" + assert plan["route_frontier"]["selected_matches_recommendation"] is True + ranked_collective_loop = next( + row for row in plan["route_frontier"]["ranked_modes"] if row["selected_agent_mode"] == "collective_loop" + ) + assert ranked_collective_loop["budget_fit"] is False + assert ranked_collective_loop["budget_blocker"] == "remaining_budget" + assert ranked_collective_loop["pareto_frontier"] is True + assert ranked_collective_loop["budget_feasible_pareto_frontier"] is False + assert "collective_loop" in plan["route_frontier"]["pareto_modes"] + assert "collective_loop" not in plan["route_frontier"]["budget_feasible_pareto_modes"] + + +def test_route_plan_preview_alternatives_report_single_pass_floor(tmp_path: Path, monkeypatch) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 0, "cycles": 0, "reasons": ["simple_prompt"]}, + ) + manager.memory_store.add_route_usage( + session_id="route-plan-single-pass-floor", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + + plan = manager.preview_route_plan( + session_id="route-plan-single-pass-floor", + prompt="Say hello.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 5.5, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + selected = next(row for row in plan["route_alternatives"] if row["is_selected"]) + assert plan["selected_agent_mode"] == "off" + assert plan["auto_agent_policy"]["session_budget_adjustment"]["reason"] == "session_route_budget_exhausted_single_pass_floor" + assert selected["selected_agent_mode"] == "off" + assert selected["estimated_cost_units"] == 1.0 + assert selected["budget_fit"] is False + assert selected["fits_remaining_budget"] is False + assert selected["fits_pacing_cap"] is None + assert selected["minimum_route_floor"] is True + assert selected["budget_blocker"] == "remaining_budget" + assert plan["route_frontier"]["budget_fit_count"] == 0 + assert plan["route_frontier"]["recommended_agent_mode"] == "off" + assert plan["route_frontier"]["recommended_reason"] == "no_budget_feasible_route" + assert plan["route_frontier"]["recommended_budget_blocker"] == "remaining_budget" + assert plan["route_frontier"]["selected_budget_blocker"] == "remaining_budget" + assert plan["route_frontier"]["budget_blockers"]["remaining_budget"] == len(plan["route_alternatives"]) + assert plan["route_frontier"]["minimum_route_floor"] is True + + +def test_route_plan_frontier_distinguishes_remaining_fit_from_pacing_fit(tmp_path: Path, monkeypatch) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id="route-plan-pacing-floor", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 1.5, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + selected = next(row for row in plan["route_alternatives"] if row["is_selected"]) + frontier = plan["route_frontier"] + assert plan["selected_agent_mode"] == "off" + assert plan["auto_agent_policy"]["session_budget_adjustment"]["reason"] == "session_route_budget_target_pacing" + assert selected["selected_agent_mode"] == "off" + assert selected["estimated_cost_units"] == 1.0 + assert selected["fits_remaining_budget"] is True + assert selected["fits_pacing_cap"] is False + assert selected["budget_fit"] is False + assert selected["budget_blocker"] == "pacing_cap" + assert frontier["remaining_cost_units"] == 1.5 + assert frontier["pacing_cap_cost_units"] == 0.375 + assert frontier["effective_cap_cost_units"] == 0.375 + assert frontier["pacing_cap_applied"] is True + assert frontier["selected_fits_remaining_budget"] is True + assert frontier["selected_fits_pacing_cap"] is False + assert frontier["selected_budget_blocker"] == "pacing_cap" + assert frontier["recommended_agent_mode"] == "off" + assert frontier["recommended_budget_blocker"] == "pacing_cap" + assert frontier["minimum_route_floor"] is False + assert frontier["recommended_reason"] == "no_pacing_feasible_route" + assert frontier["budget_blockers"]["pacing_cap"] == 1 + assert frontier["budget_blockers"]["remaining_budget"] == len(plan["route_alternatives"]) - 1 + assert frontier["budget_feasible_pareto_modes"] == [] + + +def test_route_plan_frontier_prefers_complete_adaptive_quality_cost_evidence(tmp_path: Path) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + alternatives = [ + {"selected_agent_mode": "collective", "estimated_cost_units": 2.0, "is_selected": False}, + {"selected_agent_mode": "loop", "estimated_cost_units": 8.0, "is_selected": True}, + ] + + frontier = manager._annotate_route_frontier( + alternatives=alternatives, + selected="loop", + action_mode="text", + auto_agent_policy={ + "score": 4, + "budget_profile": "balanced", + "feedback_summary": { + "mode_scores": { + "collective": { + "adaptive": { + "sample_count": 3, + "weighted_count": 2.5, + "weighted_net": 2.5, + "quality_score": 0.86, + "quality_cost_score": 0.78, + }, + "economics": {"sample_count": 3, "avg_cost_units": 2.0}, + }, + "loop": { + "adaptive": { + "sample_count": 3, + "weighted_count": 2.5, + "weighted_net": 2.5, + "quality_score": 0.94, + "quality_cost_score": 0.46, + }, + "economics": {"sample_count": 3, "avg_cost_units": 8.0}, + }, + } + }, + }, + ) + + collective = next(row for row in alternatives if row["selected_agent_mode"] == "collective") + ranked_collective = next(row for row in frontier["ranked_modes"] if row["selected_agent_mode"] == "collective") + assert collective["quality_source"] == "adaptive_feedback" + assert collective["quality_cost_source"] == "adaptive_feedback" + assert collective["quality_evidence_status"] == "adaptive_complete" + assert collective["estimated_quality_cost_score"] == 0.78 + assert frontier["recommended_agent_mode"] == "collective" + assert frontier["recommended_reason"] == "adaptive_quality_cost_frontier_recommended" + assert frontier["recommended_estimated_quality_cost_score"] == 0.78 + assert frontier["selected_estimated_quality_cost_score"] == 0.46 + assert ranked_collective["frontier_rank"] == 1 + assert ranked_collective["estimated_quality_cost_score"] == 0.78 + assert ranked_collective["quality_cost_source"] == "adaptive_feedback" + assert ranked_collective["quality_evidence_status"] == "adaptive_complete" + + +def test_route_plan_frontier_reports_incomplete_adaptive_cost_evidence_without_promoting_it( + tmp_path: Path, +) -> None: + manager = UnifiedModelManager( + records=_auto_route_records(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + alternatives = [ + {"selected_agent_mode": "collective", "estimated_cost_units": 2.0, "is_selected": False}, + {"selected_agent_mode": "loop", "estimated_cost_units": 8.0, "is_selected": True}, + ] + + frontier = manager._annotate_route_frontier( + alternatives=alternatives, + selected="loop", + action_mode="text", + auto_agent_policy={ + "score": 4, + "budget_profile": "max", + "feedback_summary": { + "mode_scores": { + "collective": { + "adaptive": { + "sample_count": 3, + "weighted_count": 2.5, + "weighted_net": 2.5, + "quality_score": 0.99, + "quality_cost_score": 0.95, + }, + "economics": {"sample_count": 0}, + } + } + }, + }, + ) + + collective = next(row for row in alternatives if row["selected_agent_mode"] == "collective") + assert collective["quality_source"] == "heuristic_policy" + assert collective["quality_cost_source"] == "heuristic_cost_adjusted" + assert collective["quality_evidence_status"] == "adaptive_incomplete_cost_evidence" + assert collective["quality_evidence"] is None + assert collective["estimated_quality_score"] == 0.7 + assert frontier["recommended_agent_mode"] == "loop" + assert frontier["ranked_modes"][0]["selected_agent_mode"] == "loop" + + +@pytest.mark.parametrize( + ("profile", "expected_modes"), + [ + ("fast", ["off", "collective"]), + ("balanced", ["off", "collective", "loop", "collective_loop"]), + ("deep", ["off", "collective", "loop", "collective_loop"]), + ("max", ["off", "collective", "loop", "collective_loop"]), + ], +) +def test_route_plan_preview_alternatives_follow_auto_budget_profile( + tmp_path: Path, + monkeypatch, + profile: str, + expected_modes: list[str], +) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id=f"route-plan-profile-{profile}", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": profile, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert [row["selected_agent_mode"] for row in plan["route_alternatives"]] == expected_modes + + +@pytest.mark.parametrize( + ("toggle_settings", "expected_modes"), + [ + ({"auto_agent_collective": False}, ["off", "loop"]), + ({"auto_agent_loop": False}, ["off", "collective"]), + ], +) +def test_route_plan_preview_alternatives_follow_auto_agent_toggles( + tmp_path: Path, + monkeypatch, + toggle_settings: dict[str, bool], + expected_modes: list[str], +) -> None: + records = _auto_route_records() + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + settings = { + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + } + settings.update(toggle_settings) + + plan = manager.preview_route_plan( + session_id="route-plan-toggle-frontier", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings=settings, + ) + + assert [row["selected_agent_mode"] for row in plan["route_alternatives"]] == expected_modes + assert plan["selected_agent_mode"] in expected_modes + + +def test_route_plan_preview_alternatives_skip_unavailable_collective_modes(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id="route-plan-no-collective-frontier", + prompt="Research, implement, and verify a complex integration.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert [row["selected_agent_mode"] for row in plan["route_alternatives"]] == ["off", "loop"] + + +def test_route_plan_preview_image_action_excludes_loop_alternatives(tmp_path: Path) -> None: + records = ( + _record("dcgan_v2_in_progress", "dcgan_image", ("image",), None), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + plan = manager.preview_route_plan( + session_id="route-plan-image-loop-fallback", + prompt="Generate a clean test image.", + model_key="dcgan_v2_in_progress", + action_mode="image", + settings={ + "agent_mode": "loop", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + modes = [row["selected_agent_mode"] for row in plan["route_alternatives"]] + assert plan["selected_agent_mode"] == "off" + assert "loop" not in modes + assert "collective_loop" not in modes + assert plan["execution_plan"]["loop_enabled"] is False + assert plan["route_economics_estimate"]["selected_agent_mode"] == "off" + + +def test_route_plan_preview_image_support_records_post_filter_exclusions( + tmp_path: Path, + monkeypatch, +) -> None: + manager = UnifiedModelManager( + records=(_record("dcgan_v2_in_progress", "dcgan_image", ("image",), None),), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + + plan = manager.preview_route_plan( + session_id="route-plan-image-support", + prompt="Generate and refine a complex production image.", + model_key="dcgan_v2_in_progress", + action_mode="image", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = plan["auto_agent_policy"] + support = policy["logging_support"] + exclusions = {row["action"]: row["reasons"] for row in support["exclusions"]} + assert plan["selected_agent_mode"] == "off" + assert [row["action"] for row in support["candidates"]] == ["off"] + assert exclusions == { + "collective": ["capability_or_policy_filter"], + "loop": ["action_mode_unsupported"], + "collective_loop": ["action_mode_unsupported"], + } + assert policy["eligible_actions"] == ["off"] + assert policy["post_filter_action_probabilities"] == {"off": 1.0} + assert support["probability_stage"] == "post_filter" + assert len(policy["candidate_set_hash"]) == 64 + assert len(policy["distribution_hash"]) == 64 + + +def test_route_plan_preview_manual_mode_skips_auto_policy_and_writes(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + plan = manager.preview_route_plan( + session_id="route-plan-manual-loop", + prompt="Iterate on this implementation.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "loop", + "loop_max_steps": "bad", + "auto_session_budget_units": 1.5, + "auto_session_budget_target_routes": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert plan["auto_agent_policy"] is None + assert plan["selected_agent_mode"] == "loop" + assert plan["execution_plan"]["loop_enabled"] is True + assert plan["execution_plan"]["will_write_memory"] is False + assert plan["route_economics_estimate"]["planned_loop_steps"] == 4 + assert [row["selected_agent_mode"] for row in plan["route_alternatives"]] == ["loop"] + assert plan["route_alternatives"][0]["is_selected"] is True + assert manager.route_health_snapshot("route-plan-manual-loop")["route_usage"]["total_routes"] == 0 + + +def test_auto_agent_session_budget_overrides_uncertainty_upgrade(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 1, "cycles": 1, "reasons": ["borderline"]}, + ) + manager.memory_store.add_route_usage( + session_id="uncertainty-budget-session", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="uncertainty-budget-session", + prompt="Audit this answer for correctness risk.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 5.5, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["uncertainty_adjustment"]["from"] == "off" + assert policy["uncertainty_adjustment"]["to"] == "collective" + assert policy["session_budget_adjustment"]["from"] == "collective" + assert policy["session_budget_adjustment"]["to"] == "off" + assert policy["selected_agent_mode"] == "off" + assert policy["route_confidence"]["session_budget_adjusted"] is True + assert payload["agent_trace"]["resolved_agent_mode"] == "off" + + +def test_auto_agent_session_budget_can_fall_back_to_single_pass_when_exhausted(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 5, "cycles": 5, "reasons": ["workflow_depth"]}, + ) + manager.memory_store.add_route_usage( + session_id="session-budget-exhausted", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="session-budget-exhausted", + prompt="Complex implementation task.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 5.5, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert payload["agent_trace"]["resolved_agent_mode"] == "off" + assert policy["selected_agent_mode"] == "off" + assert adjustment["from"] == "collective_loop" + assert adjustment["to"] == "off" + assert adjustment["remaining_cost_units"] == 0.5 + assert policy["session_budget"]["would_exceed_remaining"] is True + assert policy["route_economics_estimate"]["selected_agent_mode"] == "off" + assert payload["response"] == "reply via omni_collective_v8" + + +def test_auto_agent_session_budget_reports_single_pass_floor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 0, "cycles": 0, "reasons": ["simple_prompt"]}, + ) + manager.memory_store.add_route_usage( + session_id="session-budget-floor", + route_id="prior-route", + prompt="Prior expensive work.", + selected_agent_mode="collective_loop", + route_economics={"actual": {"elapsed_ms": 1000.0, "model_calls": 5, "cost_units": 5.0}}, + auto_agent_policy={"selected_agent_mode": "collective_loop", "score": 5}, + route_reason="prior", + model_key="omni_collective_v8", + ) + monkeypatch.setattr(manager, "_build_backend", lambda record: _FakeBackend(record)) + + payload = manager.handle_prompt( + session_id="session-budget-floor", + prompt="Say hello.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_session_budget_units": 5.5, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["session_budget_adjustment"] + assert policy["selected_agent_mode"] == "off" + assert adjustment["direction"] == "floor" + assert adjustment["from"] == "off" + assert adjustment["to"] == "off" + assert adjustment["reason"] == "session_route_budget_exhausted_single_pass_floor" + assert policy["session_budget"]["would_exceed_remaining"] is True + assert "Session budget is exhausted" in payload["route_reason"] + + +def test_auto_agent_route_feedback_is_session_scoped_and_clearable(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + calls = [] + + def fake_run_loop_agent_text(**kwargs): + calls.append(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = ( + "Research the latest evidence, design and implement a multi-step runtime integration, " + "debug regressions, verify with tests, and explain benchmark tradeoffs." + ) + + for idx in range(2): + feedback = manager.record_route_feedback( + session_id="feedback-session", + feedback={ + "route_id": f"route-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective_loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "collective_loop", "score": 5}, + "route_economics": { + "actual": { + "elapsed_ms": 100 + idx, + "model_calls": 8, + "tool_calls": 0, + "cost_units": 8.0, + } + }, + }, + ) + assert feedback["summary"]["mode_scores"]["collective_loop"]["net"] == -2 + health = manager.route_health_snapshot("feedback-session") + assert health["economics"]["sample_count"] == 2 + assert health["economics"]["avg_cost_units"] == 8.0 + + adjusted = manager.handle_prompt( + session_id="feedback-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + adjusted_policy = adjusted["agent_trace"]["auto_agent_policy"] + assert adjusted_policy["selected_agent_mode"] == "loop" + assert adjusted_policy["feedback_adjustment"]["direction"] == "downgrade" + assert calls[-1]["collective_mode"] is False + + clean = manager.handle_prompt( + session_id="clean-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + assert clean["agent_trace"]["auto_agent_policy"]["selected_agent_mode"] == "collective_loop" + assert calls[-1]["collective_mode"] is True + + manager.clear("feedback-session") + reset = manager.handle_prompt( + session_id="feedback-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + assert reset["agent_trace"]["auto_agent_policy"]["selected_agent_mode"] == "collective_loop" + assert calls[-1]["collective_mode"] is True + + +def test_auto_agent_explicit_needs_deeper_feedback_moves_one_route_up(tmp_path: Path) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + prompt = "Implement and verify the runtime integration." + manager.record_route_feedback( + session_id="needs-deeper-session", + feedback={ + "route_id": "collective-route", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "down", + "feedback_intent": "needs_deeper", + }, + ) + + summary = manager.memory_store.route_feedback_summary("needs-deeper-session", prompt) + selected, adjustment = manager._apply_auto_route_feedback( + selected="collective", + score=2, + feedback_summary=summary, + allowed_modes=("off", "collective", "loop", "collective_loop"), + ) + + assert selected == "loop" + assert adjustment["direction"] == "upgrade" + assert adjustment["reason"] == "explicit_feedback_requested_deeper_route" + assert summary["mode_scores"]["collective"]["quality_negative"] == 0 + + +def test_auto_agent_explicit_cost_feedback_moves_one_route_down_without_quality_regression(tmp_path: Path) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + prompt = "Implement and verify the runtime integration." + manager.record_route_feedback( + session_id="lower-cost-session", + feedback={ + "route_id": "loop-route", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "feedback_intent": "too_costly", + "route_economics": {"actual": {"elapsed_ms": 7000, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + + summary = manager.memory_store.route_feedback_summary("lower-cost-session", prompt) + selected, adjustment = manager._apply_auto_route_feedback( + selected="loop", + score=4, + feedback_summary=summary, + allowed_modes=("off", "collective", "loop", "collective_loop"), + ) + + assert selected == "collective" + assert adjustment["direction"] == "downgrade" + assert adjustment["reason"] == "explicit_feedback_requested_lower_cost_route" + assert summary["mode_scores"]["loop"]["adaptive"]["quality_score"] is None + + +def test_auto_agent_adaptive_feedback_downgrades_recent_regression_over_stale_positives(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration." + for idx in range(6): + manager.record_route_feedback( + session_id="adaptive-feedback-session", + feedback={ + "route_id": f"old-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 100.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-feedback-session", + feedback={ + "route_id": f"new-bad-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 120.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-feedback-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["feedback_adjustment"] + assert policy["selected_agent_mode"] == "collective" + assert captured["settings"]["agent_mode"] == "collective" + assert adjustment["direction"] == "downgrade" + assert adjustment["reason"] == "recent_weighted_feedback_regression" + assert adjustment["from"] == "loop" + assert adjustment["to"] == "collective" + assert adjustment["weighted_net"] < 0 + assert policy["feedback_summary"]["mode_scores"]["loop"]["net"] == 4 + assert policy["feedback_summary"]["mode_scores"]["loop"]["adaptive"]["regression_signal"] is True + assert "recent_weighted_feedback_regression" in payload["route_reason"] + + +def test_auto_agent_adaptive_feedback_ignores_unrelated_recent_fallback(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + unrelated_prompt = "Optimize the marketing dashboard." + target_prompt = "Implement a runtime integration." + for idx in range(6): + manager.record_route_feedback( + session_id="adaptive-unrelated-session", + feedback={ + "route_id": f"old-good-{idx}", + "prompt": unrelated_prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 100.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-unrelated-session", + feedback={ + "route_id": f"new-bad-{idx}", + "prompt": unrelated_prompt, + "selected_agent_mode": "loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 120.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-unrelated-session", + prompt=target_prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["used_recent_fallback"] is True + assert policy["feedback_summary"]["mode_scores"]["loop"]["adaptive"]["regression_signal"] is True + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_max_budget_still_honors_adaptive_quality_regression(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 2, "cycles": 2, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration." + for idx in range(6): + manager.record_route_feedback( + session_id="adaptive-max-session", + feedback={ + "route_id": f"old-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 100.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-max-session", + feedback={ + "route_id": f"new-bad-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 120.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-max-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": "max", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["budget_profile"] == "max" + assert policy["selected_agent_mode"] == "collective" + assert policy["feedback_adjustment"]["reason"] == "recent_weighted_feedback_regression" + assert captured["settings"]["agent_mode"] == "collective" + + +def test_auto_agent_adaptive_quality_cost_prefers_cheaper_neighbor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration and verify the tests." + + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-cheaper-neighbor-session", + feedback={ + "route_id": f"loop-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 7000.0, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + for idx in range(3): + manager.record_route_feedback( + session_id="adaptive-cheaper-neighbor-session", + feedback={ + "route_id": f"collective-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 900.0, "model_calls": 2, "cost_units": 2.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-cheaper-neighbor-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["feedback_adjustment"] + assert policy["selected_agent_mode"] == "collective" + assert captured["settings"]["agent_mode"] == "collective" + assert adjustment["direction"] == "downgrade" + assert adjustment["reason"] == "adaptive_quality_cost_preferred_neighbor" + assert adjustment["quality_cost_delta"] > 0 + assert policy["feedback_summary"]["mode_scores"]["collective"]["adaptive"]["quality_cost_score"] > ( + policy["feedback_summary"]["mode_scores"]["loop"]["adaptive"]["quality_cost_score"] + ) + + +def test_auto_agent_adaptive_quality_cost_requires_candidate_sample_floor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = "Implement a runtime integration and verify the tests." + for idx in range(3): + manager.record_route_feedback( + session_id="adaptive-sample-floor-session", + feedback={ + "route_id": f"loop-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 6500.0, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + manager.record_route_feedback( + session_id="adaptive-sample-floor-session", + feedback={ + "route_id": "collective-lucky-good", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 500.0, "model_calls": 1, "cost_units": 1.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-sample-floor-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["mode_scores"]["collective"]["adaptive"]["weighted_count"] == 1.0 + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_adaptive_quality_cost_ignores_missing_cost_evidence(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = "Implement a runtime integration and verify the tests." + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-missing-cost-session", + feedback={ + "route_id": f"loop-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 6500.0, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + for idx in range(3): + manager.record_route_feedback( + session_id="adaptive-missing-cost-session", + feedback={ + "route_id": f"collective-no-cost-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 800.0, "model_calls": 2}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-missing-cost-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["mode_scores"]["collective"]["economics"].get("avg_cost_units") is None + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_adaptive_quality_cost_can_promote_stronger_neighbor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 3, "cycles": 3, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = "Plan and implement a careful multi-step runtime change." + + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-stronger-neighbor-session", + feedback={ + "route_id": f"collective-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 8000.0, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-stronger-neighbor-session", + feedback={ + "route_id": f"loop-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 3}, + "route_economics": {"actual": {"elapsed_ms": 1800.0, "model_calls": 3, "cost_units": 3.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-stronger-neighbor-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["feedback_adjustment"] + assert policy["selected_agent_mode"] == "loop" + assert captured["settings"]["agent_mode"] == "loop" + assert adjustment["direction"] == "upgrade" + assert adjustment["reason"] == "adaptive_quality_cost_preferred_neighbor" + assert adjustment["quality_cost_delta"] > 0 + + +def test_auto_agent_max_budget_ignores_cost_only_adaptive_neighbor(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 2, "cycles": 2, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = "Implement a runtime integration." + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-max-cost-only-session", + feedback={ + "route_id": f"loop-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": {"actual": {"elapsed_ms": 8000.0, "model_calls": 8, "cost_units": 8.0}}, + }, + ) + for idx in range(2): + manager.record_route_feedback( + session_id="adaptive-max-cost-only-session", + feedback={ + "route_id": f"collective-good-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 900.0, "model_calls": 2, "cost_units": 2.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-max-cost-only-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": "max", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["budget_profile"] == "max" + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_adaptive_quality_cost_ignores_unrelated_recent_fallback(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + unrelated_prompt = "Summarize customer churn by segment." + target_prompt = "Implement a runtime integration and verify the tests." + for idx in range(3): + manager.record_route_feedback( + session_id="adaptive-cost-unrelated-session", + feedback={ + "route_id": f"collective-good-{idx}", + "prompt": unrelated_prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": {"actual": {"elapsed_ms": 900.0, "model_calls": 2, "cost_units": 2.0}}, + }, + ) + + payload = manager.handle_prompt( + session_id="adaptive-cost-unrelated-session", + prompt=target_prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["used_recent_fallback"] is True + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_economics_pressure_downgrades_without_positive_signal(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration and verify the tests." + + for idx, selected_mode in enumerate(("loop", "collective")): + manager.record_route_feedback( + session_id="economic-pressure-session", + feedback={ + "route_id": f"route-{idx}", + "prompt": prompt, + "selected_agent_mode": selected_mode, + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": selected_mode, "score": 4}, + "route_economics": { + "actual": { + "elapsed_ms": 9000 + idx, + "model_calls": 7, + "tool_calls": 0, + "cost_units": 7.0, + } + }, + }, + ) + + payload = manager.handle_prompt( + session_id="economic-pressure-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["feedback_adjustment"] + assert policy["selected_agent_mode"] == "collective" + assert captured["settings"]["agent_mode"] == "collective" + assert adjustment["direction"] == "downgrade" + assert adjustment["reason"] == "session_route_economics_exceeded_budget_health" + assert adjustment["economic_pressure"] == ["cost", "latency"] + assert adjustment["pressure_scope"] == "loop" + assert adjustment["sample_count"] == 1 + assert policy["feedback_summary"]["economics"]["avg_cost_units"] == 7.0 + assert policy["feedback_summary"]["mode_scores"]["loop"]["economics"]["avg_cost_units"] == 7.0 + assert "session_route_economics_exceeded_budget_health" in payload["route_reason"] + + +def test_auto_agent_economics_pressure_is_mode_scoped(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 2, "cycles": 2, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration and verify the tests." + + for idx in range(2): + manager.record_route_feedback( + session_id="mode-scoped-economic-pressure-session", + feedback={ + "route_id": f"loop-route-{idx}", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": { + "actual": { + "elapsed_ms": 12000 + idx, + "model_calls": 9, + "tool_calls": 0, + "cost_units": 9.0, + } + }, + }, + ) + manager.record_route_feedback( + session_id="mode-scoped-economic-pressure-session", + feedback={ + "route_id": "collective-route", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": { + "actual": { + "elapsed_ms": 1200, + "model_calls": 2, + "tool_calls": 0, + "cost_units": 2.0, + } + }, + }, + ) + + payload = manager.handle_prompt( + session_id="mode-scoped-economic-pressure-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "collective" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["mode_scores"]["loop"]["economics"]["avg_cost_units"] == 9.0 + assert policy["feedback_summary"]["mode_scores"]["collective"]["economics"]["avg_cost_units"] == 2.0 + assert captured["settings"]["agent_mode"] == "collective" + + +def test_auto_agent_economics_pressure_uses_selected_mode_not_aggregate(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_panel", "consulted_models": ["a", "b"]}, + ) + + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + prompt = "Implement a runtime integration and verify the tests." + + manager.record_route_feedback( + session_id="selected-mode-economic-pressure-session", + feedback={ + "route_id": "loop-route", + "prompt": prompt, + "selected_agent_mode": "loop", + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": "loop", "score": 4}, + "route_economics": { + "actual": { + "elapsed_ms": 9000, + "model_calls": 7, + "tool_calls": 0, + "cost_units": 7.0, + } + }, + }, + ) + for idx in range(6): + manager.record_route_feedback( + session_id="selected-mode-economic-pressure-session", + feedback={ + "route_id": f"collective-route-{idx}", + "prompt": prompt, + "selected_agent_mode": "collective", + "rating": "up", + "auto_agent_policy": {"selected_agent_mode": "collective", "score": 2}, + "route_economics": { + "actual": { + "elapsed_ms": 400, + "model_calls": 1, + "tool_calls": 0, + "cost_units": 1.0, + } + }, + }, + ) + + payload = manager.handle_prompt( + session_id="selected-mode-economic-pressure-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + adjustment = policy["feedback_adjustment"] + assert policy["selected_agent_mode"] == "collective" + assert adjustment["reason"] == "adaptive_quality_cost_preferred_neighbor" + assert adjustment["quality_cost_delta"] > 0 + assert policy["feedback_summary"]["economics"]["avg_cost_units"] < 6.0 + assert policy["feedback_summary"]["mode_scores"]["loop"]["economics"]["avg_cost_units"] == 7.0 + assert captured["settings"]["agent_mode"] == "collective" + + +def test_auto_agent_economics_pressure_ignores_unrelated_recent_feedback(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + unrelated_prompt = "Summarize the quarterly sales dashboard." + target_prompt = "Implement a runtime integration and verify the tests." + + for idx, selected_mode in enumerate(("loop", "collective")): + manager.record_route_feedback( + session_id="unrelated-economic-pressure-session", + feedback={ + "route_id": f"route-{idx}", + "prompt": unrelated_prompt, + "selected_agent_mode": selected_mode, + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": selected_mode, "score": 4}, + "route_economics": { + "actual": { + "elapsed_ms": 12000 + idx, + "model_calls": 9, + "tool_calls": 0, + "cost_units": 9.0, + } + }, + }, + ) + + payload = manager.handle_prompt( + session_id="unrelated-economic-pressure-session", + prompt=target_prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["selected_agent_mode"] == "loop" + assert policy["feedback_adjustment"] is None + assert policy["feedback_summary"]["used_recent_fallback"] is True + assert policy["feedback_summary"]["relevant_feedback"] == 0 + assert captured["settings"]["agent_mode"] == "loop" + + +def test_auto_agent_max_budget_ignores_passive_economics_pressure(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr( + runtime_module.chat_app, + "estimate_auto_reasoning_cycles", + lambda prompt: {"score": 4, "cycles": 4, "reasons": ["workflow_depth"]}, + ) + captured = {} + + def fake_run_loop_agent_text(**kwargs): + captured.update(kwargs) + chosen_record = kwargs["chosen_record"] + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=kwargs["route_reason"], + response=f"auto routed through {kwargs['settings']['agent_mode']}", + prompt_used=kwargs["prompt"], + agent_trace={"agent_mode": "collective_loop_agent", "loop_steps": []}, + ) + + monkeypatch.setattr(manager, "_run_loop_agent_text", fake_run_loop_agent_text) + prompt = "Implement a runtime integration and verify the tests." + + for idx, selected_mode in enumerate(("loop", "collective")): + manager.record_route_feedback( + session_id="max-economic-pressure-session", + feedback={ + "route_id": f"route-{idx}", + "prompt": prompt, + "selected_agent_mode": selected_mode, + "rating": "down", + "auto_agent_policy": {"selected_agent_mode": selected_mode, "score": 4}, + "route_economics": { + "actual": { + "elapsed_ms": 25000 + idx, + "model_calls": 15, + "tool_calls": 0, + "cost_units": 15.0, + } + }, + }, + ) + + payload = manager.handle_prompt( + session_id="max-economic-pressure-session", + prompt=prompt, + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "auto", + "auto_agent_budget": "max", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + policy = payload["agent_trace"]["auto_agent_policy"] + assert policy["budget_profile"] == "max" + assert policy["selected_agent_mode"] == "collective_loop" + assert policy["feedback_adjustment"] is None + assert captured["settings"]["agent_mode"] == "collective_loop" + + +def test_model_store_catalog_marks_installed_and_selectable_records(tmp_path: Path) -> None: + models_dir = tmp_path / "models" + models_dir.mkdir() + installed = models_dir / "dcgan_v2_in_progress.zip" + installed.write_bytes(b"zip") + records = ( + _record("dcgan_v2_in_progress", "dcgan_image", ("image",), None), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + models_dir=models_dir, + common_summary_path=tmp_path / "missing_summary.json", + ) + + manager._fetch_model_store_manifest_locked = lambda force_refresh=False: { # type: ignore[method-assign] + "models": [ + {"file_name": "dcgan_v2_in_progress.zip", "size_bytes": 3, "size_mb": 0.0, "family": "gan"}, + {"file_name": "supermix_omni_collective_v8_preview_20260407_001155.zip", "size_bytes": 10, "size_mb": 0.0, "family": "fusion"}, + ] + } + + payload = manager.model_store_catalog(force_refresh=True) + by_name = {row["file_name"]: row for row in payload["models"]} + assert by_name["dcgan_v2_in_progress.zip"]["installed"] is True + assert by_name["dcgan_v2_in_progress.zip"]["selectable"] is True + assert by_name["supermix_omni_collective_v8_preview_20260407_001155.zip"]["known"] is True + assert by_name["supermix_omni_collective_v8_preview_20260407_001155.zip"]["installed"] is False + + +def test_model_store_catalog_skips_unsafe_remote_manifest_names(tmp_path: Path) -> None: + models_dir = tmp_path / "models" + models_dir.mkdir() + manager = UnifiedModelManager( + records=(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + models_dir=models_dir, + common_summary_path=tmp_path / "missing_summary.json", + ) + + manager._fetch_model_store_manifest_locked = lambda force_refresh=False: { # type: ignore[method-assign] + "models": [ + {"file_name": "../escape.zip", "size_bytes": 10}, + {"file_name": "nested/escape.zip", "size_bytes": 10}, + {"file_name": "dcgan_v2_in_progress.zip", "size_bytes": 3}, + ] + } + + payload = manager.model_store_catalog(force_refresh=True) + + assert [row["file_name"] for row in payload["models"]] == ["dcgan_v2_in_progress.zip"] + assert not (tmp_path / "escape.zip").exists() + + +def test_model_store_install_rejects_unsafe_artifact_names(tmp_path: Path) -> None: + manager = UnifiedModelManager( + records=(), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + models_dir=tmp_path / "models", + common_summary_path=tmp_path / "missing_summary.json", + ) + manager._fetch_model_store_manifest_locked = lambda force_refresh=False: { # type: ignore[method-assign] + "models": [{"file_name": "dcgan_v2_in_progress.zip", "size_bytes": 3}] + } + + for bad_name in ("../escape.zip", "nested/escape.zip", "bad:name.zip", "not-a-zip.txt"): + with pytest.raises(ValueError): + manager.install_model_store_artifact(bad_name) + + +def test_loop_agent_runs_until_reviewer_marks_complete(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + counters = {"planner": 0, "worker": 0, "review": 0} + + def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): + if "planner sub-agent" in prompt: + counters["planner"] += 1 + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=( + "DONE: no\n" + "STEP_GOAL: produce the strongest final answer\n" + "SUCCESS_SIGNAL: the user request is fully handled\n" + "WORKING_NOTES: tighten the answer and check completeness" + ), + prompt_used=prompt, + ), [] + if "worker sub-agent" in prompt: + counters["worker"] += 1 + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=( + f"DONE: {'yes' if counters['worker'] >= 2 else 'no'}\n" + f"OUTPUT: draft answer pass {counters['worker']}\n" + "NEXT_FOCUS: close the remaining gaps" + ), + prompt_used=prompt, + ), [] + counters["review"] += 1 + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=( + f"COMPLETE: {'yes' if counters['review'] >= 2 else 'no'}\n" + f"SCORE: {0.97 if counters['review'] >= 2 else 0.45}\n" + f"CONFIDENCE: {0.96 if counters['review'] >= 2 else 0.40}\n" + "RISK_SCORE: 0.05\n" + f"FINAL_RESPONSE: final answer pass {counters['review']}\n" + "REASON: the loop has converged on a complete response\n" + "EVIDENCE: reviewer observed all requested work\n" + "NEXT_STEP: none" + ), + prompt_used=prompt, + ), [] + + monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + + payload = manager.handle_prompt( + session_id="loop-session", + prompt="Finish this task autonomously.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "loop", + "loop_max_steps": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert payload["model_key"] == "omni_collective_v8" + assert payload["agent_trace"]["agent_mode"] == "loop_agent" + assert payload["agent_trace"]["loop_completed"] is True + assert payload["agent_trace"]["loop_budget"] == 4 + assert len(payload["agent_trace"]["loop_steps"]) == 2 + assert payload["agent_trace"]["loop_steps"][0]["review_score"] == 0.45 + assert payload["agent_trace"]["loop_steps"][1]["review_score"] == 0.97 + assert payload["agent_trace"]["loop_steps"][1]["stop_decision"] == "stop" + assert payload["agent_trace"]["loop_stop_reason_code"] == "reviewer_complete" + assert payload["agent_trace"]["loop_stop_step"] == 2 + assert payload["agent_trace"]["loop_stop_score"] == 0.97 + economics = payload["agent_trace"]["route_economics"] + assert economics["estimate"]["selected_agent_mode"] == "loop" + assert economics["estimate"]["estimated_model_calls"] == 12 + assert economics["actual"]["model_calls"] == 6 + assert economics["actual"]["loop_steps"] == 2 + assert payload["timing"]["route_elapsed_ms"] == economics["actual"]["elapsed_ms"] + assert "final answer pass 2" in payload["response"] + assert counters == {"planner": 2, "worker": 2, "review": 2} + + +def test_loop_agent_records_budget_stop_telemetry(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): + if "planner sub-agent" in prompt: + response = ( + "DONE: no\n" + "STEP_GOAL: refine the incomplete answer\n" + "SUCCESS_SIGNAL: the answer is complete\n" + "WORKING_NOTES: continue iterating" + ) + elif "worker sub-agent" in prompt: + response = "DONE: no\nOUTPUT: partial answer\nNEXT_FOCUS: finish the remaining proof" + else: + response = ( + "COMPLETE: no\n" + "SCORE: 0.4\n" + "CONFIDENCE: 0.5\n" + "RISK_SCORE: 0.2\n" + "FINAL_RESPONSE: partial answer\n" + "REASON: still missing proof\n" + "EVIDENCE: reviewer found a remaining gap\n" + "NEXT_STEP: add the missing proof" + ) + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=response, + prompt_used=prompt, + ), [] + + monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + + payload = manager.handle_prompt( + session_id="loop-budget-session", + prompt="Finish this task autonomously.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "loop", + "loop_max_steps": 2, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert payload["agent_trace"]["loop_completed"] is False + assert payload["agent_trace"]["loop_stop_reason_code"] == "budget_exhausted" + assert payload["agent_trace"]["loop_stop_step"] == 2 + assert payload["agent_trace"]["loop_stop_score"] == 0.4 + assert payload["agent_trace"]["loop_steps"][-1]["stop_decision"] == "continue" + assert "loop budget ended" in payload["response"] + + +def test_loop_agent_can_stop_on_high_score_without_yes_no_label(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + calls = {"review": 0} + + def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): + if "planner sub-agent" in prompt: + response = ( + "DONE: no\n" + "STEP_GOAL: produce a final answer\n" + "SUCCESS_SIGNAL: the answer satisfies the request\n" + "WORKING_NOTES: one pass should be enough" + ) + elif "worker sub-agent" in prompt: + response = "DONE: yes\nOUTPUT: complete answer\nNEXT_FOCUS: none" + else: + calls["review"] += 1 + response = ( + "SCORE: 0.91\n" + "CONFIDENCE: 0.92\n" + "RISK_SCORE: 0.10\n" + "FINAL_RESPONSE: complete answer\n" + "REASON: score evidence clears the completion threshold\n" + "EVIDENCE: all requested parts are present\n" + "NEXT_STEP: none" + ) + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=response, + prompt_used=prompt, + ), [] + + monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + + payload = manager.handle_prompt( + session_id="loop-score-session", + prompt="Finish this task autonomously.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "loop", + "loop_max_steps": 4, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + step = payload["agent_trace"]["loop_steps"][0] + assert payload["agent_trace"]["loop_completed"] is True + assert payload["agent_trace"]["loop_stop_reason_code"] == "score_threshold" + assert payload["agent_trace"]["loop_stop_step"] == 1 + assert payload["agent_trace"]["loop_stop_score"] == 0.91 + assert step["review_complete"] is None + assert step["stop_decision"] == "stop" + assert calls["review"] == 1 + + +def test_collective_loop_agent_uses_collective_worker(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("v40_benchmax", "omni_collective_v5", ("chat", "vision"), 0.2433), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + collective_calls = {"count": 0} + + def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): + if "planner sub-agent" in prompt: + response = ( + "DONE: no\n" + "STEP_GOAL: consult the panel and produce the final answer\n" + "SUCCESS_SIGNAL: the task is fully addressed\n" + "WORKING_NOTES: use the panel once and stop if complete" + ) + else: + response = ( + "COMPLETE: yes\n" + "FINAL_RESPONSE: collective loop final\n" + "REASON: the task is complete\n" + "NEXT_STEP: none" + ) + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response=response, + prompt_used=prompt, + ), [] + + def fake_run_agent_text(*, session_id, prompt, chosen_record, settings, route_reason, action_mode, memory_bundle): + collective_calls["count"] += 1 + return ChatResult( + kind="text", + model_key=chosen_record.key, + model_label=chosen_record.label, + route_reason=route_reason, + response="DONE: yes\nOUTPUT: panel-produced answer\nNEXT_FOCUS: none", + prompt_used=prompt, + agent_trace={ + "agent_mode": "collective_panel", + "consulted_models": ["omni_collective_v8", "v40_benchmax"], + "consultation_rows": [{"model_key": "v40_benchmax", "model_label": "v40_benchmax", "response": "panel note"}], + "tool_events": [], + }, + ) + + monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + monkeypatch.setattr(manager, "_run_agent_text", fake_run_agent_text) + + payload = manager.handle_prompt( + session_id="collective-loop-session", + prompt="Solve this with the collective loop.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "collective_loop", + "loop_max_steps": 3, + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert payload["agent_trace"]["agent_mode"] == "collective_loop_agent" + assert payload["agent_trace"]["loop_worker_mode"] == "collective" + assert payload["agent_trace"]["loop_completed"] is True + assert collective_calls["count"] == 1 + assert payload["agent_trace"]["consulted_models"] == ["omni_collective_v8", "v40_benchmax"] + + +def test_collective_agent_skips_broken_consultant_and_still_returns(tmp_path: Path, monkeypatch) -> None: + records = ( + _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133), + _record("qwen_v28", "qwen_adapter", ("chat",), 0.42), + ) + manager = UnifiedModelManager( + records=records, + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + def fake_run_text_model(record, *, session_id, prompt, settings, route_reason, tool_cache, allow_tool_calls): + if record.key == "qwen_v28": + raise FileNotFoundError("missing adapter_model.safetensors") + return ChatResult( + kind="text", + model_key=record.key, + model_label=record.label, + route_reason=route_reason, + response="safe panel answer", + prompt_used=prompt, + ), [] + + monkeypatch.setattr(manager, "_run_text_model", fake_run_text_model) + + payload = manager.handle_prompt( + session_id="collective-session", + prompt="Answer despite one broken consultant.", + model_key="omni_collective_v8", + action_mode="text", + settings={ + "agent_mode": "collective", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + + assert payload["response"] == "safe panel answer" assert payload["agent_trace"]["agent_mode"] == "collective_panel" assert payload["agent_trace"]["consulted_models"] == ["omni_collective_v8"] assert payload["agent_trace"]["skipped_models"][0]["model_key"] == "qwen_v28" + + +def test_auto_route_feedback_ignores_unrelated_fallback_quality_votes(tmp_path: Path) -> None: + record = _record("omni_collective_v8", "omni_collective_v8", ("chat", "vision"), 0.2133) + manager = UnifiedModelManager( + records=(record,), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + allowed_modes = ("off", "collective", "loop", "collective_loop") + + downgraded, downgrade_adjustment = manager._apply_auto_route_feedback( + selected="loop", + score=4, + feedback_summary={ + "used_recent_fallback": True, + "mode_scores": { + "loop": { + "quality_net": -2, + "quality_negative": 2, + "adaptive": {}, + } + }, + }, + allowed_modes=allowed_modes, + ) + assert downgraded == "loop" + assert downgrade_adjustment is None + + upgraded, upgrade_adjustment = manager._apply_auto_route_feedback( + selected="collective", + score=4, + feedback_summary={ + "used_recent_fallback": True, + "mode_scores": { + "collective": {"adaptive": {}}, + "loop": {"quality_net": 2, "quality_positive": 2}, + }, + }, + allowed_modes=allowed_modes, + ) + assert upgraded == "collective" + assert upgrade_adjustment is None + + +def test_route_economics_counts_tool_followup_model_invocation(tmp_path: Path, monkeypatch) -> None: + record = _record("tool-model", "champion_chat", ("chat",), 0.9) + manager = UnifiedModelManager( + records=(record,), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + + class _ToolBackend(_FakeBackend): + def __init__(self, backend_record: ModelRecord) -> None: + super().__init__(backend_record) + self.calls = 0 + + def chat(self, session_id: str, prompt: str, settings: dict) -> ChatResult: + self.calls += 1 + response = "TOOL:web_search: current docs" if self.calls == 1 else "Answer grounded in the tool result." + return ChatResult( + kind="text", + model_key=self.record.key, + model_label=self.record.label, + route_reason=str(settings.get("route_reason") or ""), + response=response, + prompt_used=prompt, + ) + + backend = _ToolBackend(record) + monkeypatch.setattr(manager, "ensure_backend", lambda _key: (record, backend)) + + def fake_web_query(query, tool_cache, settings): + event = runtime_module.ToolEvent( + name="web_search", + query=query, + results=[{"title": "Docs", "url": "https://example.test/docs", "snippet": "Current."}], + ) + tool_cache[str(query).lower()] = event + return event + + monkeypatch.setattr(manager, "_run_web_query_cached", fake_web_query) + payload = manager.handle_prompt( + session_id="tool-economics-session", + prompt="Answer using the available evidence.", + model_key=record.key, + action_mode="text", + settings={ + "agent_mode": "off", + "memory_enabled": False, + "web_search_enabled": True, + "cmd_open_enabled": False, + }, + ) + + actual = payload["agent_trace"]["route_economics"]["actual"] + assert backend.calls == 2 + assert actual["model_calls"] == 2 + assert actual["tool_calls"] == 1 + assert actual["cost_units"] == 2.25 + + second_payload = manager.handle_prompt( + session_id="tool-economics-session", + prompt="Answer without a tool.", + model_key=record.key, + action_mode="text", + settings={ + "agent_mode": "off", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + }, + ) + assert backend.calls == 3 + assert second_payload["agent_trace"]["route_economics"]["actual"]["model_calls"] == 1 + + +def test_image_wrapper_counts_refiner_and_pipeline_in_route_economics(tmp_path: Path, monkeypatch) -> None: + record = _record("image-wrapper", "image_wrapper", ("chat", "image"), 0.9) + backend = runtime_module.ImageWrapperBackend.__new__(runtime_module.ImageWrapperBackend) + backend.record = record + + class _ImageEngine: + def generate_image(self, **kwargs): + return { + "timing_ms": 25.0, + "model_calls": 2, + "refiner_model_calls": 1, + "image_url": "/generated/test.png", + "output_path": str(tmp_path / "test.png"), + "prompt_used": kwargs["prompt"], + "refined_prompt": "refined image prompt", + } + + backend.image_engine = _ImageEngine() + manager = UnifiedModelManager( + records=(record,), + extraction_root=tmp_path / "extract", + generated_dir=tmp_path / "generated", + ) + monkeypatch.setattr(manager, "ensure_backend", lambda _key: (record, backend)) + + payload = manager.handle_prompt( + session_id="image-economics-session", + prompt="Paint a lunar observatory.", + model_key=record.key, + action_mode="image", + settings={ + "agent_mode": "off", + "memory_enabled": False, + "web_search_enabled": False, + "cmd_open_enabled": False, + "use_text_refiner": True, + }, + ) + + actual = payload["agent_trace"]["route_economics"]["actual"] + assert payload["timing"]["model_calls"] == 2 + assert payload["timing"]["refiner_model_calls"] == 1 + assert actual["model_calls"] == 2 + assert actual["cost_units"] == 2.0 + + +def test_image_prompt_refiner_marks_its_model_invocation(monkeypatch) -> None: + from source.chat_image_variant_app import ImageVariantEngine + + engine = ImageVariantEngine.__new__(ImageVariantEngine) + + class _TextEngine: + @staticmethod + def status(): + return {"loaded": True} + + engine.text_engine = _TextEngine() + monkeypatch.setattr(engine, "_refine_prompt_with_text_model", lambda _prompt: "A moon base under blue light") + + refined = engine._build_final_prompt("A moon base", style="auto", use_text_refiner=True) + unrefined = engine._build_final_prompt("A moon base", style="auto", use_text_refiner=False) + + assert refined["refiner_model_calls"] == 1 + assert unrefined["refiner_model_calls"] == 0 diff --git a/test_route_policy_shadow_web_status.py b/test_route_policy_shadow_web_status.py new file mode 100644 index 000000000..1e271f019 --- /dev/null +++ b/test_route_policy_shadow_web_status.py @@ -0,0 +1,241 @@ +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import torch + + +SOURCE_DIR = Path(__file__).resolve().parent / "source" +sys.path.insert(0, str(SOURCE_DIR)) + +import multimodel_runtime +import supermix_multimodel_web_app +from chat_web_app import Engine, build_app +from route_policy_shadow_registry import RouteShadowAssignmentRegistry + + +class _MemoryStoreStub: + def __init__(self, path: Path): + self.path = path + + +class _RouteLedgerStub: + def __init__(self, path: Path): + self.path = path + + +def _build_manager(monkeypatch, tmp_path: Path): + monkeypatch.setattr(multimodel_runtime, "configure_torch_runtime", lambda **_kwargs: None) + monkeypatch.setattr( + multimodel_runtime, + "resolve_device", + lambda *_args, **_kwargs: (torch.device("cpu"), {"resolved": "cpu"}), + ) + monkeypatch.setattr(multimodel_runtime, "ConversationMemoryStore", _MemoryStoreStub) + monkeypatch.setattr(multimodel_runtime, "RoutePolicyLedger", _RouteLedgerStub) + extraction_root = tmp_path / "tmp" / "ext" + return multimodel_runtime.UnifiedModelManager( + records=(), + extraction_root=extraction_root, + generated_dir=tmp_path / "tmp" / "gen", + models_dir=tmp_path / "models", + common_summary_path=tmp_path / "common.json", + ) + + +def test_unified_manager_uses_canonical_memory_path_without_creating_registry( + monkeypatch, tmp_path +): + manager = _build_manager(monkeypatch, tmp_path) + expected = (tmp_path / "tmp" / "memory" / "route-policy-shadow-registry.sqlite3").resolve() + + assert manager.route_shadow_registry_path == expected + assert not expected.exists() + snapshot = manager.route_shadow_registry_snapshot() + + assert snapshot == { + "ok": True, + "available": False, + "status": "not_initialized", + "registry_location": "memory/route-policy-shadow-registry.sqlite3", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + assert not expected.exists() + + +def test_unified_manager_snapshot_uses_registry_read_only_mode(monkeypatch, tmp_path): + manager = _build_manager(monkeypatch, tmp_path) + RouteShadowAssignmentRegistry(manager.route_shadow_registry_path) + before = manager.route_shadow_registry_path.read_bytes() + + snapshot = manager.route_shadow_registry_snapshot() + after = manager.route_shadow_registry_path.read_bytes() + + assert snapshot["ok"] is True + assert snapshot["available"] is True + assert snapshot["status"] == "verified" + assert snapshot["read_only"] is True + assert snapshot["campaign_count"] == 0 + assert before == after + + +def test_unified_manager_caches_verified_snapshot_until_durable_state_changes( + monkeypatch, tmp_path +): + manager = _build_manager(monkeypatch, tmp_path) + RouteShadowAssignmentRegistry(manager.route_shadow_registry_path) + calls = {"count": 0} + original_snapshot = RouteShadowAssignmentRegistry.snapshot + + def counted_snapshot(registry, *args, **kwargs): + calls["count"] += 1 + return original_snapshot(registry, *args, **kwargs) + + monkeypatch.setattr(RouteShadowAssignmentRegistry, "snapshot", counted_snapshot) + + first = manager.route_shadow_registry_snapshot() + first["campaigns"].append({"cache_poison": True}) + second = manager.route_shadow_registry_snapshot() + + assert calls["count"] == 1 + assert second["campaigns"] == [] + + registry_stat = manager.route_shadow_registry_path.stat() + newer = max(time.time_ns(), registry_stat.st_mtime_ns + 1_000_000_000) + os.utime(manager.route_shadow_registry_path, ns=(newer, newer)) + third = manager.route_shadow_registry_snapshot() + + assert calls["count"] == 2 + assert third["ok"] is True + assert third["read_only"] is True + + +class _CanonicalWebManager: + def route_shadow_registry_snapshot(self): + return { + "ok": True, + "available": False, + "status": "not_initialized", + "registry_location": "memory/route-policy-shadow-registry.sqlite3", + "read_only": True, + "campaign_count": 0, + "campaigns": [], + "event_chain": None, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + + +class _FailingCanonicalWebManager: + def route_shadow_registry_snapshot(self): + raise RuntimeError("registry audit failed") + + +def test_canonical_studio_registry_api_is_get_only_and_ui_is_read_only(): + app = supermix_multimodel_web_app.build_app(_CanonicalWebManager()) + client = app.test_client() + + response = client.get("/api/route_shadow_registry/status") + assert response.status_code == 200 + assert response.headers["Cache-Control"] == "no-store" + assert response.get_json()["route_shadow_registry"]["read_only"] is True + assert client.post("/api/route_shadow_registry/status", json={}).status_code == 405 + + html = client.get("/").get_data(as_text=True) + assert 'id="routeShadowRegistryRefresh"' in html + assert "/api/route_shadow_registry/status" in html + assert "Shadow assignment registry - read only" in html + assert "seal campaign" not in html.lower() + assert "reveal seed" not in html.lower() + + +def test_canonical_studio_registry_error_is_not_cached(): + app = supermix_multimodel_web_app.build_app(_FailingCanonicalWebManager()) + response = app.test_client().get("/api/route_shadow_registry/status") + + assert response.status_code == 500 + assert response.headers["Cache-Control"] == "no-store" + assert response.get_json() == {"ok": False, "error": "registry audit failed"} + + +def test_canonical_studio_defaults_to_loopback_binding(): + parser = supermix_multimodel_web_app.build_arg_parser() + + assert parser.parse_args([]).host == "127.0.0.1" + assert parser.parse_args(["--host", "0.0.0.0"]).host == "0.0.0.0" + + +def test_source_legacy_browser_registry_status_is_get_only(tmp_path): + registry_path = tmp_path / "missing-shadow.sqlite3" + engine = Engine( + torch.device("cpu"), + {"resolved": "cpu"}, + {"pool_mode": "all", "route_shadow_registry_path": str(registry_path)}, + ) + client = build_app(engine, "weights.pth", "meta.json").test_client() + + response = client.get("/api/route_shadow_registry/status") + assert response.status_code == 200 + assert response.headers["Cache-Control"] == "no-store" + payload = response.get_json()["route_shadow_registry"] + assert payload["available"] is False + assert payload["read_only"] is True + assert not registry_path.exists() + assert client.post("/api/route_shadow_registry/status", json={}).status_code == 405 + assert "Shadow registry - read only" in client.get("/").get_data(as_text=True) + + +def test_packaged_legacy_browser_registry_status_is_get_only(tmp_path): + runtime_dir = Path("runtime_python").resolve() + script = f""" +import json +import sys +from pathlib import Path +import torch +sys.path.insert(0, {str(runtime_dir)!r}) +import chat_web_app + +registry_path = Path({str(tmp_path / 'missing-runtime-shadow.sqlite3')!r}) +engine = chat_web_app.Engine( + torch.device('cpu'), + {{'resolved': 'cpu'}}, + {{'pool_mode': 'all', 'route_shadow_registry_path': str(registry_path)}}, +) +client = chat_web_app.build_app(engine, 'weights.pth', 'meta.json').test_client() +response = client.get('/api/route_shadow_registry/status') +print(json.dumps({{ + 'status_code': response.status_code, + 'cache_control': response.headers.get('Cache-Control'), + 'payload': response.get_json(), + 'post_status': client.post('/api/route_shadow_registry/status', json={{}}).status_code, + 'created': registry_path.exists(), + 'ui': 'Shadow registry - read only' in client.get('/').get_data(as_text=True), +}}, sort_keys=True)) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path.cwd(), + check=True, + capture_output=True, + text=True, + env={**os.environ, "PYTHONPATH": ""}, + ) + payload = json.loads(completed.stdout.strip().splitlines()[-1]) + + assert payload["status_code"] == 200 + assert payload["cache_control"] == "no-store" + assert payload["payload"]["route_shadow_registry"]["available"] is False + assert payload["payload"]["route_shadow_registry"]["read_only"] is True + assert payload["post_status"] == 405 + assert payload["created"] is False + assert payload["ui"] is True diff --git a/test_runtime_compute_controls.py b/test_runtime_compute_controls.py new file mode 100644 index 000000000..135931939 --- /dev/null +++ b/test_runtime_compute_controls.py @@ -0,0 +1,738 @@ +import importlib.util +import json +import os +import subprocess +import sys +from argparse import Namespace +from pathlib import Path + +import torch + + +sys.path.insert(0, os.path.join(os.getcwd(), "source")) + +import chat_app +from chat_web_app import Engine, build_app + + +class _RuntimeComputeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = [] + self.register_buffer("last_cycles_used", torch.tensor(0.0)) + self.register_buffer("last_ponder_cost", torch.tensor(0.0)) + self.register_buffer("last_consistency_loss", torch.tensor(0.0)) + self.register_buffer("last_gating_entropy", torch.tensor(0.0)) + self.register_buffer("last_prediction_streak", torch.tensor(0.0)) + self.register_buffer("last_prediction_confidence_delta", torch.tensor(0.0)) + self.last_exit_reason = "not_run" + + def forward( + self, + x, + reasoning_cycles=None, + adaptive_compute=False, + exit_tol=1e-3, + exit_entropy_threshold=0.2, + prediction_stability_patience=2, + prediction_stability_tol=5e-3, + ): + self.calls.append( + { + "reasoning_cycles": reasoning_cycles, + "adaptive_compute": adaptive_compute, + "exit_tol": exit_tol, + "exit_entropy_threshold": exit_entropy_threshold, + "prediction_stability_patience": prediction_stability_patience, + "prediction_stability_tol": prediction_stability_tol, + } + ) + self.last_cycles_used = torch.tensor(float(reasoning_cycles or 3)) + self.last_ponder_cost = torch.tensor(float(reasoning_cycles or 3) - 0.25) + self.last_consistency_loss = torch.tensor(0.125) + self.last_gating_entropy = torch.tensor(0.75) + self.last_prediction_streak = torch.tensor(float(prediction_stability_patience)) + self.last_prediction_confidence_delta = torch.tensor(float(prediction_stability_tol)) + self.last_exit_reason = "prediction_stable" if adaptive_compute and prediction_stability_patience else "max_cycles" + logits = torch.zeros(x.shape[0], x.shape[1], 10, device=x.device) + logits[..., 2] = 4.0 + return logits + + +class _LegacyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, x): + self.calls += 1 + return torch.zeros(x.shape[0], x.shape[1], 10, device=x.device) + + +def _candidate(text: str): + vec = chat_app.text_to_model_input(text, feature_mode="legacy")[0, 0].tolist() + return { + "text": text, + "vec": vec, + "ctx_vec": vec, + "bucket_score": 1.0, + "count": 1, + } + + +def _stub_engine_load(monkeypatch, metadata_by_name): + monkeypatch.setattr( + chat_app, + "load_metadata", + lambda path: metadata_by_name[Path(path).name], + ) + monkeypatch.setattr(chat_app, "safe_load_state_dict", lambda _path: {}) + monkeypatch.setattr( + chat_app, + "detect_model_size_from_state_dict", + lambda _state: "cognitive_leap_expert", + ) + monkeypatch.setattr( + chat_app, + "build_model", + lambda **_kwargs: _RuntimeComputeModel(), + ) + monkeypatch.setattr( + chat_app, + "load_weights_for_model", + lambda _model, _state, model_size: ([], []), + ) + + +def _write_load_files(tmp_path: Path, *meta_names: str): + weights = tmp_path / "weights.pth" + weights.write_bytes(b"stub") + metas = [] + for name in meta_names: + path = tmp_path / name + path.write_text("{}", encoding="utf-8") + metas.append(path) + return weights, metas + + +def smoke_test_runtime_compute_helper_applies_only_supported_kwargs(): + x = torch.zeros(1, 1, 128) + recursive = _RuntimeComputeModel() + + _, diag = chat_app.forward_with_runtime_compute( + recursive, + x, + reasoning_cycles="128", + adaptive_compute="on", + exit_tol="0.25", + exit_entropy_threshold="0.75", + prediction_stability_patience="4", + prediction_stability_tol="0.02", + ) + + assert recursive.calls[-1] == { + "reasoning_cycles": chat_app.MAX_RUNTIME_REASONING_CYCLES, + "adaptive_compute": True, + "exit_tol": 0.25, + "exit_entropy_threshold": 0.75, + "prediction_stability_patience": 4, + "prediction_stability_tol": 0.02, + } + assert diag["applied"] is True + assert diag["requested_reasoning_cycles"] == chat_app.MAX_RUNTIME_REASONING_CYCLES + assert diag["cycles_used"] == float(chat_app.MAX_RUNTIME_REASONING_CYCLES) + assert diag["ponder_cost"] == float(chat_app.MAX_RUNTIME_REASONING_CYCLES) - 0.25 + assert diag["consistency_loss"] == 0.125 + assert diag["gating_entropy"] == 0.75 + assert diag["exit_entropy_threshold"] == 0.75 + assert diag["prediction_streak"] == 4.0 + assert abs(diag["prediction_confidence_delta"] - 0.02) < 1e-6 + assert diag["exit_reason"] == "prediction_stable" + + legacy = _LegacyModel() + _, legacy_diag = chat_app.forward_with_runtime_compute( + legacy, + x, + reasoning_cycles=8, + adaptive_compute=True, + exit_tol=0.5, + ) + + assert legacy.calls == 1 + assert legacy_diag["applied"] is False + assert legacy_diag["cycles_used"] is None + + +def smoke_test_auto_reasoning_budget_selects_cycles_from_context(): + x = torch.zeros(1, 1, 128) + recursive = _RuntimeComputeModel() + + _, diag = chat_app.forward_with_runtime_compute( + recursive, + x, + reasoning_cycles="auto", + adaptive_compute=False, + auto_reasoning_context=( + "Debug this traceback, inspect the runtime integration, compare benchmarks, " + "and verify the fix with tests." + ), + ) + + assert diag["reasoning_budget_mode"] == "auto" + assert diag["auto_reasoning_policy"]["cycles"] == 16 + assert "code_or_debug" in diag["auto_reasoning_policy"]["reasons"] + assert recursive.calls[-1]["reasoning_cycles"] == 16 + + +def smoke_test_web_engine_forwards_runtime_compute_controls(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("Runtime compute response.")]} + engine.available_labels = [2] + + result = engine.chat( + session_id="s", + user_text="use more reasoning", + reasoning_cycles=5, + adaptive_compute=True, + adaptive_exit_tol=0.125, + adaptive_exit_entropy=0.5, + prediction_stability_patience=3, + prediction_stability_tol=0.01, + ) + + assert result["ok"] is True + assert result["compute"]["applied"] is True + assert result["compute"]["cycles_used"] == 5.0 + assert model.calls[-1] == { + "reasoning_cycles": 5, + "adaptive_compute": True, + "exit_tol": 0.125, + "exit_entropy_threshold": 0.5, + "prediction_stability_patience": 3, + "prediction_stability_tol": 0.01, + } + + +def smoke_test_web_engine_auto_runtime_compute_controls(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("Auto compute response.")]} + engine.available_labels = [2] + + result = engine.chat( + session_id="auto", + user_text="debug this runtime error and verify the benchmark integration with tests", + reasoning_cycles="auto", + ) + + assert result["ok"] is True + assert result["compute"]["reasoning_budget_mode"] == "auto" + assert result["compute"]["selected_reasoning_cycles"] == 16 + assert model.calls[-1]["reasoning_cycles"] == 16 + + +def smoke_test_web_api_accepts_runtime_compute_payload(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("API compute response.")]} + engine.available_labels = [2] + + app = build_app(engine, "weights.pth", "meta.json") + client = app.test_client() + response = client.post( + "/api/chat", + json={ + "session_id": "api", + "message": "think harder", + "reasoning_cycles": 7, + "adaptive_compute": True, + "adaptive_exit_tol": 0.05, + "adaptive_exit_entropy": 0.4, + "prediction_stability_patience": 5, + "prediction_stability_tol": 0.03, + }, + ) + + assert response.status_code == 200, response.get_data(as_text=True) + payload = response.get_json() + assert payload["compute"]["applied"] is True + assert payload["compute"]["cycles_used"] == 7.0 + assert model.calls[-1] == { + "reasoning_cycles": 7, + "adaptive_compute": True, + "exit_tol": 0.05, + "exit_entropy_threshold": 0.4, + "prediction_stability_patience": 5, + "prediction_stability_tol": 0.03, + } + + +def smoke_test_web_api_accepts_auto_runtime_compute_payload(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("API auto response.")]} + engine.available_labels = [2] + + app = build_app(engine, "weights.pth", "meta.json") + client = app.test_client() + response = client.post( + "/api/chat", + json={ + "session_id": "api-auto", + "message": "analyze this traceback, compare benchmarks, verify tests, and fix runtime integration", + "reasoning_cycles": "auto", + }, + ) + + assert response.status_code == 200, response.get_data(as_text=True) + payload = response.get_json() + assert payload["compute"]["reasoning_budget_mode"] == "auto" + assert payload["compute"]["selected_reasoning_cycles"] == 16 + assert model.calls[-1]["reasoning_cycles"] == 16 + + +def smoke_test_compute_sweep_reports_budget_rows_without_mutating_session(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("Sweep response.")]} + engine.available_labels = [2] + engine.sessions["sweep"] = [("prior", "answer")] + + before_history = list(engine.sessions["sweep"]) + result = engine.compute_sweep( + session_id="sweep", + user_text="compare compute budgets", + cycles=[1, "3", 999], + adaptive_compute=True, + adaptive_exit_tol=0.2, + adaptive_exit_entropy=0.6, + prediction_stability_patience=6, + prediction_stability_tol=0.04, + ) + + assert result["ok"] is True + assert [row["requested_cycles"] for row in result["rows"]] == [ + 1, + 3, + chat_app.MAX_RUNTIME_REASONING_CYCLES, + ] + assert [row["cycles_used"] for row in result["rows"]] == [ + 1.0, + 3.0, + float(chat_app.MAX_RUNTIME_REASONING_CYCLES), + ] + assert all(row["predicted_label"] == 2 for row in result["rows"]) + assert engine.sessions["sweep"] == before_history + assert model.calls[-1] == { + "reasoning_cycles": chat_app.MAX_RUNTIME_REASONING_CYCLES, + "adaptive_compute": True, + "exit_tol": 0.2, + "exit_entropy_threshold": 0.6, + "prediction_stability_patience": 6, + "prediction_stability_tol": 0.04, + } + + +def smoke_test_web_api_accepts_compute_sweep_payload(): + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + model = _RuntimeComputeModel() + with engine.lock: + engine.model = model + engine.feature_mode = "legacy" + engine.model_size = "cognitive_leap_expert" + engine.buckets = {2: [_candidate("API sweep response.")]} + engine.available_labels = [2] + + app = build_app(engine, "weights.pth", "meta.json") + client = app.test_client() + response = client.post( + "/api/compute_sweep", + json={ + "session_id": "api-sweep", + "message": "sweep this prompt", + "cycles": [1, 2], + "adaptive_compute": True, + "adaptive_exit_tol": 0.15, + "adaptive_exit_entropy": 0.55, + "prediction_stability_patience": 4, + "prediction_stability_tol": 0.025, + }, + ) + + assert response.status_code == 200, response.get_data(as_text=True) + payload = response.get_json() + assert payload["ok"] is True + assert [row["requested_cycles"] for row in payload["rows"]] == [1, 2] + assert payload["rows"][0]["compute"]["applied"] is True + assert payload["rows"][1]["compute"]["exit_entropy_threshold"] == 0.55 + assert payload["rows"][1]["compute"]["prediction_stability_patience"] == 4 + assert abs(payload["rows"][1]["compute"]["prediction_stability_tol"] - 0.025) < 1e-6 + assert "api-sweep" not in engine.sessions + + +def test_engine_load_applies_whitelisted_metadata_compute_defaults_and_request_overrides( + tmp_path: Path, + monkeypatch, +): + metadata = { + "feature_mode": "legacy", + "model_size": "cognitive_leap_expert", + "runtime_defaults": { + "reasoning_cycles": "7", + "adaptive_compute": "on", + "adaptive_exit_tol": "0.25", + "adaptive_exit_entropy": "0.75", + "prediction_stability_patience": "4", + "prediction_stability_tol": "0.02", + "pool_mode": "topk", + "untrusted_extra": "ignored", + }, + "buckets": {"2": [_candidate("Metadata compute response.")]}, + } + _stub_engine_load(monkeypatch, {"meta.json": metadata}) + weights, (meta_path,) = _write_load_files(tmp_path, "meta.json") + engine = Engine( + torch.device("cpu"), + {"resolved": "cpu"}, + {"model_size": "auto", "pool_mode": "all", "max_turns": 2}, + ) + + loaded = engine.load(str(weights), str(meta_path)) + + assert loaded["reasoning_cycles"] == 7 + assert loaded["adaptive_compute"] is True + assert loaded["adaptive_exit_tol"] == 0.25 + assert loaded["adaptive_exit_entropy"] == 0.75 + assert loaded["prediction_stability_patience"] == 4 + assert loaded["prediction_stability_tol"] == 0.02 + assert engine.defaults["pool_mode"] == "all" + assert "untrusted_extra" not in engine.defaults + + engine.chat(session_id="meta-defaults", user_text="use metadata defaults") + assert engine.model.calls[-1] == { + "reasoning_cycles": 7, + "adaptive_compute": True, + "exit_tol": 0.25, + "exit_entropy_threshold": 0.75, + "prediction_stability_patience": 4, + "prediction_stability_tol": 0.02, + } + + engine.chat( + session_id="request-overrides", + user_text="override every compute default", + reasoning_cycles=5, + adaptive_compute=True, + adaptive_exit_tol=0.125, + adaptive_exit_entropy=0.5, + prediction_stability_patience=3, + prediction_stability_tol=0.01, + ) + assert engine.model.calls[-1] == { + "reasoning_cycles": 5, + "adaptive_compute": True, + "exit_tol": 0.125, + "exit_entropy_threshold": 0.5, + "prediction_stability_patience": 3, + "prediction_stability_tol": 0.01, + } + + disabled = engine.chat( + session_id="request-disables-adaptive", + user_text="disable adaptive compute for this request", + adaptive_compute=False, + ) + assert disabled["compute"]["adaptive_compute"] is False + assert engine.model.calls[-1]["adaptive_compute"] is False + + +def test_engine_load_rebuilds_defaults_and_constructor_values_win( + tmp_path: Path, + monkeypatch, +): + first_meta = { + "feature_mode": "legacy", + "model_size": "cognitive_leap_expert", + "runtime_defaults": { + "reasoning_cycles": 8, + "adaptive_compute": True, + "adaptive_exit_tol": 0.4, + "adaptive_exit_entropy": 0.8, + "prediction_stability_patience": 6, + "prediction_stability_tol": 0.04, + }, + } + second_meta = { + "feature_mode": "legacy", + "model_size": "cognitive_leap_expert", + "runtime_defaults": {"adaptive_compute": False}, + } + metadata_by_name = {"first.json": first_meta, "second.json": second_meta} + _stub_engine_load(monkeypatch, metadata_by_name) + weights, (first_path, second_path) = _write_load_files( + tmp_path, + "first.json", + "second.json", + ) + + engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, {"pool_mode": "all"}) + engine.load(str(weights), str(first_path)) + assert engine.status()["reasoning_cycles"] == 8 + assert engine.status()["prediction_stability_patience"] == 6 + + reloaded = engine.load(str(weights), str(second_path)) + assert reloaded["reasoning_cycles"] == "default" + assert reloaded["adaptive_compute"] is False + assert reloaded["adaptive_exit_tol"] == 1e-3 + assert reloaded["adaptive_exit_entropy"] == chat_app.DEFAULT_ADAPTIVE_EXIT_ENTROPY + assert reloaded["prediction_stability_patience"] == chat_app.DEFAULT_PREDICTION_STABILITY_PATIENCE + assert reloaded["prediction_stability_tol"] == chat_app.DEFAULT_PREDICTION_STABILITY_TOL + + explicit = { + "pool_mode": "all", + "reasoning_cycles": 3, + "adaptive_compute": False, + "adaptive_exit_tol": 0.1, + "adaptive_exit_entropy": 0.3, + "prediction_stability_patience": 2, + "prediction_stability_tol": 0.005, + } + explicit_engine = Engine(torch.device("cpu"), {"resolved": "cpu"}, explicit) + explicit_status = explicit_engine.load(str(weights), str(first_path)) + assert explicit_status["reasoning_cycles"] == 3 + assert explicit_status["adaptive_compute"] is False + assert explicit_status["adaptive_exit_tol"] == 0.1 + assert explicit_status["adaptive_exit_entropy"] == 0.3 + assert explicit_status["prediction_stability_patience"] == 2 + assert explicit_status["prediction_stability_tol"] == 0.005 + + +def test_cli_omits_unspecified_compute_defaults_so_metadata_can_win(): + import chat_web_app + + unspecified = Namespace( + reasoning_cycles=None, + adaptive_compute=None, + adaptive_exit_tol=None, + adaptive_exit_entropy=None, + prediction_stability_patience=None, + prediction_stability_tol=None, + ) + assert chat_web_app._runtime_compute_cli_overrides(unspecified) == {} + + explicit = Namespace( + reasoning_cycles="auto", + adaptive_compute=True, + adaptive_exit_tol=0.02, + adaptive_exit_entropy=0.4, + prediction_stability_patience=4, + prediction_stability_tol=0.01, + ) + assert chat_web_app._runtime_compute_cli_overrides(explicit) == { + "reasoning_cycles": "auto", + "adaptive_compute": True, + "adaptive_exit_tol": 0.02, + "adaptive_exit_entropy": 0.4, + "prediction_stability_patience": 4, + "prediction_stability_tol": 0.01, + } + + explicit_disabled = Namespace(adaptive_compute=False) + assert chat_web_app._runtime_compute_cli_overrides(explicit_disabled) == { + "adaptive_compute": False, + } + + +def test_runtime_compute_smoke_suite(): + smoke_test_runtime_compute_helper_applies_only_supported_kwargs() + smoke_test_auto_reasoning_budget_selects_cycles_from_context() + smoke_test_web_engine_forwards_runtime_compute_controls() + smoke_test_web_engine_auto_runtime_compute_controls() + smoke_test_web_api_accepts_runtime_compute_payload() + smoke_test_web_api_accepts_auto_runtime_compute_payload() + smoke_test_compute_sweep_reports_budget_rows_without_mutating_session() + smoke_test_web_api_accepts_compute_sweep_payload() + + +def test_packaged_context_mix_v4_matches_source_runtime(): + runtime_path = Path("runtime_python/chat_pipeline.py").resolve() + spec = importlib.util.spec_from_file_location("supermix_runtime_chat_pipeline", runtime_path) + assert spec is not None and spec.loader is not None + runtime_pipeline = importlib.util.module_from_spec(spec) + spec.loader.exec_module(runtime_pipeline) + + context = ( + "system: reasoning_budget=deep\n" + "system: knowledge_recency=latest\n" + "system: source_quality=research\n" + "user: continue and improve the newest arxiv research implementation" + ) + source_vector = chat_app.text_to_model_input(context, feature_mode="context_mix_v4") + runtime_vector = runtime_pipeline.text_to_model_input(context, feature_mode="context_mix_v4") + + assert runtime_pipeline.resolve_feature_mode("context_mix_v3", smarter_auto=True) == "context_mix_v4" + assert torch.allclose(source_vector, runtime_vector) + + +def test_packaged_web_runtime_rebuilds_metadata_compute_defaults(): + runtime_dir = Path("runtime_python").resolve() + script = f""" +import json +import sys +import torch +sys.path.insert(0, {str(runtime_dir)!r}) +import chat_web_app + +metadata = {{ + "runtime_defaults": {{ + "reasoning_cycles": 8, + "adaptive_compute": True, + "adaptive_exit_tol": 0.4, + "prediction_stability_patience": 6, + }} +}} +plain = chat_web_app.Engine(torch.device("cpu"), {{"resolved": "cpu"}}, {{"pool_mode": "all"}}) +first = plain._build_effective_defaults(metadata) +clean = plain._build_effective_defaults({{"runtime_defaults": {{"adaptive_compute": False}}}}) +explicit = chat_web_app.Engine( + torch.device("cpu"), + {{"resolved": "cpu"}}, + {{"pool_mode": "all", "reasoning_cycles": 3, "adaptive_compute": False}}, +) +overridden = explicit._build_effective_defaults(metadata) +print(json.dumps({{"first": first, "clean": clean, "overridden": overridden}}, sort_keys=True)) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path.cwd(), + check=True, + capture_output=True, + text=True, + ) + payload = json.loads(completed.stdout.strip().splitlines()[-1]) + + assert payload["first"]["reasoning_cycles"] == 8 + assert payload["first"]["adaptive_compute"] is True + assert payload["first"]["adaptive_exit_tol"] == 0.4 + assert payload["first"]["prediction_stability_patience"] == 6 + assert payload["clean"]["reasoning_cycles"] is None + assert payload["clean"]["adaptive_compute"] is False + assert payload["clean"]["adaptive_exit_tol"] == 1e-3 + assert payload["overridden"]["reasoning_cycles"] == 3 + assert payload["overridden"]["adaptive_compute"] is False + + +def test_packaged_web_runtime_escapes_default_paths(): + runtime_dir = Path("runtime_python").resolve() + script = f""" +import json +import sys +import torch +sys.path.insert(0, {str(runtime_dir)!r}) +import chat_web_app + +weights = "C:/models/o'hare\\\" onfocus=\\\"alert(1)" +meta = "C:/models/&config.json" +engine = chat_web_app.Engine(torch.device("cpu"), {{"resolved": "cpu"}}, {{"pool_mode": "all"}}) +html = chat_web_app.build_app(engine, weights, meta).test_client().get("/").get_data(as_text=True) +print(json.dumps({{ + "raw_weights_present": weights in html, + "raw_meta_present": meta in html, + "escaped_apostrophe": "o'hare" in html, + "escaped_quote": "" onfocus="" in html, + "escaped_meta": "<meta>&config.json" in html, +}})) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path.cwd(), + check=True, + capture_output=True, + text=True, + ) + payload = json.loads(completed.stdout.strip().splitlines()[-1]) + + assert payload == { + "raw_weights_present": False, + "raw_meta_present": False, + "escaped_apostrophe": True, + "escaped_quote": True, + "escaped_meta": True, + } + + +def test_packaged_web_runtime_exposes_get_only_shadow_registry_status(tmp_path): + runtime_dir = Path("runtime_python").resolve() + missing_registry = (tmp_path / "missing-shadow.sqlite3").resolve() + script = f""" +import json +import sys +import torch +sys.path.insert(0, {str(runtime_dir)!r}) +import chat_web_app + +engine = chat_web_app.Engine( + torch.device("cpu"), + {{"resolved": "cpu"}}, + {{"pool_mode": "all", "route_shadow_registry_path": {str(missing_registry)!r}}}, +) +app = chat_web_app.build_app(engine, "", "") +client = app.test_client() +html = client.get("/").get_data(as_text=True) +status = client.get("/api/route_shadow_registry/status") +payload = status.get_json() +print(json.dumps({{ + "status_code": status.status_code, + "available": payload["route_shadow_registry"]["available"], + "read_only": payload["route_shadow_registry"]["read_only"], + "execution_enabled": payload["route_shadow_registry"]["execution_enabled"], + "mutation_method": client.post("/api/route_shadow_registry/status", json={{}}).status_code, + "unknown_seal": client.post("/api/route_shadow_registry/seal", json={{}}).status_code, + "ui": "routeShadowStatus" in html and "refreshRouteShadowRegistry" in html, +}})) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path.cwd(), + check=True, + capture_output=True, + text=True, + ) + payload = json.loads(completed.stdout.strip().splitlines()[-1]) + + assert payload == { + "status_code": 200, + "available": False, + "read_only": True, + "execution_enabled": False, + "mutation_method": 405, + "unknown_seal": 404, + "ui": True, + } + assert not missing_registry.exists() + + +if __name__ == "__main__": + test_runtime_compute_smoke_suite() + print("Runtime compute control smoke tests PASSED!") diff --git a/test_supermix_multimodel_web_app.py b/test_supermix_multimodel_web_app.py index b621e7b40..6d74af3ae 100644 --- a/test_supermix_multimodel_web_app.py +++ b/test_supermix_multimodel_web_app.py @@ -1,3 +1,4 @@ +import json import os import sys import tempfile @@ -13,6 +14,13 @@ def __init__(self, zip_path: Path, summary_path: Path): self.records = [] self.generated_dir = zip_path.parent self.uploads_dir = zip_path.parent + self.chat_payloads = [] + self.preview_payloads = [] + self.study_payloads = [] + self.review_bundle_payloads = [] + self.review_audit_payloads = [] + self.shadow_status_calls = 0 + self.feedback_payloads = [] self._store_rows = [ { "file_name": "supermix_omni_collective_v8_preview_20260407_001155.zip", @@ -111,6 +119,477 @@ def install_model_store_artifact(self, file_name: str): "error": "", } + def handle_prompt(self, *, session_id: str, prompt: str, model_key: str, action_mode: str, settings: dict): + self.chat_payloads.append( + { + "session_id": session_id, + "prompt": prompt, + "model_key": model_key, + "action_mode": action_mode, + "settings": dict(settings), + } + ) + return { + "ok": True, + "kind": "text", + "model_key": model_key, + "model_label": "Stub", + "route_reason": "stub", + "response": "stub response", + "agent_trace": {}, + } + + def preview_route_plan(self, *, session_id: str, prompt: str, model_key: str, action_mode: str, settings: dict): + if model_key == "missing-model": + raise KeyError("Unknown model key: missing-model") + if model_key == "no-local-models": + raise RuntimeError("No local models were discovered.") + self.preview_payloads.append( + { + "session_id": session_id, + "prompt": prompt, + "model_key": model_key, + "action_mode": action_mode, + "settings": dict(settings), + } + ) + return { + "ok": True, + "dry_run": True, + "selected_agent_mode": "collective", + "route_economics_estimate": { + "selected_agent_mode": "collective", + "estimated_cost_units": 3.0, + }, + "route_alternatives": [ + { + "selected_agent_mode": "off", + "is_selected": False, + "estimated_cost_units": 1.0, + "frontier_rank": 2, + "budget_fit": True, + "budget_blocker": None, + "pareto_frontier": True, + "budget_feasible_pareto_frontier": True, + }, + { + "selected_agent_mode": "collective", + "is_selected": True, + "estimated_cost_units": 3.0, + "frontier_rank": 1, + "budget_fit": True, + "budget_blocker": None, + "pareto_frontier": True, + "budget_feasible_pareto_frontier": True, + }, + ], + "route_frontier": { + "selected_agent_mode": "collective", + "recommended_agent_mode": "collective", + "recommended_reason": "selected_route_is_frontier_recommended", + "recommended_budget_blocker": None, + "selected_budget_blocker": None, + "budget_blockers": {"remaining_budget": 0, "pacing_cap": 0, "effective_cap": 0, "none": 2}, + "pareto_modes": ["off", "collective"], + "budget_feasible_pareto_modes": ["off", "collective"], + "budget_cap_cost_units": 12.0, + "remaining_cost_units": 12.0, + "pacing_cap_cost_units": None, + "effective_cap_cost_units": 12.0, + "selected_matches_recommendation": True, + "ranked_modes": [ + { + "selected_agent_mode": "collective", + "frontier_rank": 1, + "budget_feasible_pareto_frontier": True, + }, + { + "selected_agent_mode": "off", + "frontier_rank": 2, + "budget_feasible_pareto_frontier": True, + }, + ], + }, + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + }, + } + + def preview_route_study( + self, + *, + session_id: str, + prompt: str, + model_key: str, + action_mode: str, + settings: dict, + exploration_rate: float, + planned_routes: int, + scenario_confidence: float, + assumed_feedback_rate: float, + target_observed_labels: int, + target_policy_profile: str, + protocol_design_mode: str, + carryover_scope: str, + interference_scope: str, + temporal_variation: str, + planned_clusters: int, + max_routes_per_cluster: int, + analysis_every_clusters: int, + block_length_routes: int, + washout_routes: int, + ): + request = { + "session_id": session_id, + "prompt": prompt, + "model_key": model_key, + "action_mode": action_mode, + "settings": dict(settings), + "exploration_rate": exploration_rate, + "planned_routes": planned_routes, + "scenario_confidence": scenario_confidence, + "assumed_feedback_rate": assumed_feedback_rate, + "target_observed_labels": target_observed_labels, + "target_policy_profile": target_policy_profile, + "protocol_design_mode": protocol_design_mode, + "carryover_scope": carryover_scope, + "interference_scope": interference_scope, + "temporal_variation": temporal_variation, + "planned_clusters": planned_clusters, + "max_routes_per_cluster": max_routes_per_cluster, + "analysis_every_clusters": analysis_every_clusters, + "block_length_routes": block_length_routes, + "washout_routes": washout_routes, + } + self.study_payloads.append(request) + label_probability = 0.05 * float(assumed_feedback_rate) + return { + "ok": True, + "dry_run": True, + "baseline_agent_mode": "collective", + "route_study": { + "schema_version": "route-exploration-plan-v1", + "design_hash": "a" * 64, + "study": { + "study_id": "auto-route-adjacent-explorer-v1", + "study_version": "1.0.0", + }, + "charter": { + "enrollment": { + "eligible": True, + "reason": "eligible_adjacent_post_filter_support", + "baseline_action": "collective", + "adjacent_feasible_actions": ["off", "loop"], + }, + "probability_design": { + "decision_type": "randomized", + "applied_exploration_rate": exploration_rate, + "minimum_positive_exploration_probability": exploration_rate / 2, + "action_probabilities": { + "off": exploration_rate / 2, + "collective": 1 - exploration_rate, + "loop": exploration_rate / 2, + }, + "assignment_performed": False, + }, + "traffic_scenario": { + "planned_routes": planned_routes, + "observed_label_scenario": { + "analysis_type": "simultaneous_alternate_label_traffic_not_power_or_mnar_correction", + "target_scope": "at_least_target_observed_labels_on_every_alternate_action", + "alternate_actions": ["off", "loop"], + "assumed_feedback_rate": assumed_feedback_rate, + "target_observed_labels_per_alternate_action": target_observed_labels, + "per_route_observed_label_probability_by_alternate_action": { + "off": label_probability, + "loop": label_probability, + }, + "expected_routes_for_target_by_alternate_action": { + "off": target_observed_labels / label_probability, + "loop": target_observed_labels / label_probability, + }, + "exact_simultaneous_target": { + "method": "exact_joint_multinomial_tail_inversion_two_alternates", + "confidence_level": scenario_confidence, + "minimum_routes_for_target_on_every_alternate_action": 1971, + }, + }, + }, + "resource_forecast": { + "expected_for_planned_routes": {"cost_units": 6700.0}, + "by_action": { + "off": {"latency_tier": "low"}, + "collective": {"latency_tier": "moderate"}, + "loop": {"latency_tier": "frontier"}, + }, + }, + "causal_boundaries": { + "execution_enabled": False, + "off_policy_estimate_computed": False, + "automatic_promotion_allowed": False, + "activation_available": False, + "activation_blockers": [ + "preassignment_seed_commitment_not_sealed", + "target_policy_class_not_precommitted", + "session_carryover_not_addressed", + "external_ope_not_validated", + ], + }, + }, + }, + "route_protocol_preflight": { + "schema_version": "route-study-protocol-preflight-v1", + "protocol_hash": "b" * 64, + "protocol": { + "version": "1.0.0", + "label": "Stateful Route Experiment Preflight v1", + "state": "draft_for_independent_review", + "activation_available": False, + }, + "charter": { + "target_policy_class": { + "profile_name": target_policy_profile, + "class_hash": "c" * 64, + }, + "population": { + "planned_clusters": planned_clusters, + "max_routes_per_cluster": max_routes_per_cluster, + }, + "stateful_design": { + "selected_design_mode": protocol_design_mode, + "selected_design_status": "declaration_incomplete", + "assignment_unit": "session_hash", + "carryover_scope": carryover_scope, + "interference_scope": interference_scope, + "temporal_variation": temporal_variation, + "selected_design_blocking_reasons": [ + "carryover_scope_unknown", + "interference_scope_unknown", + "temporal_variation_unknown", + ], + }, + "randomness": { + "assignment_implementation_available": False, + "assignment_performed": False, + }, + "blocker_register": [ + { + "code": "session_carryover_not_addressed", + "status": "unresolved", + "activation_blocking": True, + }, + { + "code": "interference_not_addressed", + "status": "unresolved", + "activation_blocking": True, + }, + ], + "causal_boundaries": { + "activation_available": False, + "automatic_promotion_allowed": False, + "activation_blockers": [ + "session_carryover_not_addressed", + "interference_not_addressed", + ], + }, + }, + }, + "route_protocol_preflight_reason": "draft_for_independent_review", + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } + + def build_route_protocol_review_bundle(self, payload: dict): + from route_policy_protocol import ( + audit_route_study_review_bundle, + build_route_study_review_bundle_from_input, + ) + + self.review_bundle_payloads.append(payload) + bundle = build_route_study_review_bundle_from_input(payload) + return { + "ok": True, + "dry_run": True, + "route_protocol_review_bundle": bundle, + "verification": audit_route_study_review_bundle(bundle), + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } + + def audit_route_protocol_review_bundle(self, bundle: dict): + from route_policy_protocol import audit_route_study_review_bundle + + self.review_audit_payloads.append(bundle) + return { + "ok": True, + "dry_run": True, + "verification": audit_route_study_review_bundle(bundle), + "execution_plan": { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + }, + } + + def record_route_feedback(self, *, session_id: str, feedback: dict): + self.feedback_payloads.append({"session_id": session_id, "feedback": dict(feedback)}) + return { + "ok": True, + "feedback": { + "selected_agent_mode": feedback.get("selected_agent_mode"), + "rating": feedback.get("rating"), + }, + "summary": { + "total_feedback": len(self.feedback_payloads), + "economics": { + "sample_count": len(self.feedback_payloads), + "avg_cost_units": 4.5, + "avg_elapsed_ms": 120.0, + }, + }, + } + + def route_health_snapshot(self, session_id: str): + return { + "total_feedback": len(self.feedback_payloads), + "economics": { + "sample_count": len(self.feedback_payloads), + "avg_cost_units": 4.5 if self.feedback_payloads else None, + "avg_elapsed_ms": 120.0 if self.feedback_payloads else None, + }, + "adaptive": { + "sample_count": len(self.feedback_payloads), + "quality_score": 0.42 if self.feedback_payloads else None, + "quality_cost_score": 0.31 if self.feedback_payloads else None, + "regression_signal": bool(self.feedback_payloads), + }, + "route_usage": { + "total_routes": len(self.feedback_payloads), + "economics": { + "sample_count": len(self.feedback_payloads), + "avg_cost_units": 4.5 if self.feedback_payloads else None, + "total_cost_units": 4.5 if self.feedback_payloads else None, + "avg_elapsed_ms": 120.0 if self.feedback_payloads else None, + }, + }, + } + + def route_policy_lab_snapshot(self, session_id: str, profile: str = "balanced"): + return { + "analysis_kind": "associational_matched_route_replay", + "evidence_source": "durable_sqlite_ledger", + "profile": {"name": profile, "thresholds": {"off": 0, "collective": 2, "loop": 4, "collective_loop": 5}}, + "support": { + "usage": {"unique_route_ids": 1}, + "exact_joined_route_ids": 1, + }, + "candidate_action_agreement": {"agreement_rate": 1.0}, + "matched_observed": {"quality_sample_count": 1, "approval_rate": 1.0}, + "propensity_readiness": {"valid_routes": 0, "checked_evaluable_usage_routes": 3}, + "evaluation_readiness": { + "schema_version": "route-readiness-v2", + "thresholds": { + "minimum_global_effective_sample_size": 20.0, + "minimum_per_action_effective_sample_size": 10.0, + }, + "target_overlap": { + "effective_sample_size": 0.0, + "weakest_target_action": "collective", + "weakest_action_effective_sample_size": 0.0, + "minimum_target_probability": None, + }, + "outcome_observation": { + "quality_observed_routes": 2, + "evaluable_routes": 3, + }, + "ready_for_external_ope": False, + "policy_value_estimated": False, + }, + "outcome_contract_maturity": { + "schema_version": "route-outcome-maturity-v1", + "included_routes": 3, + "precommitted_routes": 2, + "by_outcome": { + "user_quality_rating": { + "mature_contract_count": 3, + "observed_event_count": 2, + } + }, + "descriptive_only": True, + "policy_value_estimate": None, + }, + "durable_ledger": { + "ledger_schema_version": 2, + "counts": {"started": 3, "completed": 2, "failed": 1, "inflight": 0}, + "feedback_coverage": {"terminal_coverage_rate": 0.666667}, + }, + "promotion_gate": { + "status": "blocked", + "deployment": "shadow_only", + "reason_code": "no_valid_randomized_overlap", + "blocking_reason_codes": [ + "no_valid_randomized_overlap", + "insufficient_global_overlap_ess", + ], + "checks": { + "candidate_delta_present": True, + "population_integrity_complete": False, + "execution_integrity_complete": False, + "logging_integrity_complete": False, + "minimum_overlap_routes_met": False, + "target_probability_floor_met": False, + "global_overlap_ess_met": False, + "per_action_overlap_met": False, + "outcome_evidence_integrity": False, + "quality_observation_ready": False, + "durable_lifecycle_present": True, + "lifecycle_reconciled": False, + }, + "passed_checks": 2, + "total_checks": 12, + "automatic_promotion_allowed": False, + }, + } + + def route_shadow_registry_snapshot(self): + self.shadow_status_calls += 1 + return { + "ok": True, + "available": True, + "status": "verified", + "read_only": True, + "campaign_count": 1, + "campaigns": [ + { + "campaign_id": "shadow:" + "a" * 64, + "state": "commitments_closed", + "commitment_count": 2, + "verified_assignment_count": 0, + "mismatched_assignment_count": 0, + } + ], + "event_chain": {"ok": True, "verified_events": 4}, + "execution_enabled": False, + "activation_available": False, + "automatic_promotion_allowed": False, + } + def test_three_d_model_view_endpoint_and_downloads(): with tempfile.TemporaryDirectory() as td: @@ -177,8 +656,153 @@ def test_index_contains_discovery_ui(): assert 'id="composeMediaBtn"' in html assert 'id="composeWorkbenchBtn"' in html assert 'id="loopBudget"' in html + assert 'id="autoBudget"' in html + assert 'id="reasoningCycles"' in html + assert 'id="adaptiveCompute"' in html + assert 'id="sessionBudget"' in html + assert 'id="sessionBudgetTargetRoutes"' in html + assert 'data-mode="auto"' in html + assert "Adaptive Router" in html + assert "Scoring task complexity" in html + assert "auto_agent_budget" in html + assert "auto_session_budget_units" in html + assert "auto_session_budget_target_routes" in html + assert "budget_profile" in html + assert "session_budget_adjustment" in html + assert 'id="routeFeedback"' in html + assert 'id="routeGoodBtn"' in html + assert 'id="routeBadBtn"' in html + assert 'id="routeDeeperBtn"' in html + assert 'id="routeCostBtn"' in html + assert 'id="routeSlowBtn"' in html + assert 'id="routeHealthConfidence"' in html + assert 'id="routeHealthPreference"' in html + assert 'id="routePlanBtn"' in html + assert 'id="routeHealth"' in html + assert 'id="routeHealthCount"' in html + assert 'id="routeHealthQuality"' in html + assert 'id="policyLab"' in html + assert 'id="policyLabProfile"' in html + assert 'id="policyLabGate"' in html + assert 'id="policyLabLifecycle"' in html + assert 'id="policyLabFeedbackCoverage"' in html + assert 'id="policyLabOverlapEss"' in html + assert 'id="policyLabWeakestAction"' in html + assert 'id="policyLabReadinessChecks"' in html + assert 'id="policyLabOutcomeCoverage"' in html + assert 'id="policyLabContractCoverage"' in html + assert 'id="policyLabEvidenceSource"' in html + assert 'id="policyLabChecks"' in html + assert 'id="policyLabBlockers"' in html + assert 'id="routeStudy"' in html + assert 'id="routeStudyPreview"' in html + assert 'id="routeStudyHorizon"' in html + assert 'id="routeStudyEpsilon"' in html + assert 'id="routeStudyResponseRate"' in html + assert 'id="routeStudyTargetLabels"' in html + assert 'id="routeProtocolTarget"' in html + assert 'id="routeProtocolDesign"' in html + assert 'id="routeProtocolCarryover"' in html + assert 'id="routeProtocolInterference"' in html + assert 'id="routeProtocolTemporal"' in html + assert 'id="routeProtocolClusters"' in html + assert 'id="routeProtocolBlock"' in html + assert 'id="routeProtocolWashout"' in html + assert 'id="routeStudyStatus" role="status" aria-live="polite"' in html + assert 'id="routeStudyDistribution"' in html + assert 'id="routeProtocolMode"' in html + assert 'id="routeProtocolPolicy"' in html + assert 'id="routeProtocolReview"' in html + assert 'id="routeProtocolHash"' in html + assert 'id="routeProtocolBlockers"' in html + assert 'id="routeBundleAdd"' in html + assert 'id="routeBundleBuild"' in html + assert 'id="routeBundleDownload"' in html + assert 'id="routeBundleImport"' in html + assert 'id="routeBundleClear"' in html + assert 'id="routeBundleInventory"' in html + assert 'id="routeBundleVerification"' in html + assert 'id="routeShadowRegistry"' in html + assert 'id="routeShadowRegistryRefresh"' in html + assert 'id="routeShadowRegistryStatus" role="status" aria-live="polite"' in html + assert "full_source_bound_reconstruction" in html + assert "Bounded Exposure Rehearsal" in html + assert "Rehearsal only - execution off" in html + assert 'id="controlPanel"' in html + assert 'id="panelToggle"' in html + assert 'aria-controls="controlPanel"' in html + assert 'aria-label="Open control panel"' in html + assert 'id="panelClose"' in html + assert '' in html + assert "function setPanelOpen" in html + assert "function openPanelTab" in html + assert "openPanelTab('mode')" in html + assert "renderPolicyLabUnavailable" in html + assert ".panel.is-open" in html + assert ".panel { display:none; }" not in html + assert "/api/route_plan" in html + assert "/api/route_feedback" in html + assert "/api/route_health" in html + assert "/api/route_policy_lab" in html + assert "/api/route_study_plan" in html + assert "/api/route_study_protocol_bundle" in html + assert "/api/route_shadow_registry/status" in html + assert "function refreshRouteShadowRegistry" in html + assert "cannot seal, assign, reveal, activate, or promote" in html + assert "/api/route_shadow_registry/seal" not in html + assert "function renderRouteStudy" in html + assert "function previewRouteStudy" in html + assert "function buildRouteReviewBundle" in html + assert "function importRouteReviewFile" in html + assert "exact_simultaneous_target" in html + assert "% joint" in html + assert "Hypothetical repetition of this prompt-specific support" in html + assert "Activation blockers" in html + assert "Declarations are not validation" in html + assert "route_protocol_preflight" in html + assert "protocol_design_mode" in html + assert "willRandomizeRoute" not in html + assert "Activate study" not in html + assert "No policy value was estimated" in html + assert "Nondurable compatibility evidence; readiness is blocked." in html + assert "maturity is diagnostic only" in html + assert "buildRoutePayload" in html + assert '' in html assert '' in html assert '' in html + assert "function scorePct" in html + assert "auto_agent_policy" in html + assert "loop_stop_reason_code" in html + assert "route_economics" in html + assert "function computePills" in html + assert "requested_reasoning_cycles" in html + assert "prediction_confidence_delta" in html + assert "route_alternatives" in html + assert "route_frontier" in html + assert "recommended_agent_mode" in html + assert "recommended_budget_blocker" in html + assert "recommended_estimated_quality_cost_score" in html + assert "selected_budget_blocker" in html + assert "estimated_quality_cost_score" in html + assert "quality_evidence_status" in html + assert "quality_source" in html + assert "budget_feasible_pareto_modes" in html + assert "budget_feasible_pareto_frontier" in html + assert "frontier_rank" in html + assert "remaining_cost_units" in html + assert "pacing_cap_cost_units" in html + assert "effective_cap_cost_units" in html + assert "route_usage" in html + assert "routeEconomicsPills" in html + assert "renderRouteHealth" in html + assert "routeQualityText" in html + assert "adaptive_quality_cost_preferred_neighbor" in html + assert "Pareto" in html + assert "Planned calls" in html + assert "Session budget" in html + assert "Paced" in html + assert "review_score" in html + assert "trace-score" in html assert 'id="toggleSidebarBtn"' in html assert 'id="toggleThreadDensityBtn"' in html assert 'id="responseDeck"' in html @@ -194,6 +818,451 @@ def test_index_contains_discovery_ui(): assert 'id="challengeLastReplyBtn"' in html +def test_chat_endpoint_passes_session_budget_setting_to_manager(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/chat", + json={ + "session_id": "chat-session", + "message": "Run a budgeted route.", + "model_key": "auto", + "action_mode": "text", + "settings": { + "agent_mode": "auto", + "auto_agent_budget": "balanced", + "auto_session_budget_units": 7.5, + "auto_session_budget_target_routes": 4, + }, + }, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert stub.chat_payloads[0]["session_id"] == "chat-session" + assert stub.chat_payloads[0]["settings"]["auto_session_budget_units"] == 7.5 + assert stub.chat_payloads[0]["settings"]["auto_session_budget_target_routes"] == 4 + + +def test_route_plan_endpoint_passes_payload_to_preview_without_chat(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/route_plan", + json={ + "session_id": "preview-session", + "message": "Preview this budgeted route.", + "model_key": "auto", + "action_mode": "text", + "settings": { + "agent_mode": "auto", + "auto_session_budget_units": 12, + "auto_session_budget_target_routes": 4, + }, + }, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert payload["dry_run"] is True + assert payload["execution_plan"]["will_run_inference"] is False + assert payload["route_alternatives"][1]["selected_agent_mode"] == "collective" + assert payload["route_alternatives"][1]["is_selected"] is True + assert payload["route_frontier"]["recommended_agent_mode"] == "collective" + assert payload["route_frontier"]["ranked_modes"][0]["frontier_rank"] == 1 + assert payload["route_frontier"]["remaining_cost_units"] == 12.0 + assert payload["route_frontier"]["effective_cap_cost_units"] == 12.0 + assert payload["route_frontier"]["selected_budget_blocker"] is None + assert payload["route_frontier"]["budget_blockers"]["none"] == 2 + assert payload["route_frontier"]["budget_feasible_pareto_modes"] == ["off", "collective"] + assert stub.chat_payloads == [] + assert stub.preview_payloads[0]["session_id"] == "preview-session" + assert stub.preview_payloads[0]["settings"]["auto_session_budget_target_routes"] == 4 + + +def test_route_study_endpoint_rehearses_without_chat_or_assignment(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/route_study_plan", + json={ + "session_id": "study-session", + "message": "Research and verify this integration.", + "model_key": "auto", + "action_mode": "text", + "settings": {"agent_mode": "auto", "auto_agent_budget": "balanced"}, + "exploration_rate": 0.10, + "planned_routes": 2000, + "scenario_confidence": 0.95, + "assumed_feedback_rate": 0.30, + "target_observed_labels": 20, + "target_policy_profile": "quality_first", + "protocol_design_mode": "clustered_switchback", + "carryover_scope": "within_session", + "interference_scope": "shared_resource", + "temporal_variation": "nonstationary", + "planned_clusters": 480, + "max_routes_per_cluster": 12, + "analysis_every_clusters": 60, + "block_length_routes": 20, + "washout_routes": 4, + }, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert payload["dry_run"] is True + assert payload["route_study"]["study"]["study_id"] == "auto-route-adjacent-explorer-v1" + assert payload["route_study"]["charter"]["probability_design"]["assignment_performed"] is False + assert payload["route_protocol_preflight"]["protocol"]["activation_available"] is False + assert payload["route_protocol_preflight"]["charter"]["stateful_design"][ + "selected_design_mode" + ] == "clustered_switchback" + assert payload["route_protocol_preflight"]["charter"]["randomness"][ + "assignment_performed" + ] is False + assert payload["execution_plan"] == { + "will_run_inference": False, + "will_write_memory": False, + "will_write_ledger": False, + "will_assign_route": False, + "will_randomize": False, + "activation_available": False, + } + assert stub.chat_payloads == [] + assert stub.preview_payloads == [] + assert len(stub.study_payloads) == 1 + request = stub.study_payloads[0] + assert request["session_id"] == "study-session" + assert request["exploration_rate"] == 0.10 + assert request["planned_routes"] == 2000 + assert request["assumed_feedback_rate"] == 0.30 + assert request["target_observed_labels"] == 20 + assert request["target_policy_profile"] == "quality_first" + assert request["protocol_design_mode"] == "clustered_switchback" + assert request["carryover_scope"] == "within_session" + assert request["interference_scope"] == "shared_resource" + assert request["temporal_variation"] == "nonstationary" + assert request["planned_clusters"] == 480 + assert request["max_routes_per_cluster"] == 12 + assert request["analysis_every_clusters"] == 60 + assert request["block_length_routes"] == 20 + assert request["washout_routes"] == 4 + + +def test_route_protocol_bundle_endpoints_build_and_reconstruct_without_chat(): + from route_policy_protocol_cli import _example_bundle_input + + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + build_input = _example_bundle_input() + response = client.post("/api/route_study_protocol_bundle", json=build_input) + + assert response.status_code == 200 + payload = response.get_json() + bundle = payload["route_protocol_review_bundle"] + assert payload["verification"]["verification_level"] == ( + "full_source_bound_reconstruction" + ) + assert payload["verification"]["support_stratum_count"] == 2 + assert payload["execution_plan"]["will_run_inference"] is False + assert payload["execution_plan"]["will_write_ledger"] is False + assert stub.chat_payloads == [] + assert stub.preview_payloads == [] + assert stub.study_payloads == [] + assert len(stub.review_bundle_payloads) == 1 + + audit_response = client.post( + "/api/route_study_protocol_bundle/audit", json={"bundle": bundle} + ) + assert audit_response.status_code == 200 + audit = audit_response.get_json()["verification"] + assert audit["ok"] is True + assert audit["source_plan_reconstruction_performed"] is True + assert len(stub.review_audit_payloads) == 1 + assert stub.chat_payloads == [] + + unsafe = client.post( + "/api/route_study_protocol_bundle", + json={**build_input, "prompt": "must not enter review artifacts"}, + ) + assert unsafe.status_code == 400 + assert "non-prompt-free fields: prompt" in unsafe.get_json()["error"] + + too_many = client.post( + "/api/route_study_protocol_bundle", + json={**build_input, "study_plans": build_input["study_plans"] * 51}, + ) + assert too_many.status_code == 400 + assert "at most 100 browser strata" in too_many.get_json()["error"] + + duplicate_json = json.dumps(build_input, separators=(",", ":"))[:-1] + ( + ',"study_plans":[]}' + ) + duplicate = client.post( + "/api/route_study_protocol_bundle", + data=duplicate_json, + content_type="application/json", + ) + assert duplicate.status_code == 400 + assert "duplicate object key: study_plans" in duplicate.get_json()["error"] + assert len(stub.review_bundle_payloads) == 2 + + +def test_route_shadow_registry_endpoint_is_get_only_and_read_only(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + response = client.get("/api/route_shadow_registry/status") + + assert response.status_code == 200 + assert response.headers["Cache-Control"] == "no-store" + payload = response.get_json() + assert payload["ok"] is True + assert payload["route_shadow_registry"]["read_only"] is True + assert payload["route_shadow_registry"]["campaign_count"] == 1 + assert payload["route_shadow_registry"]["execution_enabled"] is False + assert payload["route_shadow_registry"]["activation_available"] is False + assert stub.shadow_status_calls == 1 + assert client.post("/api/route_shadow_registry/status", json={}).status_code == 405 + assert client.post("/api/route_shadow_registry/seal", json={}).status_code == 404 + assert stub.chat_payloads == [] + + +def test_route_plan_endpoint_returns_compact_error_for_bad_model_key(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/route_plan", + json={ + "session_id": "preview-session", + "message": "Preview this route.", + "model_key": "missing-model", + "action_mode": "text", + "settings": {"agent_mode": "auto"}, + }, + ) + + assert response.status_code == 400 + payload = response.get_json() + assert payload["ok"] is False + assert "missing-model" in payload["error"] + assert stub.chat_payloads == [] + assert stub.preview_payloads == [] + + +def test_route_plan_endpoint_returns_compact_error_when_no_models_are_discovered(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/route_plan", + json={ + "session_id": "preview-session", + "message": "Preview this route.", + "model_key": "no-local-models", + "action_mode": "text", + "settings": {"agent_mode": "auto"}, + }, + ) + + assert response.status_code == 400 + payload = response.get_json() + assert payload["ok"] is False + assert "No local models were discovered" in payload["error"] + assert stub.chat_payloads == [] + assert stub.preview_payloads == [] + + +def test_route_feedback_endpoint_passes_compact_payload_to_manager(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + app = build_app(stub) + client = app.test_client() + + response = client.post( + "/api/route_feedback", + json={ + "session_id": "route-session", + "route_id": "route-1", + "prompt": "Research and implement.", + "response": "done", + "selected_agent_mode": "collective_loop", + "rating": "down", + "feedback_intent": "too_costly", + "auto_agent_policy": {"selected_agent_mode": "collective_loop", "score": 5}, + "route_economics": { + "actual": { + "elapsed_ms": 120.0, + "model_calls": 4, + "cost_units": 4.5, + }, + }, + "model_key": "omni_collective_v8", + "route_reason": "Auto orchestration selected collective_loop.", + }, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert payload["summary"]["total_feedback"] == 1 + assert stub.feedback_payloads[0]["session_id"] == "route-session" + assert stub.feedback_payloads[0]["feedback"]["route_id"] == "route-1" + assert stub.feedback_payloads[0]["feedback"]["rating"] == "down" + assert stub.feedback_payloads[0]["feedback"]["feedback_intent"] == "too_costly" + assert "selected_agent_mode" not in stub.feedback_payloads[0]["feedback"] + assert "route_economics" not in stub.feedback_payloads[0]["feedback"] + + +def test_route_health_endpoint_returns_session_route_economics(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + stub = _StubManager(zip_path, summary_path) + stub.feedback_payloads.append({"session_id": "route-session", "feedback": {}}) + app = build_app(stub) + client = app.test_client() + + response = client.post("/api/route_health", json={"session_id": "route-session"}) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert payload["route_health"]["total_feedback"] == 1 + assert payload["route_health"]["economics"]["avg_cost_units"] == 4.5 + assert payload["route_health"]["adaptive"]["quality_score"] == 0.42 + assert payload["route_health"]["adaptive"]["regression_signal"] is True + assert payload["route_health"]["route_usage"]["total_routes"] == 1 + assert payload["route_health"]["route_usage"]["economics"]["total_cost_units"] == 4.5 + + +def test_route_policy_lab_endpoint_is_read_only_and_shadow_gated(): + with tempfile.TemporaryDirectory() as td: + root = Path(td) + zip_path = root / "supermix_3d_generation_micro_v1_20260403.zip" + summary_path = root / "three_d_generation_micro_v1_summary.json" + zip_path.write_bytes(b"zip-bytes") + summary_path.write_bytes(b"{}") + + app = build_app(_StubManager(zip_path, summary_path)) + client = app.test_client() + response = client.post( + "/api/route_policy_lab", + json={"session_id": "route-session", "profile": "efficiency"}, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["ok"] is True + assert payload["policy_lab"]["profile"]["name"] == "efficiency" + assert payload["policy_lab"]["analysis_kind"] == "associational_matched_route_replay" + assert payload["policy_lab"]["promotion_gate"]["deployment"] == "shadow_only" + assert payload["policy_lab"]["promotion_gate"]["automatic_promotion_allowed"] is False + assert payload["policy_lab"]["evaluation_readiness"]["ready_for_external_ope"] is False + assert payload["policy_lab"]["evaluation_readiness"]["policy_value_estimated"] is False + gate = payload["policy_lab"]["promotion_gate"] + assert gate["blocking_reason_codes"] == [ + "no_valid_randomized_overlap", + "insufficient_global_overlap_ess", + ] + assert gate["passed_checks"] == 2 + assert gate["total_checks"] == 12 + assert set(gate["checks"]) == { + "candidate_delta_present", + "population_integrity_complete", + "execution_integrity_complete", + "logging_integrity_complete", + "minimum_overlap_routes_met", + "target_probability_floor_met", + "global_overlap_ess_met", + "per_action_overlap_met", + "outcome_evidence_integrity", + "quality_observation_ready", + "durable_lifecycle_present", + "lifecycle_reconciled", + } + assert payload["policy_lab"]["evidence_source"] == "durable_sqlite_ledger" + assert payload["policy_lab"]["durable_ledger"]["counts"] == { + "started": 3, + "completed": 2, + "failed": 1, + "inflight": 0, + } + + def test_model_store_endpoints(): with tempfile.TemporaryDirectory() as td: root = Path(td) From d912f23b17ed9c3892c92dfa227f947af5a38135 Mon Sep 17 00:00:00 2001 From: Kai Piper Date: Tue, 21 Jul 2026 21:39:12 +0100 Subject: [PATCH 6/6] build: package and verify the research runtime --- .github/workflows/runtime-quality-gates.yml | 133 +++ ARCHITECTURE.md | 121 ++- CONTRIBUTING.md | 44 +- README.md | 145 ++- SupermixStudioDesktop.spec | 90 +- installer/SupermixStudioDesktop.iss | 6 +- installer/postinstall_notes_studio.txt | 4 +- source/RESEARCH_UPGRADES.md | 836 ++++++++++++++++++ source/build_supermix_studio_desktop_exe.ps1 | 70 +- ...uild_supermix_studio_desktop_installer.ps1 | 6 + source/generate_studio_runtime_manifest.py | 239 +++++ source/studio_runtime_manifest.json | 143 +++ test_studio_runtime_manifest.py | 102 +++ 13 files changed, 1883 insertions(+), 56 deletions(-) create mode 100644 .github/workflows/runtime-quality-gates.yml create mode 100644 source/generate_studio_runtime_manifest.py create mode 100644 source/studio_runtime_manifest.json create mode 100644 test_studio_runtime_manifest.py diff --git a/.github/workflows/runtime-quality-gates.yml b/.github/workflows/runtime-quality-gates.yml new file mode 100644 index 000000000..5bdfb223b --- /dev/null +++ b/.github/workflows/runtime-quality-gates.yml @@ -0,0 +1,133 @@ +name: Runtime quality gates + +on: + workflow_dispatch: + pull_request: + push: + branches: [main] + +permissions: + contents: read + +concurrency: + group: runtime-quality-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + runtime-parity: + name: Runtime parity (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + defaults: + run: + shell: bash + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install focused test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install pytest flask pillow + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + + - name: Compile runtime surfaces + run: | + python -m py_compile \ + source/chat_app.py \ + source/chat_web_app.py \ + source/model_variants.py \ + source/multimodel_memory.py \ + source/route_policy_explorer.py \ + source/route_policy_protocol.py \ + source/route_policy_ledger.py \ + source/route_policy_lab.py \ + source/route_policy_shadow_registry.py \ + source/route_policy_study_cli.py \ + source/route_policy_protocol_cli.py \ + source/route_policy_shadow_cli.py \ + source/generate_studio_runtime_manifest.py \ + source/multimodel_runtime.py \ + source/qwen_supermix_pipeline.py \ + source/supermix_multimodel_web_app.py \ + source/supermix_multimodel_desktop_app.py \ + runtime_python/chat_app.py \ + runtime_python/chat_web_app.py \ + runtime_python/model_variants.py + + - name: Check self-contained runtime snapshot + run: python source/sync_runtime_model_variants.py --check + + - name: Check Studio distribution contract + run: python source/generate_studio_runtime_manifest.py --check + + - name: Run runtime and routing regression tests + run: | + python -m pytest -q test_runtime_model_variants_parity.py test_runtime_compute_controls.py + python -m pytest -q test_cognitive_leap_ultra_expert.py test_v51_materialize_demo.py + python -m pytest -q test_route_policy_explorer.py test_route_policy_protocol.py test_route_policy_ledger.py test_route_policy_lab.py test_route_policy_shadow_registry.py test_route_policy_study_cli.py test_route_policy_protocol_cli.py test_route_policy_shadow_cli.py test_route_policy_shadow_web_status.py test_studio_runtime_manifest.py test_multimodel_memory_tools.py + python -m pytest -q test_multimodel_runtime.py + python -m pytest -q test_supermix_multimodel_web_app.py + python -m pytest -q test_training_improvements.py + + - name: Freeze and smoke-test protocol console + if: runner.os == 'Windows' + shell: powershell + run: | + python -m pip install pyinstaller + New-Item -ItemType Directory -Path build\ci_route_study -Force | Out-Null + python -m PyInstaller --noconfirm --clean --onefile --console --name SupermixRouteStudy --paths source --distpath dist\ci_route_study --workpath build\ci_route_study\work --specpath build\ci_route_study source\route_policy_protocol_cli.py + $payload = (& .\dist\ci_route_study\SupermixRouteStudy.exe --example --compact) | ConvertFrom-Json + if ($payload.schema_version -ne 'route-study-protocol-preflight-v1') { throw 'Frozen protocol schema mismatch' } + if ($payload.protocol.activation_available -ne $false) { throw 'Frozen protocol unexpectedly enables activation' } + if ($payload.charter.randomness.assignment_performed -ne $false) { throw 'Frozen protocol unexpectedly assigns a route' } + $bundlePath = Join-Path $PWD 'build\ci_route_study\review-bundle.json' + & .\dist\ci_route_study\SupermixRouteStudy.exe --example-bundle --output $bundlePath + if ($LASTEXITCODE -ne 0) { throw 'Frozen review-bundle build failed' } + $bundle = Get-Content -LiteralPath $bundlePath -Raw | ConvertFrom-Json + if ($bundle.schema_version -ne 'route-study-review-bundle-v1') { throw 'Frozen review-bundle schema mismatch' } + if ($bundle.source_study_plans.Count -ne 2) { throw 'Frozen review bundle lost source strata' } + $audit = (& .\dist\ci_route_study\SupermixRouteStudy.exe --audit-bundle $bundlePath --compact) | ConvertFrom-Json + if ($audit.verification_level -ne 'full_source_bound_reconstruction') { throw 'Frozen review bundle did not fully reconstruct' } + if ($audit.source_plan_reconstruction_performed -ne $true) { throw 'Frozen review bundle skipped source reconstruction' } + if ($audit.activation_available -ne $false) { throw 'Frozen review bundle unexpectedly enables activation' } + + New-Item -ItemType Directory -Path build\ci_route_shadow -Force | Out-Null + python -m PyInstaller --noconfirm --clean --onefile --console --name SupermixRouteShadow --paths source --distpath dist\ci_route_shadow --workpath build\ci_route_shadow\work --specpath build\ci_route_shadow source\route_policy_shadow_cli.py + $registryPath = Join-Path $PWD 'build\ci_route_shadow\shadow.sqlite3' + $seedPath = Join-Path $PWD 'build\ci_route_shadow\private-seed.json' + $sealRaw = & .\dist\ci_route_shadow\SupermixRouteShadow.exe seal --registry $registryPath --bundle $bundlePath --seed-output $seedPath --compact + if ($LASTEXITCODE -ne 0) { throw 'Frozen shadow campaign seal failed' } + $seed = (Get-Content -LiteralPath $seedPath -Raw | ConvertFrom-Json).seed_material_base64url + if ($sealRaw -match [regex]::Escape($seed)) { throw 'Frozen shadow seal leaked private seed material' } + $seal = $sealRaw | ConvertFrom-Json + if ($seal.execution_enabled -ne $false) { throw 'Frozen shadow seal unexpectedly enables execution' } + $campaignId = $seal.public_package.campaign_seal.seal.campaign_id + $clusterPath = Join-Path $PWD 'build\ci_route_shadow\cluster.json' + @{ cluster_identifier = 'ci-shadow-cluster-1' } | ConvertTo-Json -Compress | Set-Content -LiteralPath $clusterPath -Encoding utf8 + $commit = (& .\dist\ci_route_shadow\SupermixRouteShadow.exe commit --registry $registryPath --campaign $campaignId --seed-input $seedPath --cluster-input $clusterPath --compact) | ConvertFrom-Json + if ($commit.chosen_arm_revealed -ne $false) { throw 'Frozen shadow commitment revealed an arm early' } + $closure = (& .\dist\ci_route_shadow\SupermixRouteShadow.exe close --registry $registryPath --campaign $campaignId --compact) | ConvertFrom-Json + if ($closure.closure.frozen_commitment_count -ne 1) { throw 'Frozen shadow closure count mismatch' } + $revealRaw = & .\dist\ci_route_shadow\SupermixRouteShadow.exe reveal --registry $registryPath --campaign $campaignId --seed-input $seedPath --compact + if ($revealRaw -match [regex]::Escape($seed)) { throw 'Frozen shadow reveal CLI leaked seed material' } + $reveal = $revealRaw | ConvertFrom-Json + if ($reveal.seed_material_returned -ne $false) { throw 'Frozen shadow reveal returned seed material' } + $verify = (& .\dist\ci_route_shadow\SupermixRouteShadow.exe verify --registry $registryPath --campaign $campaignId --compact) | ConvertFrom-Json + if (($verify.complete -ne $true) -or ($verify.matched -ne 1) -or ($verify.mismatched -ne 0)) { throw 'Frozen shadow assignment reconstruction failed' } + $registryHashBeforeStatus = (Get-FileHash -LiteralPath $registryPath -Algorithm SHA256).Hash + $status = (& .\dist\ci_route_shadow\SupermixRouteShadow.exe status --registry $registryPath --campaign $campaignId --compact) | ConvertFrom-Json + if ($status.campaigns[0].state -ne 'reveal_verification_complete') { throw 'Frozen shadow registry state mismatch' } + if ($status.read_only -ne $true) { throw 'Frozen shadow registry status is not explicitly read-only' } + if ($status.activation_available -ne $false) { throw 'Frozen shadow registry unexpectedly enables activation' } + $registryHashAfterStatus = (Get-FileHash -LiteralPath $registryPath -Algorithm SHA256).Hash + if ($registryHashAfterStatus -ne $registryHashBeforeStatus) { throw 'Frozen shadow status mutated durable registry bytes' } diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3bc4a1f84..e307ca42a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -1,6 +1,7 @@ -# Architecture Guide: Supermix_27 +# Architecture Guide: Supermix -This document provides an in-depth technical overview of the Supermix_27 model architecture, training pipeline, and design decisions. +This document provides an in-depth technical overview of the Supermix model +architecture, training pipeline, runtime surfaces, and design decisions. --- @@ -35,10 +36,122 @@ This document provides an in-depth technical overview of the Supermix_27 model a ``` The project is split into three main deployment surfaces: -- **`source/`**: All training scripts, model definitions, dataset builders, and benchmarks. -- **`runtime_python/`**: A self-contained copy of the inference-critical files with the production checkpoint. +- **`source/`**: Active training code and the packaged multimodel Supermix Studio + runtime, including the route evidence/control plane. +- **`runtime_python/`**: A legacy self-contained compatibility runtime. Its + generated `model_variants.py` stays synchronized, but it is not the active + multimodel Studio packaging tree. - **`web_static/`**: A static HTML/JS interface for GitHub Pages deployment (metadata retrieval only, no `.pth` inference). +The Windows Studio build packages +`source/supermix_multimodel_desktop_app.py` and includes a separate +`SupermixRouteStudy.exe` review console plus a separate +`SupermixRouteShadow.exe` commitment/reveal console. +`source/studio_runtime_manifest.json` binds the active entrypoints, route schema +versions, and module digests so source-only drift is detected before packaging. + +The route-study control plane has four intentionally separate artifact classes: + +1. an adjacent-route rehearsal plan for one prompt-specific post-filter support + stratum; +2. a prompt-free stateful protocol draft that inventories one or more compatible + strata without assigning or executing them; +3. a `route-study-review-bundle-v1` portable review bundle containing the full + canonical source plans and closed builder options; +4. a shadow-only whole-policy assignment registry that consumes a fully + reconstructable review bundle without writing the executed-decision ledger. + +Bundle verification reconstructs the protocol byte-for-byte from its source +plans. Standalone protocol audit remains structural because the compact draft +does not embed those plans. Neither verification level supplies authenticity, +a trusted timestamp, causal identification, random assignment, or permission to +activate; those remain separate external gates. + +The shadow registry is implemented by `source/route_policy_shadow_registry.py` +and stored separately at `memory/route-policy-shadow-registry.sqlite3`; it does +not share tables or lifecycle semantics with `route-policy-ledger.sqlite3`. A +campaign design binding freezes one sticky session-cluster assignment unit and +exactly two whole-policy arms at 5,000 basis points each: +`incumbent_source_policy` binds the complete source-policy cohort, while +`candidate_target_policy` binds the complete target-policy class. The +prompt-specific route actions carried by source strata describe where support +exists; they are not treatment arms and are never pooled into a whole-policy +assignment set. + +Those bindings are declarative policy-class manifests, not executable-code +attestations. V1 also treats the private cluster identifier as an opaque byte +string after input validation; it does not prove cluster-map membership or +canonicalize aliases. Executable artifact binding and cluster-map governance +remain separate prerequisites outside this registry. + +The commitment/reveal flow is deliberately one way: + +1. `seal` fully reconstructs the review bundle, generates 256 seed bits from the + operating-system CSPRNG, binds the seed commitment into a committed copy of + the bundle, and writes a public campaign seal; +2. the seed material is written first to a separately handled private capsule + and is absent from the registry before reveal; +3. `commit` derives a study-scoped cluster pseudonym internally and appends only + an opaque assignment-reveal commitment, withholding the chosen arm; +4. `close` freezes the commitment count and prevents further enrollment; +5. only then can `reveal` persist the public seed opening, after which `verify` + reconstructs the frozen assignments and checks their commitments. + +Registry artifacts use a strict, JCS-informed canonical JSON subset: duplicate +keys and non-finite values fail closed, registry numbers are I-JSON-safe +integers, and object keys are ASCII for portable ordering. RFC 5869 +extract-and-expand derives context-separated identity and assignment keys; +HMAC-SHA-256 derives pseudonyms and whole-policy draw buckets. These mechanics +make local reconstruction deterministic. They do not prove unbiased enrollment, +independent seed custody, nonce-grinding resistance, anonymity, cluster +independence, non-anticipation, finite carryover, or a valid interference model. + +All registry mutations require `SupermixRouteShadow.exe`. The Studio server +exposes only `GET /api/route_shadow_registry/status`, so the browser can verify +and display campaign state and the event chain but cannot seal, enroll, close, +reveal, or verify assignments. The executed route ledger rejects shadow and +rehearsal flags plus reserved shadow commitment namespaces. Any non-null +executed assignment evidence must use the closed +`route-execution-assignment-v1:` namespace, so the registry's bare +commitment digest cannot pass logging-support validation. This is a local type +boundary, not authenticated provenance against a hostile caller. + +The Studio server defaults to the loopback interface (`127.0.0.1`); listening +on another interface is an explicit operator choice and does not add +authentication. The browser response is always `Cache-Control: no-store`. +Within one process, an audited status snapshot is reused only while the durable +database and non-empty WAL stat signature are unchanged; any observed change or +concurrent-write race forces another full read-only audit. + +SQLite `BEGIN IMMEDIATE` transactions, append-only triggers, immutable closure +and reveal rows, campaign-order indexes, and a domain-separated event hash chain +protect normal local operation. Read-only snapshots audit those schema objects, +verify an exact schema-definition fingerprint, reconstruct every stored artifact, +and require the event inventory to match the evidence tables. Assignment +verification preflights the frozen closure, public seed opening, and each +commitment projection; processed mismatches remain verification failures. This +remains a local file controlled by the host: there is no signer, +trusted clock, external witness, verifiable-data-structure receipt, consistency +or inclusion proof, or transparency service. The registry never invokes a +model, chooses a live route, records an executed propensity, performs causal +inference, activates a policy, or enables automatic promotion. + +`runtime_python/model_variants.py` is a generated snapshot of the canonical +architecture registry. After changing `source/model_variants.py`, refresh and +verify the standalone runtime with: + +```bash +python source/sync_runtime_model_variants.py +python source/sync_runtime_model_variants.py --check +``` + +The v51 prediction-stability verifier uses top-1 persistence plus confidence +range to decide an inference-only early exit. It also computes a top-k +Jensen-Shannon drift diagnostic over consecutive full-prefix output +distributions, retaining all non-top-k mass in one bucket. That value is +non-causal, shadow telemetry: it is exposed by the benchmark but is not an exit +condition, router input, or activation signal. + --- ## ChampionNet Backbone diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 232532cf1..cec84c3b2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,6 +1,7 @@ -# Contributing to Supermix_27 +# Contributing to Supermix -Thank you for your interest in contributing to Supermix_27! This guide explains how to get started, the code conventions we follow, and how to submit your changes. +Thank you for contributing to Supermix. This guide covers the active source, +legacy compatibility, test, and packaging contracts. --- @@ -19,7 +20,7 @@ Thank you for your interest in contributing to Supermix_27! This guide explains ### Repository Structure ``` -Supermix_27/ +Supermix/ ├── source/ # Training, model definitions, dataset builders │ ├── run.py # ChampionNet backbone definition │ ├── model_variants.py # All classifier head variants @@ -28,9 +29,9 @@ Supermix_27/ │ ├── device_utils.py # Cross-platform device selection │ ├── benchmark.py # Evaluation benchmarks │ └── build_*.py # Dataset construction scripts -├── runtime_python/ # Self-contained inference runtime -│ ├── chat_web_app.py # Production web server -│ └── *.pth # Production checkpoint +├── runtime_python/ # Legacy self-contained compatibility runtime +├── source/supermix_multimodel_web_app.py # Active Studio web runtime +├── source/supermix_multimodel_desktop_app.py # Active packaged Studio entrypoint ├── web_static/ # GitHub Pages static UI ├── MODEL_CARD_V28.md # Model card with benchmarks ├── ARCHITECTURE.md # Technical architecture deep-dive @@ -58,7 +59,14 @@ Supermix_27/ 4. **Register in `finetune_chat.py`**: Add your variant to the `--model_size` choices. -5. **Update documentation**: Add your variant to `MODEL_CARD_V28.md` and `ARCHITECTURE.md`. +5. **Refresh the standalone runtime snapshot**: + + ```bash + python source/sync_runtime_model_variants.py + python source/sync_runtime_model_variants.py --check + ``` + +6. **Update documentation**: Add your variant to `MODEL_CARD_V28.md` and `ARCHITECTURE.md`. ### Adding a New Dataset @@ -116,18 +124,32 @@ docs(model_card): fill in v28 benchmark results ## Testing -Currently the project does not have a formal test suite. When contributing: +The project has a formal pytest suite and focused runtime quality workflow. When contributing: + +1. Run the focused tests for the files you changed, then run the full suite: + + ```bash + python -m pytest -q + ``` + +2. For active Studio runtime or route-control changes, verify the checked + distribution contract and generated model-variant snapshot: + + ```bash + python source/generate_studio_runtime_manifest.py --check + python source/sync_runtime_model_variants.py --check + ``` -1. **Smoke test** your changes by running a short training session: +3. **Smoke test** training-specific changes with a short training session: ```bash python source/finetune_chat.py --data source/conversation_data.smoke_manifest.json \ --data_manifest source/conversation_data.smoke_manifest.json \ --model_size --epochs 1 --batch_size 32 ``` -2. **Verify weight loading** from existing checkpoints to ensure backward compatibility. +4. **Verify weight loading** from existing checkpoints to ensure backward compatibility. -3. **Run benchmarks** if applicable: +5. **Run benchmarks** if applicable: ```bash python source/benchmark.py ``` diff --git a/README.md b/README.md index a7836c995..f5497b766 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# Supermix_29 +# Supermix -Supermix_29 is the working monorepo for the current Supermix / ChampionNet / Omni Collective line. +This is the working monorepo for the current Supermix / ChampionNet / Omni Collective line. This repository combines: @@ -15,12 +15,15 @@ It is intentionally a mixed workspace, not a minimal source-only model repo. ## Current status -As of March 29, 2026: +As of July 18, 2026: -- the latest finished omni checkpoint in this repo is `omni_collective_v4` -- the latest packaged desktop release is `studio-desktop-20260329-omni-v4-allmodels` -- the installer bundle currently includes `23` zipped model artifacts from the local model-pack directory used by the desktop build -- a `v5` continuation path exists in `source/` and is currently an in-progress local experiment, not a finished released model +- `source/` is the active Supermix Studio runtime and packaging tree +- the curated desktop build selects `11` core model artifacts and leaves expansion to the model store +- the route control plane includes durable lifecycle evidence, Policy Lab diagnostics, + bounded-exposure rehearsal, and a fail-closed stateful experiment protocol preflight +- `runtime_python/` remains a legacy compatibility snapshot for the smaller chat runtime; + it is not the source of truth for the multimodel Studio route control plane +- the Windows installer contract version is `2026.07.18` ## What is in this repo @@ -28,8 +31,9 @@ As of March 29, 2026: - active development workspace - training scripts, model definitions, dataset builders, benchmark runners, desktop packaging helpers - `runtime_python/` - - packaged local runtime path - - simpler run path than the full `source/` workspace + - legacy, self-contained compatibility runtime + - model-variant parity is generated from `source/model_variants.py`, but Studio-only + routing features intentionally ship through the active `source/` desktop build - `datasets/` - conversation, coding, reasoning, science, and related local training inputs - `output/` @@ -77,6 +81,16 @@ python source/chat_web_app.py python source/supermix_multimodel_desktop_app.py ``` +### Inspect the route experiment control plane + +```bash +python source/route_policy_study_cli.py --example +python source/route_policy_protocol_cli.py --example +``` + +Both commands are prompt-free, non-executing design tools. They do not assign a +route, write evidence, estimate policy value, or enable promotion. + ### Run the browser-only static bundle Open: @@ -85,9 +99,10 @@ Open: web_static/index.html ``` -## Current desktop release +## Historical desktop release -Latest release published from this repo: +The repository still documents this older published artifact for provenance; it +is not evidence that the current development tree has been released: - Release page: - `https://github.com/kai9987kai/Supermix_29/releases/tag/studio-desktop-20260329-omni-v4-allmodels` @@ -99,6 +114,8 @@ Latest release published from this repo: Local build outputs: - `dist/SupermixStudioDesktop/SupermixStudioDesktop.exe` +- `dist/SupermixStudioDesktop/SupermixRouteStudy.exe` +- `dist/SupermixStudioDesktop/SupermixRouteShadow.exe` - `dist/installer/SupermixStudioDesktopSetup.exe` - `dist/installer/SupermixStudioDesktopReleaseSHA256.txt` @@ -212,12 +229,114 @@ Primary desktop build helpers: - `SupermixStudioDesktop.spec` - `installer/SupermixStudioDesktop.iss` -The current bundled-model manifest is: +The model bundle manifest is generated from the selected local model store at +packaging time as `output/supermix_studio_bundled_models_manifest.json`; it is a +build artifact, not a checked source-of-truth file. + +The deterministic Studio code-and-contract manifest is: + +- [`source/studio_runtime_manifest.json`](source/studio_runtime_manifest.json) + +Verify it before packaging: + +```bash +python source/generate_studio_runtime_manifest.py --check +``` + +The desktop build also produces `SupermixRouteStudy.exe`, a console entrypoint +for exporting and auditing fail-closed protocol drafts and portable +multi-stratum review bundles. A review bundle embeds the complete canonical, +prompt-free source plans and the closed builder options so its protocol can be +reconstructed exactly: + +```powershell +SupermixRouteStudy.exe --example-bundle --output route-review.json +SupermixRouteStudy.exe --audit-bundle route-review.json --compact +``` + +`full_source_bound_reconstruction` proves internal semantic conformance to the +checked v1 builder. It is deliberately not a signature, trusted timestamp, +causal validation, protocol seal, route assignment, or activation approval. +The browser Studio can collect compatible support strata in ephemeral client +memory, build/download the same bundle, and re-import it for verification; it +never pools unweighted strata into a policy-value estimate. The installer +includes the Studio and both route consoles. The checked manifest binds the +active entrypoints, routing schema versions, module hashes, reconstruction +capability, and non-activation package guards. + +The build also produces `SupermixRouteShadow.exe`, an explicit local console for +the separate shadow-only whole-policy commitment/reveal registry. It accepts +only a fully reconstructable review bundle and freezes exactly two 50/50 arms: +`incumbent_source_policy`, bound to the complete source-policy cohort, and +`candidate_target_policy`, bound to the complete target-policy class. A +prompt-specific route plan's `eligible_actions` remain support-stratum evidence; +they are never reinterpreted as campaign treatment arms. +These arms bind declarative policy-class manifests, not executable code or a +runtime code digest. Cluster identifiers are treated as opaque private inputs; +v1 does not validate membership in an external cluster map or canonicalize +aliases. Both checks remain external prerequisites for any future experiment. + +The mutating workflow is CLI-only: + +```powershell +$seal = SupermixRouteShadow.exe seal --registry route-policy-shadow-registry.sqlite3 --bundle route-review.json --seed-output route-shadow-seed.private.json --compact | ConvertFrom-Json +$campaign = $seal.public_package.campaign_seal.seal.campaign_id +SupermixRouteShadow.exe commit --registry route-policy-shadow-registry.sqlite3 --campaign $campaign --seed-input route-shadow-seed.private.json --cluster-input cluster.private.json +SupermixRouteShadow.exe close --registry route-policy-shadow-registry.sqlite3 --campaign $campaign +SupermixRouteShadow.exe reveal --registry route-policy-shadow-registry.sqlite3 --campaign $campaign --seed-input route-shadow-seed.private.json +SupermixRouteShadow.exe verify --registry route-policy-shadow-registry.sqlite3 --campaign $campaign +SupermixRouteShadow.exe status --registry route-policy-shadow-registry.sqlite3 --campaign $campaign +``` -- [`output/supermix_studio_bundled_models_manifest.json`](output/supermix_studio_bundled_models_manifest.json) +`cluster.private.json` must contain exactly +`{"cluster_identifier":"..."}`. The CLI derives and stores a study-scoped +pseudonym, not that raw identifier. Keep both the cluster input and the +separately created seed capsule private; pseudonymity is not anonymity, and +post-reveal unlinkability is not guaranteed. Before closure the registry stores +only an opaque assignment commitment. After closure, `reveal` opens the seed and +`verify` reconstructs each frozen assignment. The browser exposes only a GET +status view of the canonical Studio registry; it has no seal, commit, close, +reveal, or verify mutation endpoint. CLI `status` also opens SQLite in read-only +mode and audits required tables, append-only/state-transition triggers, +high-volume campaign indexes, their exact schema-definition fingerprint, every +stored artifact, the event chain, and the one-to-one event/evidence inventory. +Reveal processing distinguishes matched assignments from mismatches; a fully +processed campaign is not reported as verification-complete unless the final +whole-campaign audit also passes. + +The Studio web server binds to `127.0.0.1` by default. Remote exposure requires +an explicit `--host` override and its own authentication/network controls. The +read-only status endpoint sends `Cache-Control: no-store`; the runtime reuses an +audited in-process snapshot only while the durable SQLite database and non-empty +WAL signature remain unchanged. + +This registry is isolated from `route-policy-ledger.sqlite3`. Its artifacts are +explicitly non-ledger-eligible and never execute a route, record an executed +propensity, run inference, estimate policy value, certify a causal design, +activate a policy, or promote one automatically. Its SQLite triggers and event +hash chain detect ordinary local mutation; they are not a signature, trusted +timestamp, external witness, inclusion/consistency proof, or transparency +service. Randomized executed-ledger rows additionally require the closed +`route-execution-assignment-v1:` namespace, so a bare shadow commitment +digest is rejected instead of being silently relabelled as execution evidence. +This is type separation, not proof against a hostile local caller fabricating an +executed envelope. ## Research and experiment notes +The v51 stability benchmark supports a research-derived, shadow-only +distribution-drift diagnostic: + +```powershell +python source/benchmark_v51_prediction_stability.py --distribution_top_k 5 +``` + +It reports consecutive-prefix top-k Jensen-Shannon divergence with a retained +`other` mass bucket. The metric is telemetry only and does not authorize an +early exit. A 480-request CPU pilot retained 96/96 agreement for the current +patience-2/tolerance-0.005 setting with 28.8% fewer cycles; the release gate is a +larger fresh-seed confirmation, not that pilot result. + This repo is a living experiment workspace. It contains finished artifacts, release-ready packaging, and in-progress work at the same time. That means you will see mixed generations such as: diff --git a/SupermixStudioDesktop.spec b/SupermixStudioDesktop.spec index 5557c8e8d..d44f12559 100644 --- a/SupermixStudioDesktop.spec +++ b/SupermixStudioDesktop.spec @@ -1,32 +1,74 @@ # -*- mode: python ; coding: utf-8 -*- +"""Portable PyInstaller template for a prepared Supermix Studio build tree. + +The normal build script stages models and a base model under ``build/`` before +invoking PyInstaller. Environment variables can override those locations for +release builders without embedding a developer machine path in this file. +""" + +import os +from pathlib import Path + from PyInstaller.utils.hooks import collect_all -datas = [('assets', 'assets'), ('C:\\Users\\kai99\\Desktop\\New folder (9)\\Supermix_27\\build\\studio_models_stage', 'bundled_models'), ('C:\\Users\\kai99\\Desktop\\New folder (9)\\Supermix_27\\build\\studio_base_model_stage', 'bundled_base_model'), ('C:\\Users\\kai99\\Desktop\\New folder (9)\\Supermix_27\\output\\benchmark_all_models_common_plus_summary_20260403.json', 'output'), ('C:\\Users\\kai99\\Desktop\\New folder (9)\\Supermix_27\\output\\supermix_studio_bundled_models_manifest.json', 'output')] + +repo_root = Path(SPECPATH).resolve() + + +def data_if_present(path, destination, *, required=False): + candidate = Path(path).expanduser().resolve() + if candidate.exists(): + return [(str(candidate), destination)] + if required: + raise FileNotFoundError(f"required Studio packaging input is missing: {candidate}") + return [] + + +models_stage = os.environ.get( + "SUPERMIX_STUDIO_MODELS_STAGE", + str(repo_root / "build" / "studio_models_stage"), +) +base_model_stage = os.environ.get( + "SUPERMIX_STUDIO_BASE_MODEL_STAGE", + str(repo_root / "build" / "studio_base_model_stage"), +) +benchmark_summary = os.environ.get("SUPERMIX_STUDIO_BENCHMARK_SUMMARY", "") +bundle_manifest = os.environ.get( + "SUPERMIX_STUDIO_BUNDLE_MANIFEST", + str(repo_root / "output" / "supermix_studio_bundled_models_manifest.json"), +) +runtime_manifest = repo_root / "source" / "studio_runtime_manifest.json" + +datas = [(str(repo_root / "assets"), "assets")] +datas += data_if_present(models_stage, "bundled_models", required=True) +datas += data_if_present(base_model_stage, "bundled_base_model", required=True) +if benchmark_summary: + datas += data_if_present(benchmark_summary, "output", required=True) +datas += data_if_present(bundle_manifest, "output", required=True) +datas += data_if_present(runtime_manifest, "output", required=True) + binaries = [] hiddenimports = [] -tmp_ret = collect_all('webview') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('flask') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('werkzeug') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('PIL') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('sympy') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('mpmath') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('safetensors') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('transformers') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] -tmp_ret = collect_all('peft') -datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] +for package in ( + "webview", + "flask", + "werkzeug", + "PIL", + "sympy", + "mpmath", + "safetensors", + "transformers", + "peft", +): + package_datas, package_binaries, package_hiddenimports = collect_all(package) + datas += package_datas + binaries += package_binaries + hiddenimports += package_hiddenimports a = Analysis( - ['source\\supermix_multimodel_desktop_app.py'], - pathex=['source'], + [str(repo_root / "source" / "supermix_multimodel_desktop_app.py")], + pathex=[str(repo_root / "source")], binaries=binaries, datas=datas, hiddenimports=hiddenimports, @@ -44,7 +86,7 @@ exe = EXE( a.scripts, [], exclude_binaries=True, - name='SupermixStudioDesktop', + name="SupermixStudioDesktop", debug=False, bootloader_ignore_signals=False, strip=False, @@ -55,7 +97,7 @@ exe = EXE( target_arch=None, codesign_identity=None, entitlements_file=None, - icon=['C:\\Users\\kai99\\Desktop\\New folder (9)\\Supermix_27\\assets\\supermix_qwen_icon.ico'], + icon=[str(repo_root / "assets" / "supermix_qwen_icon.ico")], ) coll = COLLECT( exe, @@ -64,5 +106,5 @@ coll = COLLECT( strip=False, upx=True, upx_exclude=[], - name='SupermixStudioDesktop', + name="SupermixStudioDesktop", ) diff --git a/installer/SupermixStudioDesktop.iss b/installer/SupermixStudioDesktop.iss index 9c472ec0a..3db47cb94 100644 --- a/installer/SupermixStudioDesktop.iss +++ b/installer/SupermixStudioDesktop.iss @@ -5,7 +5,7 @@ #define MyAppExeName "SupermixStudioDesktop.exe" #endif #ifndef MyAppVersion - #define MyAppVersion "2026.03.27" + #define MyAppVersion "2026.07.18" #endif #ifndef MySourceDir #define MySourceDir "..\dist\SupermixStudioDesktop" @@ -23,9 +23,11 @@ AppName={#MyAppName} AppVersion={#MyAppVersion} AppVerName={#MyAppName} {#MyAppVersion} AppPublisher=Supermix -DefaultDirName={autopf}\Supermix Studio +DefaultDirName={localappdata}\Programs\Supermix Studio DefaultGroupName={#MyAppName} DisableProgramGroupPage=yes +PrivilegesRequired=lowest +PrivilegesRequiredOverridesAllowed=dialog UninstallDisplayIcon={app}\{#MyAppExeName} SetupIconFile=..\assets\supermix_qwen_icon.ico WizardStyle=modern diff --git a/installer/postinstall_notes_studio.txt b/installer/postinstall_notes_studio.txt index be8bf6015..0e356089a 100644 --- a/installer/postinstall_notes_studio.txt +++ b/installer/postinstall_notes_studio.txt @@ -1,7 +1,9 @@ -Supermix Studio installs the bundled local model zips and the shared Qwen base model directly with the app. +Supermix Studio installs a curated core model set and the shared Qwen base model directly with the app. Additional supported models can be added through the model store. Notes: - The first time you open a model, the app may extract that zip into its local runtime cache before loading it. - Auto mode routes each prompt to the most appropriate bundled model instead of pinning one model for the whole session. - Image-capable models render images directly inside the chat thread. - The app stores generated images and extracted working copies under its runtime state directory in your user profile. +- SupermixRouteStudy.exe is installed beside the desktop app. It exports and audits prompt-free experiment protocol drafts; it cannot assign or execute routes, write evidence, estimate policy value, or activate promotion. +- SupermixStudioRuntimeManifest.json and the release SHA-256 file describe the packaged runtime contract and artifact integrity. diff --git a/source/RESEARCH_UPGRADES.md b/source/RESEARCH_UPGRADES.md index 35a589906..bf7a6abad 100644 --- a/source/RESEARCH_UPGRADES.md +++ b/source/RESEARCH_UPGRADES.md @@ -474,3 +474,839 @@ Repo changes: - `benchmark_results.json` now also includes `artifacts` pointers and a compact `sample_summary.worst_regression` block so tools do not need to rescan the full comparison file to show the top failure. - The training monitor research board now surfaces the selected run's top regression, prompt preview, and tuned/reference preview, and adds a direct `Open selected samples` action. - Older aggregate-only benchmark artifacts remain readable; the monitor now explicitly reports when a run needs a rerun to generate detailed sample traces instead of showing a blank state. + +## June 2026: Runtime Test-Time Compute Controls + +Repo changes: + +- Exposed the v50 `CognitiveLeapExpertHead` test-time compute path in the terminal and web chat runtimes: + - `--reasoning_cycles` / `/cycles` selects extra recursive inference depth. + - `--adaptive_compute` / `/adaptive` enables convergence early-exit. + - `--adaptive_exit_tol` / `/exit_tol` controls the early-exit tolerance. + - `--adaptive_exit_entropy` / `/exit_entropy` controls the entropy-based + early-exit threshold. +- Added `forward_with_runtime_compute(...)` in `source/chat_app.py` and synced it to `runtime_python/chat_app.py`. The helper introspects the model `forward(...)` signature, so these knobs are applied only to variants that support them and are ignored safely by older checkpoints. +- Runtime requests are capped at `MAX_RUNTIME_REASONING_CYCLES` to prevent + accidental runaway API/CLI calls while still allowing deeper-than-default + controlled experiments. +- Updated `source/chat_web_app.py` and `runtime_python/chat_web_app.py` to accept per-request `reasoning_cycles`, `adaptive_compute`, `adaptive_exit_tol`, and `adaptive_exit_entropy`, and to return `compute` diagnostics plus `timing_ms.cycles_used`. +- `compute` diagnostics now include `ponder_cost`, `consistency_loss`, and + `gating_entropy` when the active model exposes those buffers. +- Added `test_runtime_compute_controls.py` covering supported-model forwarding, legacy-model no-op behavior, `/api/chat` payload forwarding, cycle capping, entropy-threshold forwarding, and diagnostics extraction. +- Added a compute-sweep experiment path in `source/chat_web_app.py` and + `runtime_python/chat_web_app.py`: `/api/compute_sweep` compares multiple + reasoning-cycle budgets for the same draft prompt without mutating chat + history, and both browser UIs expose it through a `Sweep` button. + +Rationale: + +- Recent adaptive test-time compute and latent/recurrent reasoning work argues for controllable inference budgets instead of a single fixed-depth runtime. The v50 architecture already had that mechanism; this change makes it reachable from the actual app surfaces and observable during local experiments. + +## July 2026: Adaptive Runtime Routing Feedback + +Additional primary sources reviewed: + +1. ParetoBandit: Budget-Paced Adaptive Routing for Non-Stationary LLM Serving + https://arxiv.org/abs/2604.00136 +2. Adaptive LLM Routing under Budget Constraints + https://arxiv.org/abs/2508.21141 +3. LLM Routing with Dueling Feedback + https://arxiv.org/html/2510.00841v1 +4. SeqRoute: Global Budget-Aware Sequential LLM Routing + https://arxiv.org/abs/2605.25424 + +Repo changes: + +- Added recency-weighted route-feedback summaries with geometric forgetting so + recent route quality regressions can override stale positive feedback. +- Added route economics capture for agent-mode usage, including estimated and + actual model-call, tool-call, cost, and latency units. +- The auto agent router now uses session-scoped feedback to downgrade routes + with recent adaptive regressions, pace high-cost routes under budget profiles, + and prefer neighboring routes with stronger quality-cost evidence when enough + relevant samples exist. +- Added optional target-route session pacing via + `auto_session_budget_target_routes`, so a fixed session budget can be spread + across an expected number of future route decisions instead of only reacting + when the remaining budget is already tight. +- Added a non-mutating route-plan preview path via `preview_route_plan(...)` and + `/api/route_plan`, allowing the web UI to inspect the selected route, + estimated cost, feedback adjustment, and session-budget adjustment before any + model inference or memory writes happen. +- Route-plan previews now include `route_alternatives`, a compact cost estimate + table for each eligible route mode, so the UI can expose the local routing + frontier before spending inference calls. +- Added `route_frontier` preview metadata that annotates candidate routes with + heuristic/adaptive quality estimates, budget-fit status, Pareto-frontier + labels, stable frontier ranks, and a recommended route for the current budget + profile without changing the actual execution route. +- Route-frontier budget state now separates remaining-session fit from + route-horizon pacing fit, exposing `remaining_cost_units`, + `pacing_cap_cost_units`, and `effective_cap_cost_units` so the UI can + distinguish hard budget exhaustion from deliberate pacing preservation. +- Candidate frontier rows now include a `budget_blocker` classification + (`remaining_budget`, `pacing_cap`, or none), and all-over-cap previews use + separate recommendation reasons for hard budget exhaustion versus route-horizon + pacing preservation. +- Candidate frontier rows now expose `estimated_quality_cost_score`, + `quality_cost_source`, and `quality_evidence_status`; complete adaptive + quality-cost evidence can drive balanced/fast/deep route-frontier ranking, + while incomplete cost samples, regressions, and non-finite telemetry fall back + to heuristic cost-adjusted estimates. +- Route-frontier previews now distinguish raw `pareto_frontier` candidates from + `budget_feasible_pareto_frontier` candidates, so a high-quality route can + remain visible as an efficient tradeoff without being mistaken for a route that + fits the current session or pacing budget. +- Merged the missing v41/v42 Omni Collective continuation source from the + sibling `Supermix_27` checkout into the active `Supermix` tree: v41/v42 engine + wrappers, blueprint/prep scripts, smoke/frontier train scripts, run launchers, + and focused tests. Generated build outputs, checkpoints, dataset shards, and + stale package artifacts were intentionally left out of the merge. +- Hardened route-planning estimates with bounded numeric setting coercion for + loop budgets, web-search budgets, and web-search result limits so malformed UI + or API values cannot turn a dry-run route preview into a server error. +- The web app exposes route quality, average cost, latency, feedback controls, + route health, dry-run route planning, and a distinct Pareto-style trace pill + for cost-aware adaptive route changes. + +Notes: + +- This is a compact engineering adaptation of online routing ideas, not a full + contextual-bandit implementation. It stays prompt/session scoped, blocks + unrelated recent-feedback fallback from changing new prompts, requires actual + cost evidence for adaptive quality-cost comparisons, and lets explicit session + budget pacing make the final route-depth decision. The route-count horizon is a + lightweight global-budget adaptation inspired by sequential-routing work; it + does not attempt to solve the full online planning problem. + +## July 2026: v51 Cognitive Leap Ultra + Completion Build + +Primary sources reviewed: + +1. MiMo-V2-Flash Technical Report (efficient MoE, hybrid attention, and MTP) + https://arxiv.org/abs/2601.02780 +2. Scaling Test-time Compute for LLM Agents + https://arxiv.org/abs/2506.12928 +3. Rethinking Optimal Verification Granularity for Compute-Efficient Test-Time + Scaling + https://arxiv.org/abs/2505.11730 +4. Trust but Verify! A Survey on Verification Design for Test-time Scaling + https://arxiv.org/abs/2508.16665 + +Repo changes and verified build path: + +- Added the v51 `cognitive_leap_ultra_expert`: a recurrent mixture-of-cores, + cross-latent attention, deep-supervised refinement head with ACT, + convergence, and entropy-based inference halting. +- Added a controlled v51 trainer/benchmark and a metadata materializer so a + trained checkpoint can move directly into the terminal and browser runtimes. +- Kept per-request reasoning budgets, adaptive exits, non-mutating compute + sweeps, and compute diagnostics available through the canonical runtime and + its root compatibility entry point. +- Completed a bounded CPU training/build run, reloaded the emitted checkpoint, + and exercised it through the browser UI and `/api/chat` runtime path. +- Replaced linear legacy scheduler replay in the v8 resume fallback with a + constant-time cursor restore, preserving the learning-rate position without + dummy optimizer updates or scheduler-order warnings. + +Research boundary: + +- The v51 head is a compact adaptation of sparse expert routing and adaptive + test-time reasoning for this classifier-backed architecture. It is not a + reproduction of MiMo-V2-Flash's full autoregressive hybrid SWA/global-attention + backbone or its MTP decoding stack. +- Verifier-guided candidate search and larger-scale autoregressive training are + useful benchmark-gated follow-ons, not capabilities claimed by this build. + +## July 2026: Prediction-Stability Verifier for Local Test-Time Scaling + +Additional primary sources reviewed: + +1. Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained + Policy Optimization + https://arxiv.org/abs/2604.14853 +2. ThinkBooster: A Unified Framework for Seamless Test-Time Scaling of LLM + Reasoning + https://arxiv.org/abs/2606.06915 +3. LATTS: Locally Adaptive Test-Time Scaling + https://arxiv.org/abs/2509.20368 +4. Step-level Verifier-guided Hybrid Test-Time Scaling for Large Language Models + https://arxiv.org/abs/2507.15512 + +Repo changes: + +- Added an inference-only full-output verifier to the v51 recurrent head. At + each cycle it constructs the output that would be returned if reasoning + stopped at that point, then tracks the top prediction and confidence range. +- Adaptive inference can now stop when the prediction remains unchanged for a + configurable patience window and confidence drift stays below a configurable + tolerance. This criterion is separate from latent convergence, low entropy, + and ACT remaining-mass exits. +- Added `prediction_stability_patience` and `prediction_stability_tol` across + the shared runtime helper, terminal commands, source web API/UI, packaged + runtime web API/UI, materialized metadata, diagnostics, and launch profile. +- Compute diagnostics now report `exit_reason`, `prediction_streak`, and + `prediction_confidence_delta`, making the allocation decision inspectable. +- Restored `context_mix_v4` in the packaged `runtime_python` feature pipeline + and aligned its optimized imperative matcher with source semantics, so the + same v51 metadata no longer silently downgrades to `context_mix_v3` outside + the source launcher. +- Added `source/benchmark_v51_prediction_stability.py`, which compares fixed + and adaptive policies as request-sized inference rather than hiding local + exits inside one large batch. + +First serving-style checkpoint result (64 held-out requests, CPU): + +| policy | accuracy | mean cycles | mean latency | +|---|---:|---:|---:| +| fixed 3 cycles | 0.125 | 3.0 | 84.828 ms | +| stability verifier, max 8 | 0.125 | 2.0 | 72.167 ms | + +- Prediction agreement was 100%; every request exited with + `prediction_stable` at cycle 2. +- This reduced recurrent cycles by 33.3% and measured mean latency by 14.9% on + that serving slice without changing accuracy. The result is a runtime + efficiency result on a small synthetic checkpoint, not a general reasoning + quality claim. + +## July 2026: Preference-Aware Routing with Calibrated Evidence + +Additional primary sources reviewed: + +1. RouteLLM: Learning to Route LLMs from Preference Data (ICLR 2025) + https://proceedings.iclr.cc/paper_files/paper/2025/hash/5503a7c69d48a2f86fc00b3dc09de686-Abstract-Conference.html +2. Learning to Route LLMs from Bandit Feedback: One Policy, Many Trade-offs (BaRP, 2025) + https://arxiv.org/abs/2510.07429 +3. Learning to Route LLMs from Implicit Cost-Performance Preferences via Meta-Learning (2026) + https://arxiv.org/abs/2606.06178 +4. UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing (2026) + https://arxiv.org/abs/2605.18796 +5. Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing (2026) + https://arxiv.org/abs/2607.09015 + +Repo changes: + +- Replaced binary-only route corrections with a controlled feedback taxonomy: + `good`, `bad_quality`, `needs_deeper`, `too_costly`, and `too_slow`. +- Split feedback into quality, depth, cost-pressure, and latency-pressure axes. + Cost or latency complaints therefore no longer teach the router that an + otherwise useful answer was low quality, while `needs_deeper` can request one + higher route without poisoning quality estimates. +- Added recency-weighted effective sample sizes and Wilson-shaped heuristic + evidence bands for observed route quality. These are descriptive recent-signal + gates, not nominal 90% confidence intervals: with decay `0.6`, Kish effective + sample size approaches a ceiling of `(1 + 0.6) / (1 - 0.6) = 4` even as raw + feedback grows. Risk-adjusted quality/cost comparisons use these bands when + both neighboring modes have established recent evidence. + Emerging mean evidence can still influence routing after the existing + minimum-evidence and quality-delta gates; sparse or unrelated fallback + feedback cannot directly reroute an unmatched prompt. +- Exposed confidence bounds, evidence status, preference direction, cost, and + latency in route health. The dry-run route planner additionally exposes the + Pareto frontier and risk-adjusted scores. +- Connected the main multimodel runtime to v51 reasoning-cycle and adaptive-exit + controls. Auto routing now supplies its prompt-derived compute budget to the + selected backend, while explicit per-request settings still take precedence. +- Added checkpoint `runtime_defaults` handling with clean reload semantics and + surfaced compute telemetry (cycles used, exit reason, and prediction drift) + in the multimodel trace. +- Made absent standalone-CLI compute flags inherit checkpoint metadata instead + of accidentally overriding it with library defaults; explicit command-line + values still have highest precedence. +- Layered exact lazy-greedy marginal coverage onto the Qwen SFT and preference + rarity priors. SFT pair-budget and feasibility-aware token-budget coverage, + plus preference pair-budget coverage, now use heap upper bounds to reduce + redundant rescoring while preserving deterministic quality-anchored choices. +- Replaced the packaged model-variant proxy with a self-contained source + snapshot, added an isolated-import regression test and deterministic sync + checker, and wired the runtime parity gates into GitHub Actions. + +Research boundary: + +- This is an inspectable, session-scoped engineering adaptation of + preference-aware and uncertainty-aware routing. It is not a trained + generalist router, a Bayesian posterior over model utility, or an online + contextual-bandit regret guarantee. +- Recency-weighted Wilson-shaped bounds are used as heuristic evidence gates, + not as nominal coverage or proof that sparse, adaptive user feedback is + statistically independent. Larger offline replay sets and randomized policy + evaluation remain benchmark-gated follow-ons. + +## July 2026: Honest Route Policy Evidence Lab + +Additional primary sources reviewed: + +1. Confident Off-Policy Evaluation and Selection through Self-Normalized Importance Weighting (AISTATS 2021) + https://proceedings.mlr.press/v130/kuzborskij21a.html +2. Anytime-valid Off-policy Inference for Contextual Bandits (ICML 2021) + https://proceedings.mlr.press/v139/karampatziakis21a.html +3. Improved Offline Contextual Bandits with Second-Order Bounds: Betting and Freezing (COLT 2025) + https://proceedings.mlr.press/v291/ryu25a.html +4. Oracle-Efficient Pessimism: Offline Policy Optimization in Contextual Bandits (AISTATS 2024) + https://proceedings.mlr.press/v238/wang24a.html +5. Conservative Contextual Bandits with Interleaving (ICML 2023) + https://proceedings.mlr.press/v202/takemura23a.html + +Repo changes: + +- Added immutable shadow threshold profiles (`efficiency`, `balanced`, and + `quality_first`) plus a pure replay module in `source/route_policy_lab.py`. +- Replay joins usage and feedback only by exact server route ID. It reports + coverage, candidate agreement, observed approval, cost, and latency only for + rows where the candidate action matches the action that actually ran. + Changed actions receive no imputed reward. +- The runtime now returns the same UUID route ID that it writes to usage + memory. Browser feedback sends only that ID, the controlled feedback intent, + and an optional note; prompt, mode, model, policy, and economics are recovered + from the server usage row. Repeated feedback for one route becomes an + idempotent revision instead of an accidental duplicate. +- Route logs now carry policy/version and feature-schema identifiers, the + safety-filtered eligible set, exact post-filter action probabilities, + decision context, and the chosen action's logging propensity. The current + policy is explicitly marked deterministic with a one-hot vector. +- Added a read-only Route Policy Lab panel and API. Its promotion gate remains + `shadow_only` and blocks automatic promotion when valid randomized overlap is + absent. Even propensity-ready rows require a separately validated off-policy + estimator and review; readiness alone is not an estimate. +- Session JSON writes now replace an on-disk temporary file atomically, reducing + the chance that an interrupted write destroys the policy evidence ledger. + +Research boundary: + +- Historical deterministic logs do not identify outcomes for unchosen routes. + The Policy Lab therefore labels its results associational and does not claim + inverse-propensity, doubly robust, causal, or high-probability policy value. +- Explicit `needs_deeper`, `too_costly`, and `too_slow` feedback remains useful + as direct one-step user intent. It is separate from policy-promotion evidence. +- A future opt-in explorer must randomize only between safety-filtered adjacent + routes, record exact post-filter propensities before execution, retain failed + routes, and evaluate with session-clustered pessimistic bounds before any + promotion path can be enabled. + +## July 2026: Durable Failure-Aware Route Evidence + +Additional primary sources reviewed: + +1. Off-Policy Evaluation for Recommendations with Missing-Not-At-Random Rewards (2025 preprint) + https://arxiv.org/abs/2502.08993 +2. Conservative Contextual Bandits: Beyond Linear Representations (ICLR 2025) + https://proceedings.iclr.cc/paper_files/paper/2025/hash/dbca58f35bddc6e4003b2dd80e42f838-Abstract-Conference.html +3. Off-Policy Evaluation for Ranking Policies under Deterministic Logging Policies (ICLR 2026) + https://arxiv.org/abs/2603.21485 +4. Logging Policy Design for Off-Policy Evaluation (2026 preprint) + https://arxiv.org/abs/2605.15108 +5. Anytime-valid Optimal Policy Identification (2026 preprint) + https://arxiv.org/abs/2606.17515 + +Repo changes: + +- Added `source/route_policy_ledger.py`, a schema-versioned SQLite WAL ledger + with short transactions and atomic per-session sequence allocation. Only a + domain-separated session hash is stored; raw session IDs and prompt text are + excluded from the durable database. +- Route decisions now commit an `inflight` row after final capability, safety, + and session-budget filtering but before inference. Successes and exceptions + transition that row separately, so backend, economics, memory, and + serialization failures no longer disappear from the evidence base. +- Abrupt process termination intentionally leaves an `inflight` row rather than + inventing a failure outcome. The Policy Lab exposes completed, failed, and + in-flight counts separately. +- Explicit feedback is appended as an idempotent, route-ID-keyed revision. A + missing response remains `unknown`; approval is always shown alongside + terminal feedback coverage instead of treating silence as neutral or bad. +- The existing JSON store remains a bounded compatibility mirror for adaptive + route preferences. The durable ledger is the lifecycle source of truth and + the Policy Lab remains read-only and `shadow_only`. + +Research boundary: + +- The new ledger fixes survivorship and join integrity, but the current route + probabilities are still deterministic one-hot vectors. It therefore improves + traceable associational diagnostics without enabling IPS, SNIPS, doubly + robust, causal, regret, or anytime-valid policy-superiority claims. +- Exploration remains disabled (`epsilon = 0`). Known post-filter randomized + probabilities, overlap, versioned bounded outcomes, and a separately modeled + feedback-observation process are prerequisites for future causal evaluation. +- Runtime reliability, quality, latency, and cost remain separate outcomes. A + cheap or fast failure can never compensate for a reliability or safety floor. + +## July 2026: Durable Support Envelope + Policy Readiness Certificate v2 + +Additional primary sources reviewed: + +1. Anytime-valid Optimal Policy Identification (2026 preprint) + https://arxiv.org/abs/2606.17515 +2. Off-Policy Evaluation for Recommendations with Missing-Not-At-Random Rewards (2025 preprint) + https://arxiv.org/abs/2502.08993 +3. Off-Policy Confidence Sequences (COLT 2025) + https://proceedings.mlr.press/v291/ryu25a.html +4. Supplementary Outcomes for Off-Policy Evaluation (ICLR 2025) + https://proceedings.iclr.cc/paper_files/paper/2025/hash/098491b37deebbe6c007e69815729e09-Abstract-Conference.html + +Design contract: + +- A versioned durable support envelope binds each route decision to its final + post-filter candidate set, normalized logging distribution, chosen + propensity, policy and feature-schema identities, and canonical support + fingerprints before execution. It is an audit contract; it does not enable a + different route-selection policy. +- A read-only readiness certificate reports schema validity, lifecycle and + feedback-observation coverage, empirical action overlap, propensity floors, + and diagnostic effective sample size (ESS). Each failed check remains an + explicit reason code instead of being averaged into a single favorable score. +- The certificate fails closed unless evidence comes from a reconciled durable + lifecycle window, meets its declared valid-route floor, and reproduces the + versioned support envelope, chosen propensity, assignment commitment, and + candidate/distribution fingerprints. The bounded JSON mirror is descriptive + compatibility data only and can never satisfy lifecycle readiness. +- Explicit feedback now commits to SQLite before the JSON compatibility mirror. + Content-derived idempotency coalesces only an immediately adjacent identical + retry; an explicit request ID is matched across the route's full revision + history. A mirror failure returns an accepted-but-pending reconciliation + status instead of losing or duplicating the durable acknowledgement. +- The current behavior policy remains deterministic with one-hot propensities. + Consistent with recent work on deterministic logging, these rows cannot + identify counterfactual outcomes for unchosen routes and therefore cannot + produce an off-policy value estimate. +- Missing feedback remains `unknown`. Because reward observation can be + missing-not-at-random (MNAR), silence is neither a negative label nor evidence + that observation is independent of route choice; observation coverage is + reported separately by lifecycle state and action. A scalar observation + propensity is not enough to clear this gate unless its observation policy and + outcome definition are both versioned. + +Research and deployment boundary: + +- The certificate computes no IPS, SNIPS, doubly robust, causal, regret, or + anytime-valid policy-value estimate. Passing structural checks means only + that a separately validated evaluator could receive better-formed input. +- Exploration remains disabled (`epsilon = 0`), deployment remains + `shadow_only`, and automatic promotion remains forbidden. Any future opt-in + randomized policy requires an independently reviewed safety-filtered + assignment mechanism, bounded outcome contract, and evaluation protocol. +- Propensity, route-count, session-count, coverage, and ESS thresholds are + diagnostics for a declared cohort. They expose weak support and unstable + weights, but no fixed ESS threshold is a universal statistical guarantee or + proof of policy superiority. + +## July 2026: Fail-Closed Route Outcome Contract v1 + +Additional primary sources reviewed: + +1. Off-Policy Evaluation under Nonignorable Missing Data (ICML 2025) + https://proceedings.mlr.press/v267/wang25dt.html +2. A General Framework for Off-Policy Learning with Partially-Observed Reward + (ICLR 2025) + https://proceedings.iclr.cc/paper_files/paper/2025/hash/098491b37deebbe6c007e69815729e09-Abstract-Conference.html +3. Clarifying Uncertainty Quantification in Off-Policy Evaluation: Beyond + Effective Sample Sizes, Towards Confidence Intervals (ICML 2026 DEMO) + https://openreview.net/pdf?id=FuuLorZ6NQ +4. Beyond the Training Distribution: Evaluating Predictions Under Distribution + Shift and Selection Bias (2026 preprint) + https://arxiv.org/abs/2606.14506 +5. Anytime-valid Optimal Policy Identification (2026 preprint) + https://arxiv.org/abs/2606.17515 +6. Logging Policy Design for Off-Policy Evaluation (2026 preprint) + https://arxiv.org/abs/2605.15108 +7. Off-Policy Evaluation for Ranking Policies under Deterministic Logging + Policies (ICLR 2026) + https://arxiv.org/abs/2603.21485 + +Implementation contract: + +- Readiness now checks the entire fixed-as-of durable population, not only the + feedback join. Missing or duplicate route IDs, orphan feedback, unevaluable + usage rows, inconsistent terminal states, or a mismatch between chosen and + executed route fail closed with explicit population/execution reason codes. +- A versioned, prompt-free decision-record fingerprint binds the immutable + policy identity, decision context, selected action, eligible set, and support + projection committed before execution. The snapshot recomputes it rather than + trusting stored metadata; migrated v1 records remain explicitly legacy. +- Session JSON filenames combine a safe display slug with a full session digest. + Legacy files migrate only when their embedded session identity matches, while + atomic replacement, transactional schema migration, and one-transaction + evidence snapshots prevent collisions or mixed-revision projections. +- Feedback retries without an explicit request ID coalesce only when the newest + revision has identical content. Explicit request IDs remain globally stable + across the same route's revision history, so an older retry cannot create a + new revision and reusing its ID with different content fails closed. +- Route Outcome Contract v1 precommits versioned definitions and observation + semantics for route success, measured cost, measured latency, and user quality. + Completion records only outcomes actually available; failure records failure + and measured elapsed time without fabricating cost, while quality remains + unknown until explicit quality feedback is observed. Canonical hashes and + decision-start timing are reverified on replay, and missing, late, posthoc, or + tampered contracts fail the existing outcome-evidence readiness check. +- Maturity telemetry reports per-outcome precommit/observation coverage and the + share of legacy posthoc contracts. It is diagnostic evidence hygiene for a + declared fixed-as-of cohort, not an estimator or promotion score. + +Research and deployment boundary: + +- Fixed-as-of coverage telemetry is not an observation model. It does not + identify why quality was observed, correct selective labels or covariate + shift, or turn missing feedback into a reward; supplementary outcomes require + their own validated relationship to the target outcome. +- ESS diagnoses weight concentration. It is not confidence-interval width, + empirical coverage, estimator error, or a universal cross-estimator measure + of uncertainty. +- The lab still computes no OPE or policy-value estimate, and it enables no + automatic promotion. Deployment remains `shadow_only`, including when all + structural checks pass. +- Contracts reconstructed for legacy decisions after their outcomes are known + are posthoc compatibility records. They cannot satisfy pre-execution contract + maturity, causal identification, or anytime-valid policy-selection claims. + +## July 2026: Bounded-Exposure Adjacent-Route Rehearsal v1 + +Additional primary sources reviewed: + +1. Logging Policy Design for Off-Policy Evaluation (2026 preprint) + https://arxiv.org/abs/2605.15108 +2. Conservative Contextual Bandits: Beyond Linear Representations (ICLR 2025) + https://proceedings.iclr.cc/paper_files/paper/2025/hash/dbca58f35bddc6e4003b2dd80e42f838-Abstract-Conference.html +3. Clarifying Uncertainty Quantification in Off-Policy Evaluation: Beyond + Effective Sample Sizes, Towards Confidence Intervals (ICML 2026 DEMO) + https://openreview.net/pdf?id=FuuLorZ6NQ +4. Off-Policy Evaluation under Nonignorable Missing Data (ICML 2025) + https://proceedings.mlr.press/v267/wang25dt.html +5. Anytime-valid Optimal Policy Identification (2026 preprint) + https://arxiv.org/abs/2606.17515 + +Implementation contract: + +- The read-only planner consumes the final capability- and budget-filtered route + candidates plus a strict source contract. That contract binds the source + policy ID/version, feature schema, support schema, candidate-set hash, + deterministic distribution hash, and Route Outcome Contract schema into the + draft charter hash. +- The planner rehearses a fixed incumbent-heavy distribution over the incumbent + and at most its two nearest feasible neighbors; excluded and non-adjacent + routes receive no probability. `adjacent` is an ordinal exposure heuristic, + not a target-aware optimal logging-policy claim. +- A 10% alternate allocation is split across the enrolled neighbors. With two + neighbors this gives each a 5% planned propensity, matching the Policy Lab's + current structural probability floor. This is a design diagnostic, not a + claim that 5% is statistically sufficient. +- The canonical hash also binds the ordered candidate/exclusion projection, + exact probability vector, repeated-stratum horizon, response-rate and + target-label scenarios, resource envelope, and fail-closed causal boundaries. + It contains no prompt or session text and performs no ledger, memory, RNG, or + inference operation. +- A deterministic nonce-based primitive exists only to inspect replayable draw + mechanics. Its output is a non-ledger rehearsal receipt and non-ledger support + proposal. Caller-selected nonces are grindable; no seed was committed before + assignment, no immutable assignment unit was sealed, and the receipt is not a + randomization commitment or executed propensity record. The runtime never + calls this primitive. +- The browser and terminal surfaces expose the same planning semantics. Both + label the output rehearsal-only, keep execution and evidence writes off, and + expose expected cost, latency-tier exposure, alternate propensity, expected + traffic, and a simultaneous label-traffic scenario for a declared response + rate. + +Research and deployment boundary: + +- For one alternate the target-label forecast inverts an exact binomial tail. + For two alternates it inverts the exact joint multinomial probability that + every alternate reaches the target; the displayed confidence is simultaneous, + not the weaker marginal confidence for one preselected alternate. +- Expected assignments and the constant-response scenario hypothetically repeat + the same prompt-specific support stratum. They are not a campaign forecast, + statistical power, estimator precision, confidence in policy value, or + evidence that feedback is missing at random. A live campaign would require a + frozen population/context distribution and an identified observation process. +- ESS remains an overlap and weight-concentration diagnostic. It is not a + universal uncertainty measure, and the planner computes no ESS, IPS, SNIPS, + doubly robust estimate, confidence sequence, regret bound, or policy value. +- Reserving most planned mass for the incumbent is an operational exposure cap. + It is not the high-probability baseline-performance guarantee established by + conservative contextual-bandit algorithms, whose modeling assumptions are + not satisfied by the current heuristic route-quality signals. +- Activation is explicitly blocked until a target-policy class and estimand, + outcome/observation/maturity contract, population scope, preassignment seed + commitment and unique immutable unit, session carryover/interference strategy, + resource and stopping rules, and external estimator/review are sealed. The + cited logging-policy design is one-step; Supermix's stateful session behavior + is not assumed away. +- Rehearsed probabilities are never written as executed propensities. Existing + deterministic rows remain a separate `auto-route-v2` cohort, automatic + promotion remains forbidden, and deployment remains `shadow_only`. + +## July 2026: Stateful Route Experiment Preflight v1 + +Additional primary sources reviewed: + +1. Anytime-Valid Off-Policy Inference for Contextual Bandits (Journal of Data + Science, 2024) + https://jds.acm.org/files/JDS_Issue3_Paper1.pdf +2. Semiparametric Efficient Inference in Adaptive Experiments (CLeaR 2024) + https://proceedings.mlr.press/v236/cook24a.html +3. Cluster-Adaptive Network A/B Testing (JMLR 2024) + https://www.jmlr.org/papers/v25/22-0192.html +4. Data-Driven Switchback Experiments (2024 preprint) + https://arxiv.org/abs/2406.06768 +5. Clustered Switchback Experiments (2023 preprint) + https://arxiv.org/abs/2312.15574 +6. Sequentially-Rerandomized Switchback Experiments (2026 preprint) + https://arxiv.org/abs/2604.02489 +7. Logging Policy Design for Off-Policy Evaluation (2026 preprint) + https://arxiv.org/abs/2605.15108 +8. Off-Policy Evaluation under Nonignorable Missing Data (ICML 2025) + https://proceedings.mlr.press/v267/wang25dt.html + +Implementation contract: + +- `route_policy_protocol.py` wraps one or more valid adjacent-route rehearsal + plans without changing explorer-v1 or invalidating its strict hashes. The + wrapper is order-invariant across unique support strata and requires one + common source policy/schema cohort. +- The protocol hash freezes a versioned target-policy threshold class, admitted + support-stratum hashes, a prompt-free population rule and session-hash cluster + schema, fixed cluster/route ceilings, an outcome-independent analysis schedule, + the four Route Outcome Contract hashes, and every stateful design declaration. +- Route-level campaign randomization is screened out in v1. The only design + screens are sticky session-cluster assignment and clustered switchbacks. The + latter exposes block and washout declarations; the protocol draft itself does + not assign either mode. +- Carryover, interference, and temporal variation use closed enums rather than + free text. Unknowns remain incomplete; declarations remain unvalidated. A + compatible declaration set can reach only `assumptions_declared_unvalidated`, + never a causal-design certificate or activation state. +- A seed commitment can be bound as a lowercase digest, but the protocol draft + never generates or reveals its seed and does not implement HMAC assignment. + The separate shadow registry described below is the only component that opens + this later workflow; caller-selected nonces remain prohibited. +- The browser and `route_policy_protocol_cli.py` expose the same canonical + preflight. The CLI can audit hash integrity and fail-closed boundaries. The + Windows Studio packaging contract includes this console as + `SupermixRouteStudy.exe` and binds module/schema hashes in + `studio_runtime_manifest.json`. + +Research and deployment boundary: + +- Stateful routing can change later context, memory, and user behavior. A + one-turn contextual-bandit estimator is therefore not silently extended to a + session-level causal claim. Sticky clustering reduces within-session treatment + switching; it does not prove that the cluster captures shared memory or all + cross-cluster exposure. +- Switchback frequency and washout are bias-variance and identification choices. + Declaring them does not establish finite carryover, rapid mixing, a valid + interference graph, or unbiased exposure contrasts. +- Anytime-valid inference needs its own theorem-level conditions, including + predictable exact propensities, support, bounded outcomes, and a frozen policy + family/alpha budget. Optional-stopping validity would not itself establish + causal identification, missing-at-random feedback, or deployment safety. +- Missing ratings remain unknown. The cited MNAR method requires specific + dropout, state-sufficiency, positivity, response-model, and shadow-variable + assumptions that sparse voluntary ratings do not automatically satisfy. +- The draft retains all eight activation blockers. It performs no I/O, ledger + write, assignment, inference, OPE, winner selection, or automatic promotion. + Independent scientific review and external implementation remain mandatory. + +## July 2026: Portable Route Protocol Review Bundle v1 + +Additional integrity and transparency sources reviewed: + +1. An Architecture for Trustworthy and Transparent Digital Supply Chains, + RFC 9943 (June 2026) + https://www.rfc-editor.org/rfc/rfc9943.html +2. CBOR Object Signing and Encryption (COSE) Receipts, RFC 9942 (June 2026) + https://www.rfc-editor.org/rfc/rfc9942.html +3. Certificate Transparency Version 2.0, RFC 9162 + https://www.rfc-editor.org/rfc/rfc9162.pdf +4. Rekor transparency-log architecture + https://docs.sigstore.dev/logging/overview/ +5. Logging Policy Design for Off-Policy Evaluation (2026 preprint) + https://arxiv.org/abs/2605.15108 + +Implementation contract: + +- The compact protocol audit now checks strict nested schemas and frozen v1 + semantics for source inventories, population declarations, stateful design, + outcome contracts, stopping rules, randomness, blocker statuses, external + evaluation, prompt-free guarantees, and causal boundaries. Rehashing a draft + after enabling outcome-dependent stopping, promotion, validation claims, or + altered outcome definitions no longer passes. +- Compact drafts intentionally omit complete source plans, so their verifier is + labeled `structural_without_source_plans`. A digest alone cannot prove that a + claimed study hash came from the canonical explorer. +- `route-study-review-bundle-v1` carries every canonical prompt-free source plan, + the complete closed builder option set, and the resulting protocol. Full audit + validates each explorer plan, canonicalizes ordering, rebuilds the protocol, + requires exact equality, and then verifies the separate bundle hash. Its + verification label is `full_source_bound_reconstruction`. +- Runtime and browser bundle endpoints accept only the shared closed protocol + input schema. Prompt, raw session, and free-text fields fail closed. Browser + requests are capped at 2 MiB and 100 strata; the core/console contract retains + the 1,000-stratum ceiling. +- The Studio browser keeps source strata only in ephemeral client memory. It can + add/remove compatible strata, build/download a bundle, and import either a + bundle or closed build input for server-side reconstruction. It never pools + heterogeneous strata without predeclared population weights. +- `SupermixRouteStudy.exe --example-bundle` and `--audit-bundle` expose the same + contract, and Windows CI reconstructs a frozen two-stratum bundle. + +Integrity and deployment boundary: + +- Full reconstruction proves internal semantic conformance, not authorship or + historical existence. RFC 9943, RFC 9942, and Certificate Transparency + distinguish an artifact digest from a signed receipt backed by a verifiable + append-only data structure. Supermix has no signature, trusted timestamp, + external witness, inclusion proof, or consistency proof in this increment. +- A valid bundle still does not validate cluster independence, carryover, + interference, feedback missingness, support under live execution, or any + causal estimand. Logging-policy choice can materially change OPE error; an + internally consistent review artifact cannot repair unsupported actions. +- The portable bundle by itself seals or registers nothing. It generates no + seed, assigns no cluster, executes no route, writes no evidence, estimates no + policy value, and leaves all eight activation blockers active. The separate + shadow registry below can consume only this source-bound artifact; it does not + turn bundle verification into a transparency receipt or activation approval. + +## July 2026: Shadow Whole-Policy Commitment/Reveal Registry v1 + +Additional primary sources reviewed: + +1. JSON Canonicalization Scheme (JCS), RFC 8785 + https://www.rfc-editor.org/rfc/rfc8785.html +2. HMAC-based Extract-and-Expand Key Derivation Function (HKDF), RFC 5869 + https://www.rfc-editor.org/rfc/rfc5869.html +3. An Architecture for Trustworthy and Transparent Digital Supply Chains, + RFC 9943 (June 2026) + https://www.rfc-editor.org/rfc/rfc9943.html +4. CBOR Object Signing and Encryption (COSE) Receipts, RFC 9942 (June 2026) + https://www.rfc-editor.org/rfc/rfc9942.html +5. Analysis of Two-Stage Rollout Designs with Clustering for Causal Inference + under Network Interference (AISTATS / PMLR 258, 2025) + https://proceedings.mlr.press/v258/cortez-rodriguez25a.html +6. Randomization Tests in Switchback Experiments (2026 preprint) + https://arxiv.org/abs/2602.23257 + +Implementation contract: + +- `route_policy_shadow_registry.py` owns a schema-v1 SQLite database separate + from the executed route-decision ledger. The Studio runtime locates it at + `memory/route-policy-shadow-registry.sqlite3`; no registry artifact is + eligible for Policy Lab/OPE input or an executed logging-support row. +- Sealing first performs full source-bound reconstruction of a + `route-study-review-bundle-v1`. The resulting design binding and assignment + manifest freeze exactly two 50/50 whole-policy arms: + `incumbent_source_policy`, bound to the source-policy cohort, and + `candidate_target_policy`, bound to the target-policy class. Prompt-specific + `eligible_actions` remain support-stratum metadata and cannot silently become + cluster-level treatment arms. +- `SupermixRouteShadow.exe seal` obtains 256 bits from the operating-system + CSPRNG, commits the seed to the design, and writes the seed capsule as a + separate exclusively-created file before the public package enters SQLite. + The registry contains no seed material before explicit post-closure reveal; + command output never prints the seed. Capsule access control, backup, transfer, + and independent custody remain operator responsibilities and are not verified. +- `commit` accepts one raw cluster identifier through a private input file, + derives a study-scoped HMAC pseudonym internally, and persists neither the raw + identifier nor the chosen arm. It appends only the pseudonym and an opaque + assignment-reveal commitment. Pseudonymity is not anonymity, especially once + the seed is public and an observer can test candidate identifiers. +- `close` atomically freezes the enrolled commitment count and blocks later + commitments. `reveal` is rejected until closure, verifies the seed opening, + and then persists it. Bounded `verify` batches reconstruct each whole-policy + arm and append a matched or mismatched reveal record. Registry state therefore + advances from accepting commitments, through closed and seed-revealed, to + reveal verification complete without affecting inference. +- The assignment algorithm uses RFC 5869 extract-and-expand with separate + context strings for identity and assignment keys, then HMAC-SHA-256 for the + study pseudonym and integer basis-point draw. The registry's canonical + artifact subset is informed by RFC 8785: duplicate keys, non-finite values, + floating-point registry fields, non-ASCII object keys, and integers outside + the I-JSON-safe range fail closed. This constrained profile avoids claiming + general cross-language JCS conformance for arbitrary imported bundle JSON. +- SQLite WAL, `BEGIN IMMEDIATE`, foreign keys, immutable-row triggers, closure + guards, campaign-order indexes, and a domain-separated event hash chain protect + normal concurrent local use. Read-only snapshots also audit required schema + objects plus their exact definition fingerprint, reconstruct every stored + artifact, and match event artifacts to evidence rows. Reveal verification + preflights closure and seed artifacts, validates each commitment projection, + and only reports completion after a passing whole-campaign audit. The executed + ledger independently rejects shadow/rehearsal flags, + `ledger_eligible=false`, non-`route-support-v1` envelopes, reserved shadow + assignment-commitment namespaces, and every non-null commitment outside the + closed `route-execution-assignment-v1:` namespace. The namespace is a + type boundary, not authenticated provenance against a hostile local caller. +- All mutations are local-console operations. The Studio browser has only the + read-only `GET /api/route_shadow_registry/status` endpoint and can refresh + campaign state, commitment/reveal counts, and chain verification. It has no + seal, commit, close, seed-reveal, or assignment-verification control. + CLI `status` uses SQLite read-only mode as well. The server defaults to + `127.0.0.1`; remote binding is an explicit unauthenticated operator choice. + Browser responses are `Cache-Control: no-store`, while the server reuses a + full audit only until the database or non-empty WAL signature changes. + +Research and deployment boundary: + +- A seed commitment and deterministic reconstruction show that stored reveals + match the sealed local inputs. They do not show that enrollment was complete + or unbiased, that the custodian did not inspect or regenerate material before + publication, or that seed custody was independent. No shadow assignment is a + live route choice or an executed propensity. +- The 2025 clustering analysis exhibits a bias-variance tradeoff under network + interference: cutting interference edges and balancing cluster covariates are + not generally the same objective. Binding a declared session cluster does not + validate that cluster, identify its interference graph, or license the paper's + estimator for Supermix. +- The 2026 switchback randomization-test framework requires a known assignment + mechanism and, for its causal effects, non-anticipation and a finite carryover + horizon. Supermix v1 implements sticky session-cluster shadow commitments, not + switchback execution or those tests; diagnostics and verified reveals do not + establish the required assumptions. +- RFC 9943 and RFC 9942 describe signed statements, verifiable data structures, + transparency services, and signed receipts. This registry has none of those. + A host administrator can replace the database or remove its triggers, the + local clock is untrusted, and the event chain has no external anchor, signer, + witness, inclusion proof, consistency proof, or anti-equivocation service. +- The registry performs no model inference, live assignment, route execution, + outcome collection, OPE, causal or policy-value estimation, winner selection, + activation, or automatic promotion. All existing review and activation + blockers remain in force after a clean verification. +- The two arm bindings freeze declarative policy-class manifests, not + executable runtime code or a code-artifact digest. Raw cluster identifiers + are private opaque inputs; v1 does not validate an external cluster-map + membership rule or canonicalize identifier aliases. Those controls remain + external prerequisites, not properties implied by a verified registry. + +## July 2026: v51 prediction-stability pilot and distribution-drift shadow metric + +Recent primary sources reviewed: + +1. LESS Is More: Adaptive Early Exit for Diffusion Language Models + https://arxiv.org/abs/2606.16908 +2. Stop When Reasoning Converges: Adaptive Test-Time Scaling with PUMA + https://arxiv.org/abs/2605.17672 +3. LLMRouterBench: A Massive Benchmark and Unified Framework for LLM Routing + https://arxiv.org/abs/2601.07206 +4. TwinRouterBench: A Systematic Benchmark for Dynamic LLM Routing + https://arxiv.org/abs/2605.18859 + +Implementation and pilot evidence: + +- `benchmark_v51_prediction_stability.py` now records top-k Jensen-Shannon + divergence between consecutive full-prefix output distributions. The shared + top-k support is chosen from the midpoint distribution and all remaining + probability is retained in one `other` bucket. This is diagnostic telemetry + only: it cannot trigger an exit or change a model answer. +- A CPU pilot screened five stopping configurations over three unseen seeds and + 32 examples per seed (480 requests total). Patience 2 / tolerance 0.005 kept + 96/96 prediction agreement and zero observed accuracy delta while using 2.135 + mean cycles, a 28.8% reduction from the three-cycle baseline. Patience 1 + changed one prediction and is rejected. Stricter patience/tolerance settings + increased work without improving observed agreement. +- Latency measurements were not counterbalanced, so they are screening evidence + rather than a release claim. The next gate is 512 held-out examples over eight + fresh seeds with zero disagreements, no negative per-seed accuracy delta, at + least 20% mean cycle reduction, and positive median latency reduction before + changing any runtime default. +- The research transfer remains provisional. LESS studies diffusion language + models, PUMA combines semantic convergence with answer verification, and the + router benchmarks emphasize held-out dynamic evaluation plus strong simple + baselines. Supermix therefore exposes distribution drift as a shadow metric + beside its existing output-persistence verifier rather than treating a new + paper or a small pilot as activation evidence. diff --git a/source/build_supermix_studio_desktop_exe.ps1 b/source/build_supermix_studio_desktop_exe.ps1 index 6ecbea381..c0ad2f831 100644 --- a/source/build_supermix_studio_desktop_exe.ps1 +++ b/source/build_supermix_studio_desktop_exe.ps1 @@ -1,6 +1,6 @@ param( [string]$Name = "SupermixStudioDesktop", - [string]$ModelsDir = "C:\Users\kai99\Desktop\models", + [string]$ModelsDir = "", [switch]$SkipDependencyInstall ) @@ -9,6 +9,15 @@ $ErrorActionPreference = "Stop" $RepoRoot = Split-Path -Parent $PSScriptRoot Set-Location $RepoRoot +if (-not $ModelsDir) { + $ModelsDir = if ($env:SUPERMIX_MODELS_DIR) { + $env:SUPERMIX_MODELS_DIR + } else { + Join-Path $env:USERPROFILE "Desktop\models" + } +} +$ModelsDir = (Resolve-Path -LiteralPath $ModelsDir -ErrorAction Stop).Path + $PythonExe = Join-Path $RepoRoot ".venv-dml\Scripts\python.exe" if (-not (Test-Path $PythonExe)) { throw "Expected Python environment at $PythonExe" @@ -18,6 +27,11 @@ if (-not $SkipDependencyInstall) { & $PythonExe -m pip install pywebview pyinstaller pillow sympy | Out-Host } +& $PythonExe "source\generate_studio_runtime_manifest.py" --check +if ($LASTEXITCODE -ne 0) { + throw "The checked Studio runtime manifest is stale. Regenerate and review it before packaging." +} + python "source\generate_desktop_branding.py" | Out-Host $BaseModelDir = & $PythonExe -c "import sys; sys.path.insert(0, 'source'); import qwen_chat_desktop_app as app; print(app.resolve_local_base_model_path(''))" @@ -29,6 +43,7 @@ $BaseModelDir = $BaseModelDir.Trim() $ModelsStageDir = Join-Path $RepoRoot "build\studio_models_stage" $BaseModelStageDir = Join-Path $RepoRoot "build\studio_base_model_stage" $BundleManifestPath = Join-Path $RepoRoot "output\supermix_studio_bundled_models_manifest.json" +$RuntimeManifestPath = Join-Path $RepoRoot "source\studio_runtime_manifest.json" $BundledModelKeys = @( "v40_benchmax", "omni_collective_v41", @@ -106,6 +121,9 @@ if (-not (Test-Path $SummaryPath)) { } try { + New-Item -ItemType Directory -Path "build\studio_desktop_spec" -Force | Out-Null + New-Item -ItemType Directory -Path "build\route_study_cli" -Force | Out-Null + New-Item -ItemType Directory -Path "build\route_shadow_cli" -Force | Out-Null $PyInstallerArgs = @( "-m", "PyInstaller", "--noconfirm", @@ -129,6 +147,8 @@ try { "--add-data", "$BaseModelStageDir;bundled_base_model", "--add-data", "$SummaryPath;output", "--add-data", "$BundleManifestPath;output", + "--add-data", "$RuntimeManifestPath;output", + "--specpath", "build\studio_desktop_spec", "source\supermix_multimodel_desktop_app.py" ) @@ -141,7 +161,55 @@ try { } $ExePath = Join-Path $RepoRoot "dist\$Name\$Name.exe" + $RouteStudyCliName = "SupermixRouteStudy" + $RouteStudyCliArgs = @( + "-m", "PyInstaller", + "--noconfirm", + "--clean", + "--onefile", + "--console", + "--name", $RouteStudyCliName, + "--paths", "source", + "--distpath", "dist\$Name", + "--workpath", "build\route_study_cli", + "--specpath", "build\route_study_cli", + "source\route_policy_protocol_cli.py" + ) + Write-Host "Building $RouteStudyCliName prompt-free protocol console" + & $PythonExe @RouteStudyCliArgs + if ($LASTEXITCODE -ne 0) { + throw "Route study console build failed." + } + $RouteStudyCliPath = Join-Path $RepoRoot "dist\$Name\$RouteStudyCliName.exe" + if (-not (Test-Path $RouteStudyCliPath)) { + throw "Expected route study console at $RouteStudyCliPath" + } + $RouteShadowCliName = "SupermixRouteShadow" + $RouteShadowCliArgs = @( + "-m", "PyInstaller", + "--noconfirm", + "--clean", + "--onefile", + "--console", + "--name", $RouteShadowCliName, + "--paths", "source", + "--distpath", "dist\$Name", + "--workpath", "build\route_shadow_cli", + "--specpath", "build\route_shadow_cli", + "source\route_policy_shadow_cli.py" + ) + Write-Host "Building $RouteShadowCliName shadow-only commitment console" + & $PythonExe @RouteShadowCliArgs + if ($LASTEXITCODE -ne 0) { + throw "Route shadow console build failed." + } + $RouteShadowCliPath = Join-Path $RepoRoot "dist\$Name\$RouteShadowCliName.exe" + if (-not (Test-Path $RouteShadowCliPath)) { + throw "Expected route shadow console at $RouteShadowCliPath" + } Write-Host "Build complete: $ExePath" + Write-Host "Protocol console complete: $RouteStudyCliPath" + Write-Host "Shadow registry console complete: $RouteShadowCliPath" } finally { if (Test-Path $ModelsStageDir) { Remove-Item -Recurse -Force $ModelsStageDir } diff --git a/source/build_supermix_studio_desktop_installer.ps1 b/source/build_supermix_studio_desktop_installer.ps1 index de60ab031..bd84a85e1 100644 --- a/source/build_supermix_studio_desktop_installer.ps1 +++ b/source/build_supermix_studio_desktop_installer.ps1 @@ -39,16 +39,22 @@ $ExePath = Join-Path $RepoRoot "dist\SupermixStudioDesktop\SupermixStudioDesktop $SetupPath = Join-Path $InstallerDir "SupermixStudioDesktopSetup.exe" $BundleManifestSource = Join-Path $RepoRoot "output\supermix_studio_bundled_models_manifest.json" $BundleManifestTarget = Join-Path $InstallerDir "SupermixStudioDesktopCuratedBundleManifest.json" +$RuntimeManifestSource = Join-Path $RepoRoot "source\studio_runtime_manifest.json" +$RuntimeManifestTarget = Join-Path $InstallerDir "SupermixStudioRuntimeManifest.json" $HashFilePath = Join-Path $InstallerDir "SupermixStudioDesktopReleaseSHA256.txt" if (Test-Path $BundleManifestSource) { Copy-Item -Force $BundleManifestSource $BundleManifestTarget } +if (Test-Path $RuntimeManifestSource) { + Copy-Item -Force $RuntimeManifestSource $RuntimeManifestTarget +} $HashTargets = @() if (Test-Path $SetupPath) { $HashTargets += Get-Item $SetupPath } if (Test-Path $ExePath) { $HashTargets += Get-Item $ExePath } if (Test-Path $BundleManifestTarget) { $HashTargets += Get-Item $BundleManifestTarget } +if (Test-Path $RuntimeManifestTarget) { $HashTargets += Get-Item $RuntimeManifestTarget } if ($HashTargets.Count -gt 0) { $HashLines = foreach ($Item in $HashTargets) { diff --git a/source/generate_studio_runtime_manifest.py b/source/generate_studio_runtime_manifest.py new file mode 100644 index 000000000..cbe106593 --- /dev/null +++ b/source/generate_studio_runtime_manifest.py @@ -0,0 +1,239 @@ +"""Generate the deterministic Supermix Studio runtime distribution contract.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, Iterable, Optional, Sequence + + +MANIFEST_SCHEMA_VERSION = "supermix-studio-runtime-manifest-v1" +STUDIO_APP_VERSION = "2026.07.18" +DEFAULT_MANIFEST_PATH = Path("source/studio_runtime_manifest.json") + +RUNTIME_MODULES = ( + "source/route_policy_ledger.py", + "source/route_policy_lab.py", + "source/route_policy_explorer.py", + "source/route_policy_protocol.py", + "source/route_policy_shadow_registry.py", + "source/route_policy_study_cli.py", + "source/route_policy_protocol_cli.py", + "source/route_policy_shadow_cli.py", + "source/multimodel_runtime.py", + "source/supermix_multimodel_web_app.py", + "source/supermix_multimodel_desktop_app.py", +) + +CONTRACT_CONSTANTS = { + "source/route_policy_ledger.py": ( + "LEDGER_SCHEMA_VERSION", + "SUPPORT_SCHEMA_VERSION", + "EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION", + "OUTCOME_CONTRACT_SCHEMA_VERSION", + "OUTCOME_MATURITY_SCHEMA_VERSION", + ), + "source/route_policy_lab.py": ("READINESS_SCHEMA_VERSION",), + "source/route_policy_explorer.py": ( + "STUDY_PLAN_SCHEMA_VERSION", + "STUDY_ASSIGNMENT_SCHEMA_VERSION", + "REHEARSAL_RECEIPT_SCHEMA_VERSION", + "REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION", + "STUDY_ID", + "STUDY_VERSION", + ), + "source/route_policy_protocol.py": ( + "PROTOCOL_BUILD_INPUT_SCHEMA_VERSION", + "PROTOCOL_SCHEMA_VERSION", + "PROTOCOL_VERSION", + "REVIEW_BUNDLE_SCHEMA_VERSION", + "REVIEW_BUNDLE_VERSION", + "TARGET_POLICY_CLASS_SCHEMA_VERSION", + "POPULATION_SCHEMA_VERSION", + "STATEFUL_DESIGN_SCHEMA_VERSION", + "STOPPING_SCHEMA_VERSION", + "RANDOMNESS_SCHEMA_VERSION", + ), + "source/route_policy_shadow_registry.py": ( + "SHADOW_REGISTRY_SCHEMA_VERSION", + "SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION", + "SHADOW_DESIGN_BINDING_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_MANIFEST_SCHEMA_VERSION", + "SHADOW_CAMPAIGN_SEAL_SCHEMA_VERSION", + "SHADOW_SEED_CAPSULE_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION", + "SHADOW_CAMPAIGN_CLOSURE_SCHEMA_VERSION", + "SHADOW_SEED_REVEAL_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_REVEAL_SCHEMA_VERSION", + "SHADOW_REGISTRY_EVENT_SCHEMA_VERSION", + "SHADOW_REGISTRY_SNAPSHOT_SCHEMA_VERSION", + "SHADOW_ASSIGNMENT_ALGORITHM", + "SHADOW_CANONICALIZATION", + "SHADOW_LEGACY_BUNDLE_ENCODING", + "SHADOW_SCHEMA_OBJECTS_SHA256", + "SHADOW_TOTAL_ALLOCATION_BPS", + "SHADOW_CANDIDATE_ALLOCATION_BPS", + "SHADOW_BLOCK_ID", + ), +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _literal_constants(path: Path, names: Iterable[str]) -> Dict[str, Any]: + wanted = set(names) + found: Dict[str, Any] = {} + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in tree.body: + if not isinstance(node, (ast.Assign, ast.AnnAssign)): + continue + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value_node = node.value + for target in targets: + if isinstance(target, ast.Name) and target.id in wanted: + try: + found[target.id] = ast.literal_eval(value_node) + except (TypeError, ValueError): + pass + missing = sorted(wanted - set(found)) + if missing: + raise ValueError(f"{path.as_posix()} is missing literal contract constants: {', '.join(missing)}") + return {name: found[name] for name in sorted(found)} + + +def _git_provenance(repo_root: Path) -> Dict[str, Any]: + def run(*args: str) -> str: + try: + return subprocess.check_output( + ["git", *args], + cwd=repo_root, + text=True, + stderr=subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "" + + commit = run("rev-parse", "HEAD") or None + branch = run("branch", "--show-current") or None + dirty = bool(run("status", "--porcelain")) + return {"commit": commit, "branch": branch, "dirty": dirty} + + +def build_manifest(repo_root: Path, *, include_git: bool = False) -> Dict[str, Any]: + root = repo_root.resolve() + modules = [] + for relative in RUNTIME_MODULES: + path = root / relative + if not path.is_file(): + raise FileNotFoundError(f"required Studio runtime module is missing: {relative}") + modules.append( + { + "path": relative, + "sha256": _sha256(path), + "size_bytes": path.stat().st_size, + } + ) + + contracts: Dict[str, Any] = {} + for relative, names in CONTRACT_CONSTANTS.items(): + contracts[relative] = _literal_constants(root / relative, names) + + manifest: Dict[str, Any] = { + "schema_version": MANIFEST_SCHEMA_VERSION, + "app_version": STUDIO_APP_VERSION, + "active_runtime_tree": "source", + "legacy_compatibility_tree": "runtime_python", + "entrypoints": { + "desktop": "source/supermix_multimodel_desktop_app.py", + "web": "source/supermix_multimodel_web_app.py", + "route_study_console": "source/route_policy_protocol_cli.py", + "route_shadow_console": "source/route_policy_shadow_cli.py", + }, + "expected_windows_artifacts": { + "desktop_executable": "SupermixStudioDesktop.exe", + "route_study_console": "SupermixRouteStudy.exe", + "route_shadow_console": "SupermixRouteShadow.exe", + "installer": "SupermixStudioDesktopSetup.exe", + }, + "modules": modules, + "contracts": contracts, + "package_guards": { + "route_protocol_activation_available": False, + "route_protocol_assignment_implementation_available": False, + "route_rehearsal_writes_ledger": False, + "route_ledger_requires_executed_assignment_namespace": True, + "route_review_bundle_full_source_reconstruction": True, + "route_review_bundle_authenticity_proof_available": False, + "route_review_bundle_trusted_timestamp_available": False, + "route_shadow_registry_available": True, + "route_shadow_assignment_executes_routes": False, + "route_shadow_private_seed_persisted_before_reveal": False, + "route_shadow_authenticity_proof_available": False, + "route_shadow_trusted_timestamp_available": False, + "automatic_policy_promotion_available": False, + }, + } + if include_git: + manifest["release_provenance"] = _git_provenance(root) + return manifest + + +def _render(manifest: Dict[str, Any]) -> str: + return json.dumps(manifest, sort_keys=True, indent=2, ensure_ascii=False, allow_nan=False) + "\n" + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo-root", default="", help="repository root; defaults to this script's parent") + parser.add_argument("--output", default=str(DEFAULT_MANIFEST_PATH), help="manifest path relative to repo root") + parser.add_argument("--check", action="store_true", help="fail if the checked manifest is stale") + parser.add_argument( + "--release-provenance", + action="store_true", + help="include commit, branch, and dirty state for a build-only release manifest", + ) + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + repo_root = ( + Path(args.repo_root).expanduser().resolve() + if str(args.repo_root).strip() + else Path(__file__).resolve().parents[1] + ) + output = Path(args.output) + if not output.is_absolute(): + output = repo_root / output + try: + rendered = _render(build_manifest(repo_root, include_git=bool(args.release_provenance))) + if args.check: + current = output.read_text(encoding="utf-8") + if current != rendered: + print( + f"Studio runtime manifest is stale: {output}. Regenerate it before packaging.", + file=sys.stderr, + ) + return 1 + else: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(rendered, encoding="utf-8") + except (OSError, ValueError) as exc: + print(f"Studio runtime manifest error: {exc}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/source/studio_runtime_manifest.json b/source/studio_runtime_manifest.json new file mode 100644 index 000000000..191feb4e4 --- /dev/null +++ b/source/studio_runtime_manifest.json @@ -0,0 +1,143 @@ +{ + "active_runtime_tree": "source", + "app_version": "2026.07.18", + "contracts": { + "source/route_policy_explorer.py": { + "REHEARSAL_RECEIPT_SCHEMA_VERSION": "route-exploration-rehearsal-receipt-v1", + "REHEARSAL_SUPPORT_PROPOSAL_SCHEMA_VERSION": "route-exploration-rehearsal-support-proposal-v1", + "STUDY_ASSIGNMENT_SCHEMA_VERSION": "route-exploration-rehearsal-assignment-v1", + "STUDY_ID": "auto-route-adjacent-explorer-v1", + "STUDY_PLAN_SCHEMA_VERSION": "route-exploration-plan-v1", + "STUDY_VERSION": "1.0.0" + }, + "source/route_policy_lab.py": { + "READINESS_SCHEMA_VERSION": "route-readiness-v2" + }, + "source/route_policy_ledger.py": { + "EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION": "route-execution-assignment-v1", + "LEDGER_SCHEMA_VERSION": 3, + "OUTCOME_CONTRACT_SCHEMA_VERSION": "route-outcome-contract-v1", + "OUTCOME_MATURITY_SCHEMA_VERSION": "route-outcome-maturity-v1", + "SUPPORT_SCHEMA_VERSION": "route-support-v1" + }, + "source/route_policy_protocol.py": { + "POPULATION_SCHEMA_VERSION": "route-study-population-v1", + "PROTOCOL_BUILD_INPUT_SCHEMA_VERSION": "route-study-protocol-build-input-v1", + "PROTOCOL_SCHEMA_VERSION": "route-study-protocol-preflight-v1", + "PROTOCOL_VERSION": "1.0.0", + "RANDOMNESS_SCHEMA_VERSION": "route-study-randomness-v1", + "REVIEW_BUNDLE_SCHEMA_VERSION": "route-study-review-bundle-v1", + "REVIEW_BUNDLE_VERSION": "1.0.0", + "STATEFUL_DESIGN_SCHEMA_VERSION": "route-stateful-design-screen-v1", + "STOPPING_SCHEMA_VERSION": "route-study-stopping-v1", + "TARGET_POLICY_CLASS_SCHEMA_VERSION": "route-target-policy-class-v1" + }, + "source/route_policy_shadow_registry.py": { + "SHADOW_ASSIGNMENT_ALGORITHM": "hkdf-sha256-hmac-sha256-whole-policy-bps-v1", + "SHADOW_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION": "route-study-shadow-assignment-commitment-v1", + "SHADOW_ASSIGNMENT_MANIFEST_SCHEMA_VERSION": "route-study-shadow-assignment-manifest-v1", + "SHADOW_ASSIGNMENT_REVEAL_SCHEMA_VERSION": "route-study-shadow-assignment-reveal-v1", + "SHADOW_BLOCK_ID": 0, + "SHADOW_CAMPAIGN_CLOSURE_SCHEMA_VERSION": "route-study-shadow-campaign-closure-v1", + "SHADOW_CAMPAIGN_SEAL_SCHEMA_VERSION": "route-study-shadow-campaign-seal-v1", + "SHADOW_CANDIDATE_ALLOCATION_BPS": 5000, + "SHADOW_CANONICALIZATION": "rfc8785-jcs-restricted-ijson-integer-v1", + "SHADOW_DESIGN_BINDING_SCHEMA_VERSION": "route-study-shadow-design-binding-v1", + "SHADOW_LEGACY_BUNDLE_ENCODING": "route-review-bundle-v1-python-sorted-compact-utf8", + "SHADOW_PUBLIC_PACKAGE_SCHEMA_VERSION": "route-study-shadow-public-package-v1", + "SHADOW_REGISTRY_EVENT_SCHEMA_VERSION": "route-study-shadow-registry-event-v1", + "SHADOW_REGISTRY_SCHEMA_VERSION": 1, + "SHADOW_REGISTRY_SNAPSHOT_SCHEMA_VERSION": "route-study-shadow-registry-snapshot-v1", + "SHADOW_SCHEMA_OBJECTS_SHA256": "fe12a50841e5f983adce9ef1eb4ebf6635ab079f8c0043f4ca68500b38f9a04c", + "SHADOW_SEED_CAPSULE_SCHEMA_VERSION": "route-study-shadow-seed-capsule-v1", + "SHADOW_SEED_REVEAL_SCHEMA_VERSION": "route-study-shadow-seed-reveal-v1", + "SHADOW_TOTAL_ALLOCATION_BPS": 10000 + } + }, + "entrypoints": { + "desktop": "source/supermix_multimodel_desktop_app.py", + "route_shadow_console": "source/route_policy_shadow_cli.py", + "route_study_console": "source/route_policy_protocol_cli.py", + "web": "source/supermix_multimodel_web_app.py" + }, + "expected_windows_artifacts": { + "desktop_executable": "SupermixStudioDesktop.exe", + "installer": "SupermixStudioDesktopSetup.exe", + "route_shadow_console": "SupermixRouteShadow.exe", + "route_study_console": "SupermixRouteStudy.exe" + }, + "legacy_compatibility_tree": "runtime_python", + "modules": [ + { + "path": "source/route_policy_ledger.py", + "sha256": "d7fe4b95922448484fde62c5f0889ea9cd6d292871aa8d22de4355274492dcf9", + "size_bytes": 111375 + }, + { + "path": "source/route_policy_lab.py", + "sha256": "2e888ed9b400de692ed8589e44b86b71e70ebfd7e2a50a552eb91f4192fb5632", + "size_bytes": 57994 + }, + { + "path": "source/route_policy_explorer.py", + "sha256": "4d0790b104ffb94017f7f6fb9f2dbbd9c0757fed89059b4794267907e81430e5", + "size_bytes": 62438 + }, + { + "path": "source/route_policy_protocol.py", + "sha256": "8fc25b420e549ba5755e75e4de630d19102f629f80b774e481916641f45654ff", + "size_bytes": 61648 + }, + { + "path": "source/route_policy_shadow_registry.py", + "sha256": "6208e89a1495417c9f22b01aef6a5328fa92746e6a44e09d9eafa997bdb3e670", + "size_bytes": 103163 + }, + { + "path": "source/route_policy_study_cli.py", + "sha256": "7de1204145c837a1a445a71801ade37f5feee4a359bf72b6bd5815ace45fb3c0", + "size_bytes": 5973 + }, + { + "path": "source/route_policy_protocol_cli.py", + "sha256": "1bfcfd81d5b17235ec2f08607ee9d9edf172a03ff366548beca038d88ee4698e", + "size_bytes": 5925 + }, + { + "path": "source/route_policy_shadow_cli.py", + "sha256": "48f9157e28b66553a5e6857f97f51d9d8abba3fe20ef01c650d0a45e5607fc0a", + "size_bytes": 13330 + }, + { + "path": "source/multimodel_runtime.py", + "sha256": "0f1c179d5b5b9541b37f048bf0b937571d3b88fdd2aa4364d078f8a4578c51b9", + "size_bytes": 266501 + }, + { + "path": "source/supermix_multimodel_web_app.py", + "sha256": "7883be70e6547292c1d86f56441f9a4732682e96d848c948e7ece634416f2e5a", + "size_bytes": 160699 + }, + { + "path": "source/supermix_multimodel_desktop_app.py", + "sha256": "f3e679df20096ee170ec9fcacaaf65119c57137c52bba6a8fb0700c3823e1be9", + "size_bytes": 11632 + } + ], + "package_guards": { + "automatic_policy_promotion_available": false, + "route_ledger_requires_executed_assignment_namespace": true, + "route_protocol_activation_available": false, + "route_protocol_assignment_implementation_available": false, + "route_rehearsal_writes_ledger": false, + "route_review_bundle_authenticity_proof_available": false, + "route_review_bundle_full_source_reconstruction": true, + "route_review_bundle_trusted_timestamp_available": false, + "route_shadow_assignment_executes_routes": false, + "route_shadow_authenticity_proof_available": false, + "route_shadow_private_seed_persisted_before_reveal": false, + "route_shadow_registry_available": true, + "route_shadow_trusted_timestamp_available": false + }, + "schema_version": "supermix-studio-runtime-manifest-v1" +} diff --git a/test_studio_runtime_manifest.py b/test_studio_runtime_manifest.py new file mode 100644 index 000000000..56dfa5539 --- /dev/null +++ b/test_studio_runtime_manifest.py @@ -0,0 +1,102 @@ +import json +from pathlib import Path + +from source.generate_studio_runtime_manifest import ( + MANIFEST_SCHEMA_VERSION, + RUNTIME_MODULES, + STUDIO_APP_VERSION, + build_manifest, + main, +) + + +ROOT = Path(__file__).resolve().parent + + +def test_checked_runtime_manifest_is_current_and_complete(): + assert main(["--repo-root", str(ROOT), "--check"]) == 0 + checked = json.loads((ROOT / "source/studio_runtime_manifest.json").read_text(encoding="utf-8")) + generated = build_manifest(ROOT) + assert checked == generated + assert checked["schema_version"] == MANIFEST_SCHEMA_VERSION + assert checked["app_version"] == STUDIO_APP_VERSION + assert [row["path"] for row in checked["modules"]] == list(RUNTIME_MODULES) + assert all(len(row["sha256"]) == 64 and row["size_bytes"] > 0 for row in checked["modules"]) + assert checked["entrypoints"]["route_study_console"] == "source/route_policy_protocol_cli.py" + assert checked["expected_windows_artifacts"]["route_study_console"] == "SupermixRouteStudy.exe" + assert checked["entrypoints"]["route_shadow_console"] == "source/route_policy_shadow_cli.py" + assert checked["expected_windows_artifacts"]["route_shadow_console"] == "SupermixRouteShadow.exe" + assert checked["package_guards"] == { + "route_protocol_activation_available": False, + "route_protocol_assignment_implementation_available": False, + "route_rehearsal_writes_ledger": False, + "route_ledger_requires_executed_assignment_namespace": True, + "route_review_bundle_full_source_reconstruction": True, + "route_review_bundle_authenticity_proof_available": False, + "route_review_bundle_trusted_timestamp_available": False, + "route_shadow_registry_available": True, + "route_shadow_assignment_executes_routes": False, + "route_shadow_private_seed_persisted_before_reveal": False, + "route_shadow_authenticity_proof_available": False, + "route_shadow_trusted_timestamp_available": False, + "automatic_policy_promotion_available": False, + } + + +def test_manifest_exposes_route_contract_versions_without_importing_runtime(): + manifest = build_manifest(ROOT) + contracts = manifest["contracts"] + assert contracts["source/route_policy_ledger.py"]["OUTCOME_CONTRACT_SCHEMA_VERSION"] == ( + "route-outcome-contract-v1" + ) + assert contracts["source/route_policy_ledger.py"][ + "EXECUTED_ASSIGNMENT_COMMITMENT_SCHEMA_VERSION" + ] == "route-execution-assignment-v1" + assert contracts["source/route_policy_explorer.py"]["STUDY_PLAN_SCHEMA_VERSION"] == ( + "route-exploration-plan-v1" + ) + assert contracts["source/route_policy_protocol.py"]["PROTOCOL_SCHEMA_VERSION"] == ( + "route-study-protocol-preflight-v1" + ) + assert contracts["source/route_policy_protocol.py"]["REVIEW_BUNDLE_SCHEMA_VERSION"] == ( + "route-study-review-bundle-v1" + ) + assert contracts["source/route_policy_protocol.py"][ + "PROTOCOL_BUILD_INPUT_SCHEMA_VERSION" + ] == "route-study-protocol-build-input-v1" + assert contracts["source/route_policy_shadow_registry.py"][ + "SHADOW_REGISTRY_SCHEMA_VERSION" + ] == 1 + assert contracts["source/route_policy_shadow_registry.py"][ + "SHADOW_ASSIGNMENT_ALGORITHM" + ] == "hkdf-sha256-hmac-sha256-whole-policy-bps-v1" + assert contracts["source/route_policy_shadow_registry.py"][ + "SHADOW_CANONICALIZATION" + ] == "rfc8785-jcs-restricted-ijson-integer-v1" + assert len( + contracts["source/route_policy_shadow_registry.py"][ + "SHADOW_SCHEMA_OBJECTS_SHA256" + ] + ) == 64 + + +def test_studio_packaging_sources_are_machine_portable_and_include_console_contract(): + spec = (ROOT / "SupermixStudioDesktop.spec").read_text(encoding="utf-8") + build = (ROOT / "source/build_supermix_studio_desktop_exe.ps1").read_text(encoding="utf-8") + installer = (ROOT / "installer/SupermixStudioDesktop.iss").read_text(encoding="utf-8") + + assert "Supermix_27" not in spec + assert "C:\\Users\\" not in spec + assert "Supermix_27" not in build + assert "C:\\Users\\" not in build + assert "studio_runtime_manifest.json" in spec + assert "generate_studio_runtime_manifest.py" in build + assert "--check" in build + assert "SupermixRouteStudy" in build + assert "route_policy_protocol_cli.py" in build + assert "SupermixRouteShadow" in build + assert "route_policy_shadow_cli.py" in build + assert "studio_runtime_manifest.json" in build + assert f'#define MyAppVersion "{STUDIO_APP_VERSION}"' in installer + assert "DefaultDirName={localappdata}\\Programs\\Supermix Studio" in installer + assert "PrivilegesRequired=lowest" in installer