pipecrab_lm/

model.rs

//! The [`LanguageModel`] trait, its error types ([`LmError`], [`LmConfigError`]),
//! the chat-context types ([`Message`], [`Conversation`], [`GenParams`]), the
//! provider-neutral [`ToolDefinition`], and the streaming [`ModelDelta`] /
//! [`ModelStream`] protocol between an implementation and the stage.

use std::sync::Arc;

use async_trait::async_trait;
use pipecrab_core::ToolCall;
use pipecrab_runtime::MaybeSendSync;

/// A provider-neutral tool definition: name, description, and a JSON Schema for
/// the arguments — the shape every hosted provider expects. It carries no
/// execution callback; tool execution lives outside `pipecrab-lm`.
///
/// [`parameters`](ToolDefinition::parameters) is a [`serde_json::Value`] so a
/// framework's existing schema passes through without a JSON-string round trip
/// and reaches a hosted adapter unchanged.
#[derive(Clone, Debug, PartialEq)]
pub struct ToolDefinition {
    /// Model-facing tool name.
    pub name: Arc<str>,
    /// Human-readable description of what the tool does.
    pub description: Arc<str>,
    /// JSON Schema describing the tool arguments; a JSON object.
    pub parameters: serde_json::Value,
}

impl ToolDefinition {
    /// Build a validated tool definition.
    ///
    /// Fails if `name` is empty or `parameters` is not a JSON object. The schema
    /// is not required to carry a draft identifier or a top-level
    /// `"type": "object"` — providers generate valid schemas with differing
    /// conventions.
    pub fn new(
        name: impl Into<Arc<str>>,
        description: impl Into<Arc<str>>,
        parameters: serde_json::Value,
    ) -> Result<Self, LmConfigError> {
        let name = name.into();
        if name.is_empty() {
            return Err(LmConfigError::EmptyToolName);
        }
        if !parameters.is_object() {
            return Err(LmConfigError::ToolParametersNotObject { name });
        }
        Ok(Self {
            name,
            description: description.into(),
            parameters,
        })
    }
}

/// One turn of conversation history, preserving the structure a hosted adapter
/// needs to reconstruct valid provider history: visible assistant text, the
/// assistant's tool calls, tool-call IDs, tool results, and external events.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Message {
    /// The system prompt: instructions framing the whole conversation.
    System {
        /// The prompt text.
        content: Arc<str>,
    },
    /// A turn from the user.
    User {
        /// The user's text.
        content: Arc<str>,
    },
    /// A turn generated by the model: its visible text and any tool calls it
    /// made, in the order the provider must replay them.
    Assistant {
        /// Visible assistant text; empty for a tool-call-only turn.
        content: Arc<str>,
        /// The tool calls this turn made, correlated to later results by
        /// [`ToolCall::id`].
        tool_calls: Arc<[ToolCall]>,
    },
    /// The result of a previous [`ToolCall`], correlated by `tool_call_id`.
    ToolResult {
        /// The [`ToolCall::id`] this result answers.
        tool_call_id: Arc<str>,
        /// Model-facing name of the tool that produced the result.
        name: Arc<str>,
        /// The tool's output, as text.
        content: Arc<str>,
    },
    /// An event from another system rather than the user.
    Event {
        /// Identifier of the system that produced the event.
        source: Arc<str>,
        /// The kind of event, in the producer's own vocabulary.
        kind: Arc<str>,
        /// The event payload, as text.
        content: Arc<str>,
    },
}

impl Message {
    /// A [`System`](Message::System) message — the framing prompt.
    pub fn system(content: impl Into<Arc<str>>) -> Self {
        Self::System {
            content: content.into(),
        }
    }

    /// A [`User`](Message::User) message.
    pub fn user(content: impl Into<Arc<str>>) -> Self {
        Self::User {
            content: content.into(),
        }
    }

    /// An [`Assistant`](Message::Assistant) message with visible text and no tool
    /// calls.
    pub fn assistant(content: impl Into<Arc<str>>) -> Self {
        Self::Assistant {
            content: content.into(),
            tool_calls: Arc::from([]),
        }
    }

