src/warehouse/generator/mod.rs
//! Generative-LLM abstraction (EPIC #39) — ONE [`Generator`] trait, THREE real
//! implementations (`candle`, `mistralrs`, `onnx`) behind a [`generator`]
//! factory, plus a `mock` for tests and an off-by-default `ollama` HTTP client.
//!
//! # Why a trait
//! The H5 bake-off ([`super::agent_model_runs`]) pits a prompt against a set of
//! local models and records, per model, the answer + its token economics. Until
//! now the only real driver was [`OllamaCaller`](super::agent_model_runs) — an
//! HTTP shell-out to a running ollama daemon. EPIC #39 turns the model call into
//! a first-class interface with several in-process, pure-Rust (or ORT) backends
//! so nornir can generate without ollama. ollama is demoted to one OFF-by-default
//! convenience client; it is NOT one of the three backends.
//!
//! # The interface
//! ```ignore
//! pub trait Backend { fn id(&self) -> &str; fn available(&self) -> bool; }
//! pub trait Generator: Backend + Send + Sync {
//! fn complete(&self, req: &GenRequest) -> Result<GenAnswer>;
//! }
//! ```
//! [`GenRequest`] carries the prompt + sampling knobs; [`GenAnswer`] carries the
//! completion + the exact fields the bake-off's `ModelAnswer` needs, so a
//! [`GenAnswer`] slots straight into an `agent_model_runs` row (see
//! [`GenAnswer::into_model_answer`]).
//!
//! # The factory
//! [`generator(spec)`](generator) parses a `<backend>:<model>` spec and returns
//! a `Box<dyn Generator>`. Each backend arm is `#[cfg(feature = …)]`-gated, so a
//! lean default build only carries `mock` (+ `ollama` when `gen-ollama` is on);
//! an unrecognized backend, or one whose feature is off, yields a clear error.
//! Backends runtime-probe their lib/model in [`Backend::available`], so a
//! compiled-but-unconfigured backend constructs fine and simply reports
//! `available() == false` until its model is present.
//!
//! Selectable from the CLI via `--gen <spec>` or `$NORNIR_GEN_BACKEND`
//! ([`spec_from_env`]) — the CLI-parity LAW.
use anyhow::{Result, anyhow};
use super::agent_model_runs::{ModelAnswer, ModelCaller};
#[cfg(feature = "gen-candle")]
pub mod candle;
#[cfg(feature = "gen-mistralrs")]
pub mod mistralrs;
#[cfg(feature = "gen-ollama")]
pub mod ollama;
#[cfg(feature = "gen-onnx")]
pub mod onnx;
// The `claude` executor backend lives where the `viz::claude` client compiles.
#[cfg(any(feature = "viz", feature = "server"))]
pub mod claude;
/// The `$NORNIR_GEN_BACKEND` env var name (CLI-parity fallback for `--gen`).
pub const GEN_BACKEND_ENV: &str = "NORNIR_GEN_BACKEND";
/// One generation request: the prompt + sampling knobs. Mirrors the inputs every
/// backend needs; backends ignore knobs they don't support (and say so in their
/// docs) rather than erroring.
#[derive(Debug, Clone, PartialEq)]
pub struct GenRequest {
/// The user prompt to complete.
pub prompt: String,
/// Optional system prompt (role/instructions). Empty `None` ⇒ none.
pub system: Option<String>,
/// Hard cap on generated tokens.
pub max_tokens: usize,
/// Sampling temperature (`0.0` ⇒ greedy/deterministic).
pub temperature: f32,
/// Stop sequences: generation halts when any is produced.
pub stop: Vec<String>,
}
impl GenRequest {
/// A minimal request: just a prompt, greedy decode, 256-token cap, no stops.
pub fn new(prompt: impl Into<String>) -> Self {
GenRequest {
prompt: prompt.into(),
system: None,
max_tokens: 256,
temperature: 0.0,
stop: Vec::new(),
}
}
/// Set the system prompt. Chainable.
pub fn with_system(mut self, system: impl Into<String>) -> Self {
self.system = Some(system.into());
self
}
/// Set the token cap. Chainable.
pub fn with_max_tokens(mut self, n: usize) -> Self {
self.max_tokens = n;
self
}
/// Set the sampling temperature. Chainable.
pub fn with_temperature(mut self, t: f32) -> Self {
self.temperature = t;
self
}
/// Add a stop sequence. Chainable.
pub fn with_stop(mut self, s: impl Into<String>) -> Self {
self.stop.push(s.into());
self
}
}
/// One generation answer + its token economics — the fields the bake-off's
/// `ModelAnswer` (and thus an `agent_model_runs` row) needs, so a [`GenAnswer`]
/// slots straight in via [`into_model_answer`](GenAnswer::into_model_answer).
#[derive(Debug, Clone, PartialEq)]
pub struct GenAnswer {
/// The model's completion.
pub text: String,
/// Prompt token count.
pub tokens_in: i64,
/// Completion token count.
pub tokens_out: i64,
/// Decode throughput (completion tokens per second). `0.0` ⇒ the bake-off
/// derives it from `tokens_out / (latency_ms/1000)`.
pub tokens_per_s: f64,
/// Wall-clock duration of the call, milliseconds.
pub latency_ms: f64,
}
impl GenAnswer {
/// Convert to the bake-off's [`ModelAnswer`] (a local generation is free and
/// makes no MCP calls; the quality `score` is left `0.0` for a judge to fill).
pub fn into_model_answer(self) -> ModelAnswer {
ModelAnswer {
output: self.text,
latency_ms: self.latency_ms,
tokens_in: self.tokens_in,
tokens_out: self.tokens_out,
tokens_per_s: self.tokens_per_s,
score: 0.0,
cost_usd: 0.0,
mcp_tool_calls: 0,
}
}
}
/// Backend identity + a runtime availability probe. Split out from [`Generator`]
/// so a caller can list/probe backends without driving generation.
pub trait Backend {
/// The backend's spec id (e.g. `"candle:qwen2-0.5b"`, `"mock"`).
fn id(&self) -> &str;
/// Is this backend usable *right now*? Runtime-probes the underlying lib +
/// model presence — a compiled backend with no model downloaded reports
/// `false` instead of erroring on `complete`.
fn available(&self) -> bool;
}
/// The ONE generative-LLM interface. A [`Generator`] turns a [`GenRequest`] into
/// a [`GenAnswer`]. `Send + Sync` so the server/bake-off can share one across
/// threads.
pub trait Generator: Backend + Send + Sync {
/// Complete `req`, returning the answer + economics, or an error (a missing
/// model, a load failure, …). Callers that drive a bake-off translate an
/// error into a red `ok=false` row rather than aborting.
fn complete(&self, req: &GenRequest) -> Result<GenAnswer>;
}
/// Adapt any [`Generator`] to the bake-off's [`ModelCaller`] so the existing
/// `run_bakeoff` pipeline can drive the new backends unchanged. The `prompt`
/// becomes a default [`GenRequest`]; `agent`/`model` are recorded by the row
/// (the generator already knows which model it loaded).
pub struct GeneratorCaller<G: Generator> {
r#gen: G,
}
impl<G: Generator> GeneratorCaller<G> {
pub fn new(r#gen: G) -> Self {
Self { r#gen }
}
}
impl<G: Generator> ModelCaller for GeneratorCaller<G> {
fn call(&self, _agent: &str, _model: &str, prompt: &str) -> Result<ModelAnswer> {
let answer = self.r#gen.complete(&GenRequest::new(prompt))?;
Ok(answer.into_model_answer())
}
}
/// The boxed-trait-object form of [`GeneratorCaller`] — bridges a
/// `Box<dyn Generator>` (the factory's return type) to the bake-off's
/// [`ModelCaller`] so the CLI can drive any selected backend through the
/// existing `run_bakeoff_matrix` pipeline.
pub struct BoxGeneratorCaller {
r#gen: Box<dyn Generator>,
}
impl BoxGeneratorCaller {
pub fn new(r#gen: Box<dyn Generator>) -> Self {
Self { r#gen }
}
/// The wrapped generator's id (for CLI display).
pub fn id(&self) -> &str {
self.r#gen.id()
}
/// Probe the wrapped generator's availability (for CLI display).
pub fn available(&self) -> bool {
self.r#gen.available()
}
}
impl ModelCaller for BoxGeneratorCaller {
fn call(&self, _agent: &str, _model: &str, prompt: &str) -> Result<ModelAnswer> {
let answer = self.r#gen.complete(&GenRequest::new(prompt))?;
Ok(answer.into_model_answer())
}
}
/// Resolve the generator spec from the environment: `$NORNIR_GEN_BACKEND` if set,
/// else `None` (the CLI then falls back to its own default). The `--gen` flag
/// takes precedence over this when both are present.
pub fn spec_from_env() -> Option<String> {
std::env::var(GEN_BACKEND_ENV)
.ok()
.filter(|s| !s.trim().is_empty())
}
/// The generative-LLM factory: parse `spec` and build the matching backend.
///
/// Accepted forms:
/// - `candle:<model>` — `#[cfg(feature = "gen-candle")]`
/// - `mistralrs:<model>` — `#[cfg(feature = "gen-mistralrs")]`
/// - `onnx:<model>` — `#[cfg(feature = "gen-onnx")]`
/// - `ollama:<model>` — `#[cfg(feature = "gen-ollama")]` (NOT one of the 3)
/// - `mock` / `mock:<id>` — always available; round-trips a request to a canned
/// answer for tests + offline use.
///
/// A `<backend>:<model>` whose backend feature is off returns a clear
/// "not compiled" error (so the operator knows to rebuild with the feature),
/// and an unknown backend returns an "unknown backend" error. The returned
/// generator may still report `available() == false` if its model/lib is absent.
pub fn generator(spec: &str) -> Result<Box<dyn Generator>> {
let spec = spec.trim();
let (backend, model) = match spec.split_once(':') {
Some((b, m)) => (b.trim(), m.trim()),
None => (spec, ""),
};
match backend {
"mock" => Ok(Box::new(MockGenerator::new(if model.is_empty() {
"mock"
} else {
model
}))),
"candle" => {
#[cfg(feature = "gen-candle")]
{
Ok(Box::new(candle::CandleGenerator::new(model)?))
}
#[cfg(not(feature = "gen-candle"))]
{
Err(not_compiled("candle", "gen-candle"))
}
}
"mistralrs" => {
#[cfg(feature = "gen-mistralrs")]
{
Ok(Box::new(mistralrs::MistralRsGenerator::new(model)?))
}
#[cfg(not(feature = "gen-mistralrs"))]
{
Err(not_compiled("mistralrs", "gen-mistralrs"))
}
}
"onnx" => {
#[cfg(feature = "gen-onnx")]
{
Ok(Box::new(onnx::OnnxGenerator::new(model)?))
}
#[cfg(not(feature = "gen-onnx"))]
{
Err(not_compiled("onnx", "gen-onnx"))
}
}
"ollama" => {
#[cfg(feature = "gen-ollama")]
{
Ok(Box::new(ollama::OllamaGenerator::new(model, None)))
}
#[cfg(not(feature = "gen-ollama"))]
{
Err(not_compiled("ollama", "gen-ollama"))
}
}
"claude" => {
#[cfg(any(feature = "viz", feature = "server"))]
{
Ok(Box::new(claude::ClaudeGenerator::new(model)))
}
#[cfg(not(any(feature = "viz", feature = "server")))]
{
Err(not_compiled("claude", "viz"))
}
}
other => Err(anyhow!(
"unknown generator backend `{other}` in spec `{spec}` — \
expected one of candle:<m> | mistralrs:<m> | onnx:<m> | ollama:<m> | claude:<m> | mock"
)),
}
}
/// The error a factory arm returns when the named backend's feature is off.
#[allow(dead_code)] // unused only in a build where every gen feature is on
fn not_compiled(backend: &str, feature: &str) -> anyhow::Error {
anyhow!(
"generator backend `{backend}` is not compiled in — \
rebuild with `--features {feature}`"
)
}
// ─── mock backend ──────────────────────────────────────────────────────────
/// The `mock` backend: round-trips a [`GenRequest`] into a deterministic
/// [`GenAnswer`] with no LLM. It echoes the prompt (so a test can assert the
/// request reached the backend) and reports token counts derived from the
/// request, exercising the full request→answer path the real backends share.
///
/// `available()` is always `true` — the mock needs nothing external.
#[derive(Debug, Clone)]
pub struct MockGenerator {
id: String,
}
impl MockGenerator {
pub fn new(id: impl Into<String>) -> Self {
Self { id: id.into() }
}
/// The canned completion for `req`: a fixed echo so tests assert real I/O.
fn answer_text(req: &GenRequest) -> String {
let sys = req.system.as_deref().unwrap_or("");
format!("mock[{}]: {}{}", req.max_tokens, sys, req.prompt)
}
}
impl Backend for MockGenerator {
fn id(&self) -> &str {
&self.id
}
fn available(&self) -> bool {
true
}
}
impl Generator for MockGenerator {
fn complete(&self, req: &GenRequest) -> Result<GenAnswer> {
let text = Self::answer_text(req);
// Deterministic, request-derived economics: 1 token per whitespace word.
let tokens_in = req.prompt.split_whitespace().count() as i64;
let tokens_out = text.split_whitespace().count() as i64;
// A fixed synthetic latency so tokens_per_s is reproducible.
let latency_ms = 10.0;
let tokens_per_s = tokens_out as f64 / (latency_ms / 1000.0);
Ok(GenAnswer {
text,
tokens_in,
tokens_out,
tokens_per_s,
latency_ms,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mock_round_trips_request_into_answer() {
let r#gen = generator("mock").unwrap();
assert_eq!(r#gen.id(), "mock");
assert!(r#gen.available(), "mock is always available");
let req = GenRequest::new("capital of France?")
.with_system("be terse")
.with_max_tokens(64);
let ans = r#gen.complete(&req).unwrap();
// The answer carries the request back (real round-trip, not a no-op).
assert!(
ans.text.contains("capital of France?"),
"echoes prompt: {}",
ans.text
);
assert!(ans.text.contains("be terse"), "echoes system: {}", ans.text);
assert!(ans.text.contains("64"), "echoes max_tokens: {}", ans.text);
// Token economics are derived from the actual text, not zeroed.
assert_eq!(ans.tokens_in, 3, "3 prompt words");
assert!(ans.tokens_out > 0);
assert!(ans.tokens_per_s > 0.0);
assert!((ans.latency_ms - 10.0).abs() < 1e-9);
}
#[test]
fn mock_with_explicit_id_keeps_the_id() {
let r#gen = generator("mock:tiny").unwrap();
assert_eq!(r#gen.id(), "tiny");
}
#[test]
fn factory_errors_on_unknown_backend() {
let err = match generator("wat:model") {
Ok(_) => panic!("unknown backend must error"),
Err(e) => e.to_string(),
};
assert!(err.contains("unknown generator backend"), "{err}");
assert!(err.contains("wat"), "{err}");
}
/// A backend whose feature is OFF must report "not compiled"; one whose
/// feature is ON must NOT be the "unknown backend" error (its arm exists) —
/// it may Ok or Err on the model, but the backend is recognized. Each spec is
/// gated by ITS OWN feature so the assertion matches the build.
fn assert_backend_arm(spec: &str, compiled: bool) {
let res = generator(spec);
match res {
Ok(_) => assert!(
compiled,
"{spec} produced a generator but its feature is off"
),
Err(e) => {
let s = e.to_string();
if compiled {
assert!(
!s.contains("unknown generator backend") && !s.contains("not compiled"),
"{spec} is compiled — error must be model-level, got: {s}"
);
} else {
assert!(
s.contains("not compiled") || s.contains("rebuild with"),
"{spec} feature is off — expected 'not compiled', got: {s}"
);
}
}
}
}
#[test]
fn factory_reports_uncompiled_backends() {
assert_backend_arm("candle:m", cfg!(feature = "gen-candle"));
assert_backend_arm("mistralrs:m", cfg!(feature = "gen-mistralrs"));
assert_backend_arm("onnx:m", cfg!(feature = "gen-onnx"));
// The unknown-backend arm is unconditional: an unrecognized backend is
// always the "unknown generator backend" error.
match generator("nope:m") {
Ok(_) => panic!("`nope` is not a backend"),
Err(e) => assert!(e.to_string().contains("unknown generator backend"), "{}", e),
}
}
#[test]
fn generator_caller_bridges_to_modelcaller() {
let r#gen = MockGenerator::new("mock");
let caller = GeneratorCaller::new(r#gen);
let ans = caller.call("local-llm", "mock", "2+2?").unwrap();
assert!(ans.output.contains("2+2?"));
assert_eq!(ans.cost_usd, 0.0, "local generation is free");
assert_eq!(ans.mcp_tool_calls, 0);
assert!(ans.tokens_out > 0);
}
#[test]
fn spec_from_env_reads_the_var() {
// Save/restore so we don't disturb a parallel test.
let prev = std::env::var(GEN_BACKEND_ENV).ok();
// TODO: Audit that the environment access only happens in single-threaded code.
unsafe { std::env::set_var(GEN_BACKEND_ENV, "mock:envpick") };
assert_eq!(spec_from_env().as_deref(), Some("mock:envpick"));
// TODO: Audit that the environment access only happens in single-threaded code.
unsafe { std::env::set_var(GEN_BACKEND_ENV, " ") };
assert_eq!(spec_from_env(), None, "blank is treated as unset");
match prev {
// TODO: Audit that the environment access only happens in single-threaded code.
Some(v) => unsafe { std::env::set_var(GEN_BACKEND_ENV, v) },
// TODO: Audit that the environment access only happens in single-threaded code.
None => unsafe { std::env::remove_var(GEN_BACKEND_ENV) },
}
}
}