sync_slim.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Trim large tool payloads before cloud sync."""
from __future__ import annotations
 
import random
 
UMAP_WEB_MAX_POINTS = 8192
 
 
def downsample_umap_points(points: list[dict], max_points: int = UMAP_WEB_MAX_POINTS) -> list[dict]:
    n = len(points)
    if n <= max_points:
        return points
    idx = sorted(random.Random(42).sample(range(n), max_points))
    return [points[i] for i in idx]
 
 
def slim_tool_result_for_sync(tool_name: str, result: dict, card: dict | None = None) -> dict:
    """Drop bulky fields from synced tool.result (full data lives on card only)."""
    if tool_name == "ensure_umap_loaded":
        card_points = []
        if card and isinstance(card.get("data"), dict):
            card_points = card["data"].get("points") or []
        return {
            "model_id": result.get("model_id"),
            "layer": result.get("layer"),
            "n_features": result.get("n_features"),
            "n_points": len(card_points) or result.get("n_points"),
        }
    if tool_name == "run_sae_diff" and card is not None:
        return {
            "baseModelId": result.get("baseModelId"),
            "ftCheckpointName": result.get("ftCheckpointName"),
            "layer": result.get("layer"),
            "nFeatures": result.get("nFeatures"),
            "nChanged": result.get("nChanged"),
            "meanAbsDelta": result.get("meanAbsDelta"),
            "maxAbsDelta": result.get("maxAbsDelta"),
        }
    if tool_name == "run_sae_align" and card is not None:
        return {
            "mean_cosine": result.get("mean_cosine"),
            "n_pairs": result.get("n_pairs"),
        }
    if tool_name == "run_sae_train" and card is not None:
        return {
            "model_id": result.get("model_id"),
            "layer": result.get("layer"),
            "quick": result.get("quick"),
            "name": result.get("name"),
            "output_path": result.get("output_path"),
            "status": result.get("status"),
        }
    if tool_name == "run_weight_diff" and card is not None:
        return {
            "baseModelId": result.get("baseModelId"),
            "ftCheckpointName": result.get("ftCheckpointName"),
            "modelMode": result.get("modelMode"),
            "deltaMode": result.get("deltaMode"),
            "nMatrices": result.get("nMatrices"),
            "totalDeltaL2": result.get("totalDeltaL2"),
            "maxDeltaL2": result.get("maxDeltaL2"),
        }
    if tool_name == "run_merge_analysis" and card is not None:
        return {
            "baseModelId": result.get("baseModelId"),
            "ftCheckpointName": result.get("ftCheckpointName"),
            "mergeVerdict": result.get("mergeVerdict"),
            "nWarnings": len(result.get("warnings") or []),
            "totalDeltaL2": result.get("totalDeltaL2"),
        }
    if tool_name == "run_trajectory_analysis" and card is not None:
        return {
            "baseModelId": result.get("baseModelId"),
            "nCheckpoints": result.get("nCheckpoints"),
            "nAnalyzed": result.get("nAnalyzed"),
            "peakStep": result.get("peakStep"),
            "peakTotalDeltaL2": result.get("peakTotalDeltaL2"),
        }
    if tool_name == "run_residual_drift" and card is not None:
        return {
            "baseModelId": result.get("baseModelId"),
            "ftCheckpointName": result.get("ftCheckpointName"),
            "modelMode": result.get("modelMode"),
            "nProbes": result.get("nProbes"),
            "meanDrift": result.get("meanDrift"),
            "maxDrift": result.get("maxDrift"),
            "peakLayer": result.get("peakLayer"),
        }
    if tool_name == "run_capture_activations" and card is not None:
        return {
            "model_id": result.get("model_id"),
            "model_mode": result.get("model_mode"),
            "n_probes": result.get("n_probes"),
            "layers": result.get("layers"),
            "position": result.get("position"),
            "encode_sae": result.get("encode_sae"),
            "manifest_path": result.get("manifest_path"),
            "status": result.get("status"),
        }
    if tool_name == "run_find_feature" and card is not None:
        return {
            "model_id": result.get("model_id"),
            "layer": result.get("layer"),
            "scorer": result.get("scorer"),
            "n_honest": result.get("n_honest"),
            "n_deceptive": result.get("n_deceptive"),
            "chosen_feature_idx": result.get("chosen_feature_idx"),
            "chosen_delta": result.get("chosen_delta"),
            "persisted_key": result.get("persisted_key"),
            "status": result.get("status"),
        }
    return result