types.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.
 
from __future__ import annotations
from typing import Any, Literal, Optional, Union
from pydantic import BaseModel
 
 
# ─── CaptureBundle ────────────────────────────────────────────────────────────
 
class CaptureBundle(BaseModel):
    version: int = 1
    capture_id: str
    session_id: str
    model_id: str
    hf_name: str
    prompt: str
    response: str
    created_at: str
    top_features: list[Any]
    sae_layer: int
    logit_lens: list[Any]
    attention: dict[str, Any]
 
 
# ─── Refs + summaries ────────────────────────────────────────────────────────
 
class CaptureRef(BaseModel):
    r2_key: str
    size_bytes: int
 
 
class ArtifactRef(BaseModel):
    r2_key: str
    artifact_type: str
 
 
class CaptureSummary(BaseModel):
    model_id: str
    prompt: str
    top_feature_count: int
 
 
# ─── EngineRegistration ───────────────────────────────────────────────────────
 
class EngineRegistration(BaseModel):
    id: str
    user_id: str
    device: str
    status: Literal["online", "offline"]
    last_heartbeat: str
    default_model_hf: Optional[str]
    gpu_info: Optional[str]
 
 
# ─── SessionState ─────────────────────────────────────────────────────────────
 
class SessionState(BaseModel):
    messages: list[dict[str, Any]]
    subAgentPanels: list[dict[str, Any]]
    captures: dict[str, CaptureRef]
    artifacts: dict[str, ArtifactRef]
    activeModelId: Optional[str]
 
 
# ─── SyncEvent — everything engine pushes to web ─────────────────────────────
 
class SyncEventMessageAppend(BaseModel):
    type: Literal["message.append"]
    message: dict[str, Any]
 
class SyncEventMessagePatch(BaseModel):
    type: Literal["message.patch"]
    message_id: str
    patch: dict[str, Any]
 
class SyncEventPanelUpsert(BaseModel):
    type: Literal["panel.upsert"]
    panel: dict[str, Any]
 
class SyncEventPanelPatch(BaseModel):
    type: Literal["panel.patch"]
    panel_id: str
    patch: dict[str, Any]
 
class SyncEventCaptureReady(BaseModel):
    type: Literal["capture.ready"]
    capture_id: str
    ref: CaptureRef
    summary: CaptureSummary
 
class SyncEventArtifactReady(BaseModel):
    type: Literal["artifact.ready"]
    artifact_id: str
    ref: ArtifactRef
 
class SyncEventToolStart(BaseModel):
    type: Literal["tool.start"]
    tool_call_id: str
    tool_name: str
    args: Any
 
class SyncEventToolResult(BaseModel):
    type: Literal["tool.result"]
    tool_call_id: str
    result: Any
    card: Optional[dict[str, Any]] = None
 
class SyncEventSessionMeta(BaseModel):
    type: Literal["session.meta"]
    patch: dict[str, Any]
 
class SyncEventEngineHeartbeat(BaseModel):
    type: Literal["engine.heartbeat"]
    device: str
    gpu_info: Optional[str] = None
 
class SyncEventEngineDisconnect(BaseModel):
    type: Literal["engine.disconnect"]
 
class SyncEventEngineLog(BaseModel):
    type: Literal["engine.log"]
    line: str
 
class SyncEventCommandAck(BaseModel):
    type: Literal["command.ack"]
    command_id: str
    status: Literal["accepted", "rejected", "done"]
    error: Optional[str] = None
 
 
SyncEvent = Union[
    SyncEventMessageAppend,
    SyncEventMessagePatch,
    SyncEventPanelUpsert,
    SyncEventPanelPatch,
    SyncEventCaptureReady,
    SyncEventArtifactReady,
    SyncEventToolStart,
    SyncEventToolResult,
    SyncEventSessionMeta,
    SyncEventEngineHeartbeat,
    SyncEventEngineDisconnect,
    SyncEventEngineLog,
    SyncEventCommandAck,
]
 
def parse_sync_event(data: dict[str, Any]) -> SyncEvent:
    from pydantic import TypeAdapter
    return TypeAdapter(SyncEvent).validate_python(data)
 
 
# ─── WebCommand — every UI action web sends to engine ────────────────────────
 
