turbospark-server 0.3.0

//! The pluggable generation backend a request is served from.
//!
//! Two implementations: [`ScriptedChatModel`], which drives the
//! raw-completion loop with a fixed, pre-scripted logit sequence (portable,
//! and what the integration tests use to prove the HTTP envelopes, chat
//! templating, and SSE framing), and [`crate::RealChatModel`] (macOS only),
//! which drives a real `RealForwardRunner` forward pass.
//!
//! [`ChatModel::with_producer`] is a pass-through rather than a
//! `new_producer() -> Box<dyn LogitProducer>` factory on purpose: a real
//! runner costs a multi-gigabyte mmap plus a full Metal pipeline compile to
//! open, so there is exactly one per process and it is borrowed mutably for
//! the duration of a request, not handed out by value.

use runtime::{
    run_raw_completion_cancellable, CancelFlag, GenerationConfig, LogitProducer, RawDecodeProgress,
    RawDecodeResult, RuntimeError,
};
use tokenizer::MfTokenizer;

pub trait ChatModel: Send + Sync {
    fn tokenizer(&self) -> &MfTokenizer;
    fn vocab_size(&self) -> usize;
    fn max_context(&self) -> u32;

    /// The canonical id `GET /v1/models` advertises. A registry routes an
    /// exact request for this id (or an advertised alias) to this backend.
    fn model_id(&self) -> &str;

    /// Extra ids `GET /v1/models` advertises for this chat backend.
    ///
    /// Claude Code gateway discovery keeps only ids containing `claude` or
    /// `anthropic`. Giving every generative backend its own derived alias
    /// keeps that discovery path useful on a multi-model local server without
    /// pretending a local checkpoint is an Anthropic product model.
    fn model_aliases(&self) -> Vec<String> {
        vec![format!("claude-turbospark-{}", self.model_id())]
    }

    /// Lends a producer to `f` for one generation. Implementations may
    /// serialize concurrent calls; `run_raw_completion` resets the producer
    /// on entry, so a producer reused across calls carries no state over.
    fn with_producer(
        &self,
        f: &mut dyn FnMut(&mut dyn LogitProducer) -> Result<RawDecodeResult, RuntimeError>,
    ) -> Result<RawDecodeResult, RuntimeError>;

    /// Runs one generation, and owns the choice of WHICH decode loop.
    ///
    /// **A method rather than a call in `exec.rs`, because the speculative
    /// loop cannot be reached through [`Self::with_producer`] at all.**
    /// `run_raw_completion_speculative` is generic over
    /// `runtime::SpeculativeProducer`, which carries an associated
    /// `Checkpoint` type and so is not object-safe; `with_producer` hands out
    /// a `&mut dyn LogitProducer`. Only a backend holding the CONCRETE runner
    /// can call it, so the decision belongs to the backend.
    ///
    /// The default is the sequential loop through `with_producer`, which is
    /// byte for byte what every caller did before this method existed --
    /// `ScriptedChatModel` therefore needs no implementation and every
    /// integration test keeps the exact path it had.
    /// This install's vision configuration, or `None` where it has no tower.
    ///
    /// Read off the INSTALL by the backend: the pixel budget and the special
    /// ids are per-checkpoint, and a constant in the handler would be
    /// AGENTS.md Gotcha 38's shape. `None` is what makes an image request a
    /// reported DROP rather than a silent one.
    fn vision(&self) -> Option<crate::vision::VisionInfo> {
        None
    }

    /// Whether this model supports dense vector embedding generation.
    fn supports_embeddings(&self) -> bool {
        false
    }