    /// An [`Assistant`](Message::Assistant) message carrying visible text and the
    /// turn's tool calls.
    pub fn assistant_with_tool_calls(
        content: impl Into<Arc<str>>,
        tool_calls: impl Into<Arc<[ToolCall]>>,
    ) -> Self {
        Self::Assistant {
            content: content.into(),
            tool_calls: tool_calls.into(),
        }
    }
}

impl From<pipecrab_core::ModelMessage> for Message {
    /// Lift a core [`ModelMessage`](pipecrab_core::ModelMessage) — a tool result
    /// or an external event — into conversation history.
    fn from(message: pipecrab_core::ModelMessage) -> Self {
        match message {
            pipecrab_core::ModelMessage::ToolResult {
                tool_call_id,
                name,
                content,
            } => Message::ToolResult {
                tool_call_id,
                name,
                content,
            },
            pipecrab_core::ModelMessage::Event {
                source,
                kind,
                content,
            } => Message::Event {
                source,
                kind,
                content,
            },
        }
    }
}

/// The chat context handed to [`LanguageModel::generate`]: an ordered list of
/// [`Message`]s, oldest first. The system prompt, if any, is the first message.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Conversation {
    /// The turns so far, in order.
    pub messages: Vec<Message>,
}

/// Knobs for one [`generate`](LanguageModel::generate) call. All optional; an
/// engine applies its own default for any left `None`.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct GenParams {
    /// Upper bound on generated tokens, if any.
    pub max_tokens: Option<u32>,
    /// Sampling temperature, if any.
    pub temperature: Option<f32>,
    /// An optional grammar / JSON-schema constraint the engine interprets (e.g.
    /// GBNF). Opaque to the pipeline: the trait carries it and the engine
    /// enforces it.
    pub grammar: Option<Arc<str>>,
}

/// One item of the streaming protocol between a [`LanguageModel`] and
/// [`LmStage`](crate::LmStage): a text delta to append, or one complete tool
/// call.
///
/// Not a pipeline frame. Provider-specific streaming formats, tool-call
/// fragments, and constrained-output parsing are assembled *inside* the
/// implementation; the stage sees only normalized text and complete tool calls,
/// and turns them into [`ModelFrame`](pipecrab_core::ModelFrame)s. Every item is
/// a cancellation and barge-in preemption point.
#[derive(Clone, Debug, PartialEq)]
pub enum ModelDelta {
    /// Text to append to the current assistant response.
    Text(Arc<str>),
    /// One complete model tool call.
    ToolCall(ToolCall),
}

impl ModelDelta {
    /// Build a [`ToolCall`](ModelDelta::ToolCall) delta from structured
    /// arguments, giving adapters a JSON-value API while the core frame keeps a
    /// dependency-free JSON-text representation.
    ///
    /// Fails if `id` or `name` is empty, if `arguments` is not a JSON object, or
    /// if serializing the validated object fails.
    pub fn tool_call(
        id: impl Into<Arc<str>>,
        name: impl Into<Arc<str>>,
        arguments: serde_json::Value,
    ) -> Result<Self, LmError> {
        let id = id.into();
        if id.is_empty() {
            return Err(LmError::MissingToolCallId);
        }
        let name = name.into();
        if name.is_empty() {
            return Err(LmError::IncompleteToolCall);
        }
        if !arguments.is_object() {
            return Err(LmError::ToolArgumentsNotObject);
        }
        let arguments_json = serde_json::to_string(&arguments)
            .map_err(|error| LmError::ToolCallSerialization(error.to_string()))?;
        Ok(ModelDelta::ToolCall(ToolCall {
            id,
            name,
            arguments_json: Arc::from(arguments_json),
        }))
    }
}

/// The output of a [`LanguageModel::generate`] call: a boxed stream of
/// [`ModelDelta`] results, delivered delta by delta.
///
/// `BoxStream` on native and `LocalBoxStream` on `wasm32`, behind one cfg'd alias
/// — the same `Send`-where-it-exists split as
/// [`MaybeSend`](pipecrab_runtime::MaybeSend): the pipeline is one logical task
/// that stays `Send` for a work-stealing executor natively, while on `wasm32`
/// (one thread, `!Send` JS handles) that bound must vanish.
#[cfg(not(target_arch = "wasm32"))]
pub type ModelStream = futures::stream::BoxStream<'static, Result<ModelDelta, LmError>>;
/// The output of a [`LanguageModel::generate`] call: a boxed stream of
/// [`ModelDelta`] results, delivered delta by delta.
#[cfg(target_arch = "wasm32")]
pub type ModelStream = futures::stream::LocalBoxStream<'static, Result<ModelDelta, LmError>>;