class CmdBase(BaseModel):
    command_id: str
 
# core session
class CmdChat(CmdBase):
    type: Literal["command.chat"]
    text: str
    hidden: Optional[bool] = None
 
class CmdApprove(CmdBase):
    type: Literal["command.approve"]
    plan_id: str
 
class CmdAbort(CmdBase):
    type: Literal["command.abort"]
    job_id: Optional[str] = None
 
class CmdSetModel(CmdBase):
    type: Literal["command.set_model"]
    model_id: str
 
class CmdSetDataset(CmdBase):
    type: Literal["command.set_dataset"]
    rows: list[Any]
 
class CmdUI(CmdBase):
    type: Literal["command.ui"]
    action: str
    payload: Any
 
# inspection
class CmdInspect(CmdBase):
    type: Literal["command.inspect"]
    prompt: str
    model_id: Optional[str] = None
 
class CmdRunFullInspection(CmdBase):
    type: Literal["command.run_full_inspection"]
    prompt: Optional[str] = None
 
class CmdRunBenchmarksOnTopFeature(CmdBase):
    type: Literal["command.run_benchmarks_on_top_feature"]
    feature_idx: int
 
class CmdRunAudit(CmdBase):
    type: Literal["command.run_audit"]
 
class CmdRunConsistencyEval(CmdBase):
    type: Literal["command.run_consistency_eval"]
    query: str
    templates: list[str]
 
class CmdRunSuppressionEval(CmdBase):
    type: Literal["command.run_suppression_eval"]
    topics: dict[str, list[str]]
 
class CmdRunBoundaryEval(CmdBase):
    type: Literal["command.run_boundary_eval"]
    prompts: list[str]
 
class CmdGetFeatureLogits(CmdBase):
    type: Literal["command.get_feature_logits"]
    feature_idx: int
    top_k: Optional[int] = None
 
class CmdGetFeatureNeighbors(CmdBase):
    type: Literal["command.get_feature_neighbors"]
    feature_idx: int
    top_k: Optional[int] = None
 
class CmdRunLayerAnalysis(CmdBase):
    type: Literal["command.run_layer_analysis"]
    prompts: Optional[list[str]] = None
    top_k: Optional[int] = None
    in_domain_prompts: Optional[list[str]] = None
    ood_prompts: Optional[list[str]] = None
 
class CmdRunPerturbationSensitivity(CmdBase):
    type: Literal["command.run_perturbation_sensitivity"]
    prompt: str
    n_channels: Optional[int] = None
    method: Optional[Literal["dropout", "gaussian"]] = None
 
class CmdRunAttentionRouting(CmdBase):
    type: Literal["command.run_attention_routing"]
    prompt: str
    top_k: Optional[int] = None
 
# steering
class CmdSteer(CmdBase):
    type: Literal["command.steer"]
    capture_id: str
    feature_idx: int
    scale: float
 
class CmdRunSteerAndShow(CmdBase):
    type: Literal["command.run_steer_and_show"]
    prompt: str
    feature_idx: int
    feature_label: str
    strength: Optional[float] = None
    max_new_tokens: Optional[int] = None
 
class SteerFeature(BaseModel):
    feature_idx: int
    strength: float
    label: str
 
class CmdRunMultiSteer(CmdBase):
    type: Literal["command.run_multi_steer"]
    prompt: str
    features: list[SteerFeature]
    max_new_tokens: Optional[int] = None
 
# UMAP
class CmdEnsureUmapLoaded(CmdBase):
    type: Literal["command.ensure_umap_loaded"]
 
# dataset (LLM)
class CmdDatasetGenerate(CmdBase):
    type: Literal["command.dataset_generate"]
    topic: str
    count: Optional[int] = None
    format: Optional[str] = None
 
# simulation
class CmdRunSimulation(CmdBase):
    type: Literal["command.run_simulation"]
    topic: Optional[str] = None
    dataset: Optional[str] = None
    algo: Optional[str] = None
    count: Optional[int] = None
    rank: Optional[int] = None
    alpha: Optional[float] = None
    lr: Optional[float] = None
    epochs: Optional[int] = None
    dropout: Optional[float] = None
    targetModules: Optional[list[str]] = None
    warmupSteps: Optional[int] = None
    gradClip: Optional[float] = None
    weightDecay: Optional[float] = None
    gradAccumSteps: Optional[int] = None
    optimizer: Optional[Literal["adamw", "sgd", "lion"]] = None
    scheduler: Optional[Literal["cosine", "linear", "constant"]] = None
    maxSeqLen: Optional[int] = None
    useQlora: Optional[bool] = None
    use_rlhf: Optional[bool] = None
    rlhf_beta: Optional[float] = None
    rows: Optional[list[dict]] = None
 
