trusty_common/chat/

mod.rs

//! Provider-agnostic streaming chat abstraction with tool-use support.
//!
//! Why: trusty-memory and trusty-search both want to support more than one
//! upstream LLM (OpenRouter for cloud, Ollama / LM Studio for local). Rather
//! than each crate re-implementing the dispatch, we expose a small
//! [`ChatProvider`] trait plus concrete implementations and an auto-detector
//! for a running local model server. The trait also surfaces OpenAI-style
//! tool/function calling so downstream agents can let the model invoke tools
//! (search, memory recall, shell, etc.).
//!
//! What: defines the [`ChatProvider`] trait, [`ToolDef`] / [`ToolCall`] /
//! [`ChatEvent`] tool-use types, an [`OpenRouterProvider`] and an
//! [`OllamaProvider`] that both speak OpenAI-compatible
//! `/v1/chat/completions` with SSE streaming (including the streamed
//! `tool_calls` shape), a [`BedrockProvider`] that uses the AWS Bedrock
//! `Converse` API (behind the `bedrock` feature flag), and
//! [`auto_detect_local_provider`] which probes `{base_url}/v1/models` with a
//! 1-second timeout.
//!
//! Test: `cargo test -p trusty-common` covers default config values, the
//! unreachable-server path of `auto_detect_local_provider`, SSE delta
//! streaming, and accumulation of streamed tool-call fragments.

mod openai_compat;

#[cfg(feature = "bedrock")]
mod bedrock_impl;
#[cfg(not(feature = "bedrock"))]
mod bedrock_stub;

pub use openai_compat::{OllamaProvider, OpenRouterProvider, auto_detect_local_provider};

#[cfg(feature = "bedrock")]
pub use bedrock_impl::{
    BedrockProvider, DEFAULT_BEDROCK_MODEL, DEFAULT_BEDROCK_REGION, ENV_REGION_AWS,
    ENV_REGION_TRUSTY,
};

// Re-expose the bedrock_impl module as `bedrock_provider` so downstream
// crates can access constants (e.g. `DEFAULT_BEDROCK_MODEL`) without needing
// to depend on the bedrock feature themselves.
#[cfg(feature = "bedrock")]
pub mod bedrock_provider {
    pub use super::bedrock_impl::*;
}

#[cfg(not(feature = "bedrock"))]
pub use bedrock_stub::BedrockProvider;

// Stub constant so code that references DEFAULT_BEDROCK_MODEL compiles without
// the bedrock feature. Must stay in sync with bedrock_impl::DEFAULT_BEDROCK_MODEL.
// Claude Sonnet 4.6 drops the date stamp and -v1:0 suffix (verified vs AWS docs).
#[cfg(not(feature = "bedrock"))]
pub const DEFAULT_BEDROCK_MODEL: &str = "us.anthropic.claude-sonnet-4-6";

use crate::ChatMessage;
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc::Sender;

// ── Public re-exports so callers get the full surface from `chat::*` ──────────

/// Configuration for a local OpenAI-compatible model server (Ollama, LM
/// Studio, llama.cpp's server, etc.).
///
/// Why: callers want a single struct they can deserialize from config files
/// and pass to [`auto_detect_local_provider`] without juggling defaults.
/// What: holds an enable flag, the server's base URL (no trailing slash),
/// and the default model to request. Defaults target Ollama's standard
/// localhost binding.
/// Test: `local_model_config_defaults` asserts the default values.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalModelConfig {
    pub enabled: bool,
    pub base_url: String,
    pub model: String,
}

impl Default for LocalModelConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            base_url: "http://localhost:11434".to_string(),
            model: "qwen3:30b".to_string(),
        }
    }
}

// ─── Tool-use types ───────────────────────────────────────────────────────────

/// JSON-Schema description of a callable tool, in OpenAI function-calling
/// shape.
///
/// Why: downstream agents (trusty-memory, trusty-search) expose tools like
/// `memory_recall` or `web_search` to the LLM. The OpenAI tool format is the
/// de-facto common denominator across OpenRouter, Ollama, LM Studio, and
/// most cloud providers.
/// What: `name` and `description` are passed verbatim; `parameters` is a
/// JSON Schema object (typically `{"type":"object","properties":{...}}`).
/// Test: `tool_def_serializes_as_function` checks the wire shape.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDef {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