/// The swappable language-model capability: a chat context in, structured
/// generation out incrementally.
///
/// This is the durable interface. A native engine (e.g. a llama.cpp context) and
/// a hosted engine (a Rig agent, a browser Worker) all implement this one trait,
/// so [`LmStage`](crate::LmStage) — and the pipeline above it — never names a
/// concrete model.
///
/// # Engines are worker-handles
///
/// An implementor is expected to be a thin *handle* to a long-lived worker that
/// owns the model's mutable decode state: a dedicated thread on native (a
/// llama.cpp context is `!Send`, so the worker pattern is mandatory there), a Web
/// Worker on `wasm32`. That is why [`generate`](Self::generate) takes `&self` —
/// it hands the context to the worker and returns its delta stream — and why the
/// worker outlives any single call.
///
/// # Streaming is a barge-in requirement
///
/// [`generate`](Self::generate) yields a [`ModelDelta`] at a time rather than one
/// buffer at the end: every item of the returned [`ModelStream`] is a preemption
/// point the run loop can drop an in-flight generation at, so a user barging in
/// stops the reply within one delta instead of after the whole turn. Dropping the
/// stream is how the *stage* stops pulling; [`cancel`](Self::cancel) is how the
/// *engine* stops producing.
///
/// [`cancel`](Self::cancel) is a *control call* (see
/// [`Processor`](pipecrab_core::Processor)'s control-call carve-out): it maps to
/// the engine's abort callback / an atomic the decode loop checks, so it is
/// synchronous, non-blocking, and safe to invoke directly from a stage's
/// `decide_*` where the barge-in is decided.
///
/// `?Send` on `wasm32` matches pipecrab's single-threaded execution model, so one
/// implementation runs unchanged on a current-thread executor and in the browser,
/// where `Send` bounds cannot be satisfied.
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
pub trait LanguageModel: MaybeSendSync {
    /// Generate a reply to `conversation` under `params` with `tools` available,
    /// yielding [`ModelDelta`]s.
    ///
    /// `tools` are the tools configured on the stage. An implementation that
    /// wraps a higher-level agent (e.g. a Rig agent) keeps its own registered
    /// tools internal and uses them here directly, alongside any in `tools`.
    ///
    /// Takes `&self`: like every [`Stage::perform`](pipecrab_runtime::Stage::perform),
    /// generation must not mutate observable state, so the run loop can drop an
    /// in-flight call — at any stream item — on a barge-in interrupt without
    /// tearing anything. Every item of the returned [`ModelStream`] is such a
    /// preemption point.
    async fn generate(
        &self,
        conversation: &Conversation,
        params: &GenParams,
        tools: &[ToolDefinition],
    ) -> Result<ModelStream, LmError>;

    /// Control call: abort in-flight generation. Sync, non-blocking, idempotent.
    ///
    /// Maps to the engine's abort callback / an atomic the decode loop checks; the
    /// next [`generate`](Self::generate) starts clean. Safe to call from a stage's
    /// synchronous `decide_*` — see the trait-level note.
    fn cancel(&self);

    /// Checkpoint the worker's session state (KV cache and any decode state) to an
    /// opaque, serialized blob.
    async fn save_state(&self) -> Result<Vec<u8>, LmError>;

    /// Restore worker session state previously produced by
    /// [`save_state`](Self::save_state).
    async fn load_state(&self, blob: &[u8]) -> Result<(), LmError>;
}

/// Why configuring an [`LmStage`](crate::LmStage) or a [`ToolDefinition`] failed —
/// a static, provider-neutral error surfaced at construction, not per generation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LmConfigError {
    /// A tool definition had an empty name.
    EmptyToolName,
    /// A tool's `parameters` was not a JSON object.
    ToolParametersNotObject {
        /// The offending tool's name.
        name: Arc<str>,
    },
    /// Two tools in the effective set shared a name.
    DuplicateToolName {
        /// The duplicated name.
        name: Arc<str>,
    },
}

