open-standards/joule-code/examples/l3-jouleclaw/tests/common/mod.rs

//! Shared by the fabric-grade tests: run the L3 stage, with the fabric ON in this
//! test process and priced from a calibration file the test writes (never the
//! machine's own `~/.jouleclaw/fabric-calibration.json`), over a backend that
//! dispatches one real matmul through the fabric and reports the ledger's
//! projection, as the hybrid backend does.

use joule_code_cascade::{Request, Stage};
use joule_code_receipt::Energy;
use jouleclaw_fabric::EnergyProvenance;
use jouleclaw_llm::{FinishReason, LlmBackend, LlmError, LlmRequest, LlmResponse};
use jouleclaw_loader_gguf::fabric_dispatch::{ledger_reset, ledger_snapshot, matmul_bf16_auto, LedgerEntry};

/// One decode-shaped matmul (m = 1, a 4 KiB weight): the fabric places it on the
/// CPU SIMD member, `local/cpu-neon`.
struct OneMatmul;

impl LlmBackend for OneMatmul {
    fn model_name(&self) -> &str {
        "one-matmul"
    }
    fn complete(&self, _request: &LlmRequest) -> Result<LlmResponse, LlmError> {
        let (m, n, k) = (1usize, 32usize, 64usize);
        let a = vec![0.5f32; m * k];
        let w = vec![0u8; 2 * n * k];
        let mut out = vec![0f32; m * n];
        ledger_reset();
        matmul_bf16_auto(&a, &w, &[], m, n, k, &mut out);
        let projected_j: f64 = ledger_snapshot().iter().filter_map(|e| e.projected_uj).sum::<f64>() / 1e6;
        Ok(LlmResponse {
            text: "fn hello() {}".into(),
            finish_reason: FinishReason::Stop,
            input_tokens: 1,
            output_tokens: 1,
            energy_joules: (projected_j > 0.0).then_some(projected_j),
        })
    }
    fn typical_joules_per_call(&self) -> f64 {
        1.0
    }
}

/// Write a calibration file pricing `local/cpu-neon`'s matmul at `grade` (or
/// declaring none), switch the fabric on for this process, resolve one L3
/// request, and return its receipt energy with the ledger window it was priced
/// from. Call once per test process: the fabric is discovered once.
pub fn resolve_priced_at(tag: &str, grade: Option<EnergyProvenance>) -> (Energy, Vec<LedgerEntry>) {
    let path = std::path::Path::new(env!("CARGO_TARGET_TMPDIR")).join(format!("l3-fabric-grade-{tag}.json"));
    let declared = grade.map(|g| format!(r#","provenance":"{g:?}""#)).unwrap_or_default();
    // Priced high enough that the one 4.1e-6 GFLOP matmul projects to ~1 mJ, a
    // whole number of µJ on the receipt.
    let row = format!(
        r#"[{{"member_id":"local/cpu-neon","op":"matmul","uj_per_gflop":250000000.0,"method":"test row"{declared}}}]"#
    );
    std::fs::write(&path, row).expect("write the test calibration file");
    std::env::set_var("JOULECLAW_FABRIC_CALIBRATION", &path);
    std::env::set_var("JOULECLAW_FABRIC", "1");

    let stage = l3_jouleclaw_binding::jouleclaw_l3_stage(Box::new(OneMatmul));
    let res = stage
        .try_resolve(&Request { prior: "b3:AAAA".into(), intent: "write a hello function".into() })
        .expect("L3 resolved");
    let window = ledger_snapshot();
    let _ = std::fs::remove_file(&path);
    (res.receipt.energy, window)
}