open-standards/wai/wai-rs/tests/semantic_path.rs

//! The perceptual / generative-semantic path, tested WITHOUT any heavy model.
//!
//! The wire format, the conformance metric, and the verify loop are pure Rust
//! and fully exercised here with mock generator + embedder. The only part that
//! needs a real CLIP ONNX is `OnnxClip::embed`, which is covered by its own
//! (skipped-if-absent) test and not required for these.

#![cfg(feature = "semantic")]

use std::cell::Cell;

use wai::container::{Conditioning, Manifest, ModelRequirement, Wai};
use wai::semantic::{
    cosine_sim, regenerate_verified, SemanticEmbedder, Frame, SemanticError, SemanticGenerator,
    SemanticPackage,
};

fn pkg_from(emb: &[f32]) -> SemanticPackage {
    SemanticPackage::from_embedding(emb, b"thumb-jpeg".to_vec(), b"edges-png".to_vec(), "tag")
}

fn envelope(pkg: &SemanticPackage, min_sim: Option<f32>) -> Wai {
    envelope_m(pkg, min_sim, None)
}

fn envelope_m(pkg: &SemanticPackage, min_sim: Option<f32>, sim_metric: Option<String>) -> Wai {
    let mut requirement = ModelRequirement::new("wai.semantic.clip");
    requirement.min_sim = min_sim;
    requirement.sim_metric = sim_metric;
    let manifest = Manifest::new("image", "improve", requirement, Conditioning::new("semantic_package_v1"))
        .with_target(serde_json::json!({"w": 8, "h": 8}));
    Wai::new(manifest, pkg.pack())
}

/// A generator that returns a fixed 1×1 frame — the protocol doesn't care what
/// pixels come back, only how they score.
struct FixedGen;
impl SemanticGenerator for FixedGen {
    fn regenerate(&self, _pkg: &SemanticPackage, _w: u32, _h: u32) -> Result<Frame, SemanticError> {
        Ok(Frame { width: 1, height: 1, rgb: vec![0, 0, 0] })
    }
}

/// An embedder that always returns a fixed embedding, regardless of the frame.
struct Echo(Vec<f32>);
impl SemanticEmbedder for Echo {
    fn embed(&self, _f: &Frame) -> Result<Vec<f32>, SemanticError> {
        Ok(self.0.clone())
    }
}

#[test]
fn package_round_trips_byte_faithfully() {
    let pkg = pkg_from(&[1.0, 0.0, 0.0, 0.0]);
    let bytes = pkg.pack();
    let back = SemanticPackage::unpack(&bytes).expect("unpack");
    assert_eq!(pkg, back, "semantic package must survive pack/unpack");
    assert_eq!(back.thumb_jpeg, b"thumb-jpeg");
    assert_eq!(back.edges_png, b"edges-png");
    assert_eq!(back.tag, "tag");
    // The quantized clip vector dequantizes to the original direction.
    assert!(cosine_sim(&back.clip(), &[1.0, 0.0, 0.0, 0.0]) > 0.999);
}

#[test]
fn unpack_rejects_truncation_without_panic() {
    assert!(SemanticPackage::unpack(&[]).is_err());
    // length prefix says 5 bytes, only 2 follow
    assert!(SemanticPackage::unpack(&[5, 0, 0, 0, 1, 2]).is_err());
}