class CmdListSimulationRuns(CmdBase):
    type: Literal["command.list_simulation_runs"]
 
class CmdLoadSimulationRun(CmdBase):
    type: Literal["command.load_simulation_run"]
    run_id: str
 
class CmdCompareSimulations(CmdBase):
    type: Literal["command.compare_simulations"]
    run_id_a: str
    run_id_b: str
    label_a: Optional[str] = None
    label_b: Optional[str] = None
 
# interactive cards
class CmdRunRedTeam(CmdBase):
    type: Literal["command.run_red_team"]
    vectors: Optional[list[str]] = None
 
class CmdCheckWeights(CmdBase):
    type: Literal["command.check_weights"]
    collapse_threshold: Optional[float] = None
 
class CmdRunCustomEval(CmdBase):
    type: Literal["command.run_custom_eval"]
    name: str
    prompts: list[str]
    reference_answers: list[str]
    description: Optional[str] = None
    threshold: Optional[float] = None
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
 
# session memory
class CmdWriteSessionMemory(CmdBase):
    type: Literal["command.write_session_memory"]
    key: str
    value: Any
 
class CmdReadSessionMemory(CmdBase):
    type: Literal["command.read_session_memory"]
    key: str
 
 
WebCommand = Union[
    CmdChat, CmdApprove, CmdAbort, CmdSetModel, CmdSetDataset, CmdUI,
    CmdInspect, CmdRunFullInspection, CmdRunBenchmarksOnTopFeature, CmdRunAudit,
    CmdRunConsistencyEval, CmdRunSuppressionEval, CmdRunBoundaryEval,
    CmdGetFeatureLogits, CmdGetFeatureNeighbors,
    CmdRunLayerAnalysis, CmdRunPerturbationSensitivity, CmdRunAttentionRouting,
    CmdSteer, CmdRunSteerAndShow, CmdRunMultiSteer,
    CmdEnsureUmapLoaded,
    CmdDatasetGenerate,
    CmdRunSimulation,
    CmdListSimulationRuns, CmdLoadSimulationRun, CmdCompareSimulations,
    CmdRunRedTeam, CmdCheckWeights,
    CmdRunCustomEval,
    CmdWriteSessionMemory, CmdReadSessionMemory,
]
 
def parse_web_command(data: dict[str, Any]) -> WebCommand:
    from pydantic import TypeAdapter
    return TypeAdapter(WebCommand).validate_python(data)
 
 
__all__ = [
    "CaptureBundle",
    "CaptureRef", "ArtifactRef", "CaptureSummary",
    "EngineRegistration", "SessionState",
    "SyncEvent", "parse_sync_event",
    "WebCommand", "parse_web_command",
    "CmdChat", "CmdApprove", "CmdAbort", "CmdSetModel", "CmdSetDataset", "CmdUI",
    "CmdInspect", "CmdRunFullInspection", "CmdRunBenchmarksOnTopFeature", "CmdRunAudit",
    "CmdRunConsistencyEval", "CmdRunSuppressionEval", "CmdRunBoundaryEval",
    "CmdGetFeatureLogits", "CmdGetFeatureNeighbors",
    "CmdRunLayerAnalysis", "CmdRunPerturbationSensitivity", "CmdRunAttentionRouting",
    "CmdSteer", "CmdRunSteerAndShow", "CmdRunMultiSteer", "SteerFeature",
    "CmdEnsureUmapLoaded",
    "CmdDatasetGenerate",
    "CmdRunSimulation",
    "CmdListSimulationRuns", "CmdLoadSimulationRun", "CmdCompareSimulations",
    "CmdRunRedTeam", "CmdCheckWeights",
    "CmdRunCustomEval",
    "CmdWriteSessionMemory", "CmdReadSessionMemory",
]