    /// Encode input texts into dense vector embeddings.
    fn encode(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>, String> {
        Err(format!(
            "model {} does not support embeddings",
            self.model_id()
        ))
    }

    /// Runs one generation, and owns the choice of WHICH decode loop.
    ///
    /// **`images` CANNOT BE A SEPARATE CALL, and that is a concurrency
    /// property rather than a style choice.** A backend serializes on its one
    /// runner per call (Gotcha 1), so `set_prompt_vision` followed by
    /// `run_completion` would be two locks with a gap: a second request
    /// arriving in that gap overwrites the map, and the first generation then
    /// prefills the second's picture. Both are fluent. Passing the images
    /// here is what puts the encode, the injection and the decode inside ONE
    /// lock.
    ///
    /// `cancel` is polled once per prefill and decoded token by the runtime
    /// loop underneath; a caller with nothing to cancel on passes `&|| false`,
    /// which is the same statement sequence this ran before cancellation
    /// existed (`runtime::run_raw_completion`'s own `NEVER` constant, not
    /// exported past that crate, is exactly this).
    fn run_completion(
        &self,
        prompt_ids: &[foundation::TokenId],
        config: &GenerationConfig,
        images: Option<&crate::vision::RequestImages>,
        cancel: CancelFlag<'_>,
        on_progress: &mut dyn FnMut(RawDecodeProgress),
    ) -> Result<RawDecodeResult, RuntimeError> {
        if images.is_some() {
            // The DEFAULT backend is the scripted one, which has no tower and
            // no runner to encode with. Refused BY NAME rather than dropped:
            // a caller who sent a picture and got a text answer would read it
            // as the model ignoring the image.
            return Err(RuntimeError::Producer(
                "this backend cannot encode images; run the server against a .gturbo install \
                 whose checkpoint carries a vision tower"
                    .to_string(),
            ));
        }
        self.with_producer(&mut |producer| {
            run_raw_completion_cancellable(
                producer,
                self.tokenizer(),
                prompt_ids,
                config,
                self.max_context(),
                self.vocab_size(),
                cancel,
                &mut *on_progress,
            )
        })
    }

    /// The decode rate cap and thermal stepping this backend generates
    /// under (ROADMAP Phase P2). Process-level, not per request: there is
    /// one runner per process and a power setting is a property of the
    /// machine, not of a caller's prompt. The default is uncapped, which is
    /// what the scripted backend wants.
    fn rate_control(&self) -> runtime::RateControl {
        runtime::RateControl::default()
    }

    /// The FIFO admission gate this backend's generations pass through
    /// (ROADMAP P1 item 5, `crate::queue`). `None` -- the default -- means
    /// ungated: the scripted backend serializes on nothing today and keeps
    /// that behavior, so the integration suite's timing is untouched.
    ///
    /// A backend that owns a single contended runner (the real one) returns
    /// ONE gate per runner, and the request paths acquire it in ARRIVAL
    /// order before touching the runner -- the fairness the runner's own
    /// mutex cannot provide. Callers acquire at a closed, non-nesting set
    /// of async sites (`run_guarded`, the live-stream spawn sites, ollama's
    /// `run`, `completions::run_full`, embeddings); acquiring twice on one
    /// request would deadlock on the single permit, so the acquisition
    /// points are load-bearing and documented in `queue.rs`.
    fn generation_queue(&self) -> Option<std::sync::Arc<crate::queue::GenerationQueue>> {
        None
    }

    /// Which tool-call guardrails this backend generates under.
    ///
    /// Process-level for exactly the reason [`ChatModel::rate_control`] is,
    /// and NOT for a symmetry: whether a deployment repairs a malformed tool
    /// call is a property of the deployment, and a per-request field would let
    /// any client opt its own traffic out of it.
    ///
    /// The default is ON, which costs the scripted backend nothing: every
    /// guardrail is keyed on the request carrying tools, and a request with
    /// none takes the path it always took.
    fn guardrails(&self) -> crate::guardrails::GuardrailConfig {
        crate::guardrails::GuardrailConfig::default()
    }

    /// Default reasoning effort level when a request omits `reasoning_effort`.
    fn default_reasoning(&self) -> tokenizer::ReasoningEffort {
        tokenizer::ReasoningEffort::Off
    }

    /// Deployment-wide system prompt for requests that carry no system or
    /// developer message of their own (`--system` / `--system-file`).
    ///
    /// Process-level for `default_reasoning`'s reason, and applied at ONE
    /// place -- `handler::plan` -- because every wire format this server
    /// speaks converts to `ChatCompletionRequest` before it reaches there.
    fn default_system(&self) -> Option<&str> {
        None
    }
}

/// Always replays the same scripted logit sequence, regardless of the
/// prompt. `steps` must include one entry per prefill token plus one per
/// decode step the caller expects to reach before a stop condition.
pub struct ScriptedChatModel {
    tokenizer: MfTokenizer,
    vocab_size: usize,
    max_context: u32,
    steps: Vec<Vec<foundation::LogitValue>>,
    default_reasoning: tokenizer::ReasoningEffort,
    default_system: Option<String>,
}

impl ScriptedChatModel {
    pub fn new(
        tokenizer: MfTokenizer,
        max_context: u32,
        steps: Vec<Vec<foundation::LogitValue>>,
    ) -> Self {
        let vocab_size = tokenizer.vocab_size;
        Self {
            tokenizer,
            vocab_size,
            max_context,
            steps,
            default_reasoning: tokenizer::ReasoningEffort::Off,
            default_system: None,
        }
    }

    pub fn with_default_reasoning(mut self, reasoning: tokenizer::ReasoningEffort) -> Self {
        self.default_reasoning = reasoning;
        self
    }

    pub fn with_default_system(mut self, system: impl Into<String>) -> Self {
        self.default_system = Some(system.into());
        self
    }
}

impl ChatModel for ScriptedChatModel {
    fn tokenizer(&self) -> &MfTokenizer {
        &self.tokenizer
    }

    fn vocab_size(&self) -> usize {
        self.vocab_size
    }

    fn max_context(&self) -> u32 {
        self.max_context
    }

    fn model_id(&self) -> &str {
        "scripted"
    }

    fn default_reasoning(&self) -> tokenizer::ReasoningEffort {
        self.default_reasoning
    }

    fn default_system(&self) -> Option<&str> {
        self.default_system.as_deref()
    }

    fn with_producer(
        &self,
        f: &mut dyn FnMut(&mut dyn LogitProducer) -> Result<RawDecodeResult, RuntimeError>,
    ) -> Result<RawDecodeResult, RuntimeError> {
        let mut producer = runtime::ScriptedLogitProducer::new(self.steps.clone());
        f(&mut producer)
    }

    fn supports_embeddings(&self) -> bool {
        true
    }

    fn encode(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, String> {
        let mut out = Vec::with_capacity(texts.len());
        for text in texts {
            let len = text.len() as f32;
            let mut vec = vec![1.0, len * 0.1, 0.5, -0.2];
            let norm: f32 = vec.iter().map(|v| v * v).sum::<f32>().sqrt();
            for x in &mut vec {
                *x /= norm;
            }
            out.push(vec);
        }
        Ok(out)
    }
}