/// A tool invocation the model wants the host to perform.
///
/// Why: the streaming chat API emits `tool_calls` in fragments — first an
/// `id` + `function.name`, then a string of `function.arguments` deltas.
/// We accumulate fragments and surface one fully-formed [`ToolCall`] per
/// invocation to the caller.
/// What: `id` is the upstream's call id (echoed back in subsequent
/// `role:"tool"` messages); `name` is the function name; `arguments` is a
/// JSON string (NOT a parsed value — many models emit malformed JSON and
/// callers want the raw text for error reporting / repair).
/// Test: `accumulates_streamed_tool_call_fragments`.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub arguments: String,
}

/// Sampling knobs forwarded to an OpenAI-compatible provider.
///
/// Why (issue #3758): the streaming request wire previously sent only
/// `model`/`messages`/`tools`, so a streamed reply's style, verbosity, and
/// stopping behaviour were NOT equivalent to the blocking path's for the same
/// turn — the caller silently got provider defaults instead of its configured
/// temperature, token ceiling, and stop sequences. Carrying them in one struct
/// (rather than three provider constructor arguments) keeps `new()` stable for
/// existing callers and lets the set grow without churning call sites.
/// What: every field is optional — `None`/empty means "omit from the request
/// body and let the provider default apply", which is exactly the pre-#3758
/// behaviour, so a caller that does not opt in is unaffected. `stop` maps to
/// the OpenAI `stop` array (the direct-Anthropic dialect calls the same thing
/// `stop_sequences`).
/// Test: `sampling_params_serialize_into_request_body`,
/// `default_sampling_omits_fields`.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct SamplingParams {
    /// Sampling temperature; `None` omits the field.
    pub temperature: Option<f32>,
    /// Maximum tokens to generate; `None` omits the field.
    pub max_tokens: Option<u32>,
    /// Stop sequences; empty omits the field.
    pub stop: Vec<String>,
}

impl SamplingParams {
    /// The stop sequences as a request-body-ready `Option<&[String]>`.
    ///
    /// Why: an EMPTY `stop` array is not the same as an absent one — some
    /// OpenAI-compatible servers reject `"stop": []` outright, which is the
    /// same trap [`ChatProvider`] already documents for an empty `tools`
    /// array. Collapsing empty to `None` at the boundary keeps that decision
    /// in one place instead of at each provider.
    /// What: `None` when empty, `Some(slice)` otherwise.
    /// Test: `default_sampling_omits_fields`.
    pub fn stop_slice(&self) -> Option<&[String]> {
        if self.stop.is_empty() {
            None
        } else {
            Some(&self.stop)
        }
    }
}

/// Token-usage tally surfaced at the end of a streamed response.
///
/// Why (issue #3767): Bedrock's `ConverseStream` reports usage only once, in
/// a terminal `metadata` event — unlike the per-token deltas, it has no
/// natural home in [`ChatEvent::Delta`]. Without a dedicated variant a
/// streaming provider has no way to hand token counts back to the caller, so
/// a downstream cost-reporting consumer would silently see zero usage for
/// every streamed call. Kept as a small standalone struct (rather than
/// reusing `trusty_common::inference::types::Usage`) so the `chat` module —
/// which compiles unconditionally, unlike the `inference-client`-gated
/// `types` module — never depends on a feature-gated type.
/// What: four token buckets mirroring the shape every provider in this
/// workspace already reports (prompt/completion split, plus the two
/// prompt-cache buckets). The four fields are SEPARATE and ADDITIVE — not
/// nested — matching both the Bedrock `TokenUsage` wire shape
/// (`input_tokens`/`output_tokens`/`cache_read_input_tokens`/
/// `cache_write_input_tokens` are independent fields, verified against
/// `aws-sdk-bedrockruntime` 1.132.0) and this workspace's existing
/// convention (`crate::perf::TokenUsage`,
/// `trusty-agents/src/llm/anthropic_native/mod.rs`'s usage parsing, and
/// `chat::bedrock_impl::parse_usage`'s non-streaming counterpart). A
/// provider that has none of the cache fields simply leaves them zero.
/// Test: `bedrock_stream_reports_usage_from_metadata_event` (bedrock_impl).
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ChatUsage {
    /// Prompt (input) tokens. Separate from `cache_read_tokens` — a
    /// cost-accounting consumer that assumes this already includes cache
    /// reads will double-count; one that assumes it excludes the full
    /// prompt will undercount. Sum with `cache_read_tokens` for the total
    /// prompt-side token count when that is what's needed.
    pub prompt_tokens: u32,
    /// Completion (output) tokens produced.
    pub completion_tokens: u32,
    /// Prompt tokens served from the provider's prompt cache — additive
    /// with `prompt_tokens`, not a subset of it.
    pub cache_read_tokens: u32,
    /// Prompt tokens newly written into the prompt cache this turn.
    pub cache_creation_tokens: u32,
}