impl std::fmt::Display for LmConfigError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LmConfigError::EmptyToolName => write!(f, "tool name must be nonempty"),
            LmConfigError::ToolParametersNotObject { name } => {
                write!(f, "parameters of tool {name:?} must be a JSON object")
            }
            LmConfigError::DuplicateToolName { name } => {
                write!(f, "duplicate tool name {name:?}")
            }
        }
    }
}

impl std::error::Error for LmConfigError {}

/// Why a [`LanguageModel`] call failed.
///
/// Provider-neutral: an adapter maps its engine's native error into one of these
/// kinds, so no provider-specific error type reaches the public interface.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LmError {
    /// The generation engine itself failed — an inference error, a worker that
    /// crashed, a model that never loaded. Carries a human-readable description.
    Engine(String),
    /// The provider stream failed mid-generation.
    ProviderStream(String),
    /// A tool call's arguments were structurally invalid.
    InvalidToolArguments(String),
    /// Tool-call arguments were not a JSON object.
    ToolArgumentsNotObject,
    /// A provider tool call was incomplete — missing its name or arguments.
    IncompleteToolCall,
    /// A provider tool call carried no identifier.
    MissingToolCallId,
    /// Serializing validated tool arguments to JSON text failed.
    ToolCallSerialization(String),
}

impl std::fmt::Display for LmError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LmError::Engine(msg) => write!(f, "lm engine error: {msg}"),
            LmError::ProviderStream(msg) => write!(f, "lm provider stream error: {msg}"),
            LmError::InvalidToolArguments(msg) => write!(f, "invalid tool-call arguments: {msg}"),
            LmError::ToolArgumentsNotObject => {
                write!(f, "tool-call arguments must be a JSON object")
            }
            LmError::IncompleteToolCall => write!(f, "incomplete provider tool call"),
            LmError::MissingToolCallId => write!(f, "missing tool-call identifier"),
            LmError::ToolCallSerialization(msg) => {
                write!(f, "tool-call serialization failed: {msg}")
            }
        }
    }
}

impl std::error::Error for LmError {}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn tool_definition_rejects_empty_name() {
        assert_eq!(
            ToolDefinition::new("", "d", json!({})),
            Err(LmConfigError::EmptyToolName)
        );
    }

    #[test]
    fn tool_definition_rejects_non_object_schema() {
        let err = ToolDefinition::new("t", "d", json!("not an object")).unwrap_err();
        assert_eq!(
            err,
            LmConfigError::ToolParametersNotObject {
                name: Arc::from("t")
            }
        );
    }

    #[test]
    fn tool_definition_preserves_the_schema_value_unchanged() {
        let schema = json!({
            "type": "object",
            "properties": { "city": { "type": "string" } },
            "required": ["city"],
        });
        let def = ToolDefinition::new("weather", "look up weather", schema.clone()).unwrap();
        assert_eq!(
            def.parameters, schema,
            "the schema value is stored verbatim"
        );
    }

    #[test]
    fn tool_call_delta_normalizes_object_arguments_to_json_text() {
        let ModelDelta::ToolCall(call) =
            ModelDelta::tool_call("call-1", "weather", json!({ "city": "paris" })).unwrap()
        else {
            unreachable!("tool_call builds a ToolCall delta")
        };
        assert_eq!(&*call.id, "call-1");
        assert_eq!(&*call.name, "weather");
        // Core carries JSON *text*; it must parse back to the structured input.
        let parsed: serde_json::Value = serde_json::from_str(&call.arguments_json).unwrap();
        assert_eq!(parsed, json!({ "city": "paris" }));
    }

    #[test]
    fn tool_call_delta_rejects_non_object_arguments() {
        assert_eq!(
            ModelDelta::tool_call("call-1", "t", json!("scalar")),
            Err(LmError::ToolArgumentsNotObject)
        );
        assert_eq!(
            ModelDelta::tool_call("call-1", "t", json!([1, 2, 3])),
            Err(LmError::ToolArgumentsNotObject)
        );
    }

    #[test]
    fn tool_call_delta_requires_id_and_name() {
        assert_eq!(
            ModelDelta::tool_call("", "t", json!({})),
            Err(LmError::MissingToolCallId)
        );
        assert_eq!(
            ModelDelta::tool_call("call-1", "", json!({})),
            Err(LmError::IncompleteToolCall)
        );
    }
}