#[test]
fn clip_sim_is_a_cosine() {
    assert!((cosine_sim(&[1.0, 0.0], &[1.0, 0.0]) - 1.0).abs() < 1e-6); // identical
    assert!(cosine_sim(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-6); // orthogonal
    assert!((cosine_sim(&[1.0, 0.0], &[-1.0, 0.0]) + 1.0).abs() < 1e-6); // opposite
    // Scale-invariant: magnitude doesn't matter, direction does.
    assert!((cosine_sim(&[3.0, 0.0], &[9.0, 0.0]) - 1.0).abs() < 1e-6);
}

#[test]
fn conformance_passes_when_similar() {
    let pkg = pkg_from(&[1.0, 0.0, 0.0, 0.0]);
    let env = envelope(&pkg, Some(0.9));
    let res = regenerate_verified(&env, &FixedGen, &Echo(vec![1.0, 0.0, 0.0, 0.0]), 8, 8, 1)
        .expect("verify");
    assert!(res.conformed, "sim {} should clear 0.9", res.similarity);
    assert!(res.similarity > 0.99);
}

#[test]
fn conformance_fails_when_dissimilar() {
    let pkg = pkg_from(&[1.0, 0.0, 0.0, 0.0]);
    let env = envelope(&pkg, Some(0.5));
    // Embedder returns an orthogonal vector → similarity ~0, below threshold.
    let res = regenerate_verified(&env, &FixedGen, &Echo(vec![0.0, 1.0, 0.0, 0.0]), 8, 8, 1)
        .expect("verify");
    assert!(!res.conformed, "orthogonal output must not conform (sim {})", res.similarity);
}

/// An embedder whose output ramps from orthogonal toward the target across
/// calls, so best-of-K should strictly beat a single attempt.
struct RampEmbedder {
    target: Vec<f32>,
    ortho: Vec<f32>,
    calls: Cell<usize>,
    total: usize,
}
impl SemanticEmbedder for RampEmbedder {
    fn embed(&self, _f: &Frame) -> Result<Vec<f32>, SemanticError> {
        let i = self.calls.get();
        self.calls.set(i + 1);
        let t = (i as f32 + 1.0) / self.total as f32; // 1/total .. 1.0
        let v = self
            .target
            .iter()
            .zip(&self.ortho)
            .map(|(a, b)| t * a + (1.0 - t) * b)
            .collect();
        Ok(v)
    }
}

#[test]
fn best_of_k_beats_single_attempt() {
    let pkg = pkg_from(&[1.0, 0.0, 0.0, 0.0]);
    let env = envelope(&pkg, Some(0.0));

    let single = regenerate_verified(
        &env, &FixedGen,
        &RampEmbedder { target: vec![1.0, 0.0, 0.0, 0.0], ortho: vec![0.0, 1.0, 0.0, 0.0],
                        calls: Cell::new(0), total: 4 },
        8, 8, 1,
    ).unwrap();

    let best = regenerate_verified(
        &env, &FixedGen,
        &RampEmbedder { target: vec![1.0, 0.0, 0.0, 0.0], ortho: vec![0.0, 1.0, 0.0, 0.0],
                        calls: Cell::new(0), total: 4 },
        8, 8, 4,
    ).unwrap();

    assert!(best.similarity > single.similarity,
            "best-of-4 ({}) should beat best-of-1 ({})", best.similarity, single.similarity);
    assert!(best.similarity > 0.999, "last ramp step reaches the target");
}

#[test]
fn sim_metric_is_named_and_carried_through() {
    let pkg = pkg_from(&[1.0, 0.0, 0.0, 0.0]);
    let target = vec![1.0, 0.0, 0.0, 0.0];

    // Unset metric → "clip" by convention, and that shows up in the result.
    let def = envelope(&pkg, Some(0.0));
    assert_eq!(def.sim_metric(), "clip");
    let r = regenerate_verified(&def, &FixedGen, &Echo(target.clone()), 8, 8, 1).unwrap();
    assert_eq!(r.metric, "clip");

    // A declared metric (e.g. DreamSim space) is named and carried through —
    // the sink must score with a matching embedder for the result to mean
    // anything; the protocol just records which space was used.
    let ds = envelope_m(&pkg, Some(0.0), Some("dreamsim".into()));
    assert_eq!(ds.sim_metric(), "dreamsim");
    let r2 = regenerate_verified(&ds, &FixedGen, &Echo(target), 8, 8, 1).unwrap();
    assert_eq!(r2.metric, "dreamsim");
}