/// Streaming chat event.
///
/// Why: replaces the previous "string-only" channel so callers can
/// distinguish text deltas from tool invocations and from terminal
/// success/error without parsing magic markers out of the text stream.
/// What: `Delta` is a content chunk; `ToolCall` is a fully-accumulated tool
/// invocation; `Usage` carries the call's token tally (emitted zero or one
/// times, whenever the provider's wire format reports it — see
/// [`ChatUsage`]); `Done` signals the upstream stream terminated normally;
/// `Error` carries a human-readable message for stream-mid failures (the
/// provider also returns `Err` from `chat_stream`, but `Error` lets the
/// caller display partial-stream failures inline).
/// Test: `ollama_provider_streams_sse_deltas`.
///
/// # Stability
///
/// `#[non_exhaustive]`: adding a variant here is otherwise a SemVer-breaking
/// change for every downstream crate, because an exhaustive `match` over the
/// old variant set stops compiling with E0004. That is not hypothetical — the
/// `Usage` variant forced arm additions in five consumers at once
/// (`trusty-agents`, `trusty-analyze`, `trusty-memory`, `trusty-mpm`,
/// `trusty-search`) and is the direct cause of the 0.27.0 MINOR bump: a patch
/// release would have re-resolved the already-published, arm-less consumer
/// sources against the new variant and hard-failed `cargo install` (the same
/// failure that forced `trusty-analyze` 0.7.3 to be yanked).
///
/// Downstream matches must therefore carry a wildcard arm. This attribute does
/// NOT retroactively fix consumers published before it landed — it only stops
/// the next variant addition from repeating the break.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ChatEvent {
    Delta(String),
    ToolCall(ToolCall),
    Usage(ChatUsage),
    Done,
    Error(String),
}

/// Streaming chat provider abstraction.
///
/// Why: downstream crates (trusty-memory, trusty-search) want to support
/// multiple LLM backends without hard-coding which one to call. Providers
/// expose a uniform streaming interface so the caller can swap them at
/// runtime based on configuration / availability.
/// What: implementors stream [`ChatEvent`]s into `tx`. Pass an empty
/// `tools` vec to disable tool use entirely (the provider MUST then omit
/// the `tools` field from the upstream request — some models error on an
/// empty array). Returning `Ok(())` means the stream completed normally;
/// the caller should also expect a final [`ChatEvent::Done`].
/// Test: implementations are covered by their own unit tests in this
/// module plus integration tests in downstream crates.
#[async_trait]
pub trait ChatProvider: Send + Sync {
    /// Human-readable provider name (e.g. `"openrouter"`, `"ollama"`).
    fn name(&self) -> &str;
    /// Model identifier sent on every request.
    fn model(&self) -> &str;
    /// Stream chat events into `tx`. `tools` empty disables tool use.
    async fn chat_stream(
        &self,
        messages: Vec<ChatMessage>,
        tools: Vec<ToolDef>,
        tx: Sender<ChatEvent>,
    ) -> Result<()>;
}

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

    #[test]
    fn local_model_config_defaults() {
        let cfg = LocalModelConfig::default();
        assert!(cfg.enabled);
        assert_eq!(cfg.base_url, "http://localhost:11434");
        assert_eq!(cfg.model, "qwen3:30b");
    }

    #[test]
    fn local_model_config_deserializes_from_toml() {
        let toml_src = r#"
            enabled = true
            base_url = "http://localhost:1234"
            model = "qwen2.5-coder"
        "#;
        let cfg: LocalModelConfig = toml::from_str(toml_src).expect("parse TOML");
        assert!(cfg.enabled);
        assert_eq!(cfg.base_url, "http://localhost:1234");
        assert_eq!(cfg.model, "qwen2.5-coder");
    }
}