klieo_core/

llm.rs

//! LLM client trait and request/response types.
//!
//! Providers implement [`LlmClient`]. The runtime never depends on a
//! specific provider crate — selection is explicit at app startup.

/// Re-exported so the error this module's traits raise is reachable from the
/// module that raises it. `LlmError` is defined in [`crate::error`]; without
/// this, `klieo_core::llm::LlmError` — the path a caller writing against
/// [`LlmClient`] naturally reaches for — does not resolve.
pub use crate::error::LlmError;
use async_trait::async_trait;
use futures_core::Stream;
use serde::{Deserialize, Serialize};
use std::pin::Pin;
use std::time::Duration;

/// One message in a chat conversation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Speaker role.
    pub role: Role,
    /// Body text. May be empty when `tool_calls` carries the payload.
    pub content: String,
    /// Tool calls the assistant requested in this message.
    #[serde(default)]
    pub tool_calls: Vec<ToolCall>,
    /// When this message answers a tool call, the id of that call.
    #[serde(default)]
    pub tool_call_id: Option<String>,
}

impl Message {
    /// [`Role::System`] message; `tool_calls` is empty and `tool_call_id` is `None`.
    pub fn system(content: impl Into<String>) -> Self {
        Self::text(Role::System, content)
    }

    /// [`Role::User`] message; `tool_calls` is empty and `tool_call_id` is `None`.
    pub fn user(content: impl Into<String>) -> Self {
        Self::text(Role::User, content)
    }

    /// [`Role::Assistant`] message; `tool_calls` is empty and `tool_call_id` is `None`.
    pub fn assistant(content: impl Into<String>) -> Self {
        Self::text(Role::Assistant, content)
    }

    fn text(role: Role, content: impl Into<String>) -> Self {
        Self {
            role,
            content: content.into(),
            tool_calls: Vec::new(),
            tool_call_id: None,
        }
    }
}

/// Speaker role.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[non_exhaustive]
pub enum Role {
    /// System prompt.
    System,
    /// User-supplied input.
    User,
    /// Model output.
    Assistant,
    /// Tool call result fed back into the conversation.
    Tool,
}

/// One tool call requested by the assistant.
///
/// Marked `#[non_exhaustive]` — future field additions are
/// additive (no SemVer-major bump required). Code outside `klieo-core`
/// that constructs a `ToolCall` must use [`ToolCall::new`] instead of
/// a struct literal; pattern matching must use the `..` rest pattern.
/// In-crate construction is unaffected.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolCall {
    /// Provider-stable id; echoed back in the matching tool response.
    pub id: String,
    /// Tool name.
    pub name: String,
    /// JSON arguments.
    pub args: serde_json::Value,
}

impl ToolCall {
    /// Prefer this over struct literals outside `klieo-core` —
    /// `#[non_exhaustive]` forbids external struct-literal construction so that
    /// future field additions remain additive.
    pub fn new(id: impl Into<String>, name: impl Into<String>, args: serde_json::Value) -> Self {
        Self {
            id: id.into(),
            name: name.into(),
            args,
        }
    }
}

/// Tool catalogue entry shown to the LLM.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolDef {
    /// Tool name (must be unique within the catalogue).
    pub name: String,
    /// Human-readable description for the LLM.
    pub description: String,
    /// JSON-schema for arguments.
    pub json_schema: serde_json::Value,
}

impl ToolDef {
    /// One entry in the tool catalogue offered to the model.
    ///
    /// Prefer this constructor over struct literals in code outside
    /// `klieo-core` — `#[non_exhaustive]` forbids external struct-literal
    /// construction so that future field additions remain additive.
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        json_schema: serde_json::Value,
    ) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            json_schema,
        }
    }
}

/// Chat completion request.
#[derive(Debug, Clone)]
pub struct ChatRequest {
    /// Conversation history.
    pub messages: Vec<Message>,
    /// Tools available for the model to call.
    pub tools: Vec<ToolDef>,
    /// Sampling temperature.
    pub temperature: Option<f32>,
    /// Maximum response tokens.
    pub max_tokens: Option<u32>,
    /// Output format hint.
    pub response_format: ResponseFormat,
    /// Stop sequences.
    pub stop: Vec<String>,
    /// Per-request deadline. Provider should abort if exceeded.
    pub timeout: Option<Duration>,
}

impl ChatRequest {
    /// Build a request with the supplied messages and no tools.
    pub fn new(messages: Vec<Message>) -> Self {
        Self {
            messages,
            tools: Vec::new(),
            temperature: None,
            max_tokens: None,
            response_format: ResponseFormat::Text,
            stop: Vec::new(),
            timeout: None,
        }
    }

    /// Start building a request; messages are sent in the order they are added.
    ///
    /// ```
    /// use klieo_core::llm::ChatRequest;
    /// let req = ChatRequest::builder().system("be terse").user("hi").build();
    /// assert_eq!(req.messages.len(), 2);
    /// ```
    pub fn builder() -> ChatRequestBuilder {
        ChatRequestBuilder::default()
    }
}

/// Builder for [`ChatRequest`]: messages accumulate in call order, and any
/// option left unset falls back to the [`ChatRequest::new`] default on `build`.
#[derive(Default)]
pub struct ChatRequestBuilder {
    messages: Vec<Message>,
    tools: Vec<ToolDef>,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    response_format: Option<ResponseFormat>,
    timeout: Option<Duration>,
}

impl ChatRequestBuilder {
    /// Appends a [`Role::System`] message after any already added.
    pub fn system(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::system(content));
        self
    }

    /// Appends a [`Role::User`] message after any already added.
    pub fn user(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::user(content));
        self
    }

    /// Appends a [`Role::Assistant`] message after any already added.
    pub fn assistant(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::assistant(content));
        self
    }

    /// Appends a pre-built message — the only way to add a tool-result turn.
    pub fn message(mut self, message: Message) -> Self {
        self.messages.push(message);
        self
    }

    /// Replaces the tool catalogue; empty (no tools) when never called.
    pub fn tools(mut self, tools: Vec<ToolDef>) -> Self {
        self.tools = tools;
        self
    }

    /// Sampling temperature; left to the provider default when never called.
    pub fn temperature(mut self, temperature: f32) -> Self {
        self.temperature = Some(temperature);
        self
    }

    /// Response-token cap; provider default when never called.
    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    /// Output format; falls back to [`ResponseFormat::Text`] when never called.
    pub fn response_format(mut self, response_format: ResponseFormat) -> Self {
        self.response_format = Some(response_format);
        self
    }

    /// Per-request deadline; none (provider transport governs) when never called.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    /// Builds the request; the `stop` list is always empty and other unset
    /// options take their [`ChatRequest::new`] defaults.
    pub fn build(self) -> ChatRequest {
        ChatRequest {
            messages: self.messages,
            tools: self.tools,
            temperature: self.temperature,
            max_tokens: self.max_tokens,
            response_format: self.response_format.unwrap_or(ResponseFormat::Text),
            stop: Vec::new(),
            timeout: self.timeout,
        }
    }
}

/// Output format hint passed to the provider.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ResponseFormat {
    /// Plain text.
    Text,
    /// Best-effort JSON output.
    Json {
        /// Schema describing expected JSON shape.
        schema: serde_json::Value,
    },
    /// Strict structured output validated against the schema.
    StructuredOutput {
        /// Schema describing expected shape.
        schema: serde_json::Value,
    },
}

/// Chat completion response.
///
/// Marked `#[non_exhaustive]` — future field additions are additive.
/// All construction outside `klieo-core` must use [`ChatResponse::new`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ChatResponse {
    /// Assistant message returned by the provider.
    pub message: Message,
    /// Token usage.
    pub usage: Usage,
    /// Why the provider stopped generating.
    pub finish_reason: FinishReason,
    /// Model that served this response — provider-reported when available
    /// (Ollama), else the client-configured model the request used. Recorded
    /// into `Capture.model_version` for reproducibility; best-effort metadata,
    /// not an authorization or correctness input.
    pub model: String,
}

impl ChatResponse {
    /// Provider response for one completion.
    ///
    /// `model` is the provider-reported model name when the response body
    /// carries one (e.g. Ollama), or the configured model used for the
    /// request otherwise.
    pub fn new(
        message: Message,
        usage: Usage,
        finish_reason: FinishReason,
        model: impl Into<String>,
    ) -> Self {
        Self {
            message,
            usage,
            finish_reason,
            model: model.into(),
        }
    }
}

/// Token usage report.
///
/// Marked `#[non_exhaustive]` — future field additions are
/// additive (no SemVer-major bump required). Code outside `klieo-core`
/// that constructs a `Usage` must use [`Usage::new`] instead of
/// a struct literal; pattern matching must use the `..` rest pattern.
/// In-crate construction is unaffected.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub struct Usage {
    /// Prompt tokens. For providers with prompt caching (e.g. Anthropic) this
    /// is the **uncached** prompt count; cached tokens are reported separately
    /// in [`Usage::cache_read_tokens`] / [`Usage::cache_creation_tokens`].
    pub prompt_tokens: u32,
    /// Completion tokens. Includes reasoning/thinking tokens where the provider
    /// bills them as output (e.g. Gemini `thoughtsTokenCount`).
    pub completion_tokens: u32,
    /// Prompt tokens served from the provider's cache. Billed at a reduced tier
    /// (Anthropic ~0.1x the input rate); 0 when caching is absent or unused.
    /// `#[serde(default)]` keeps pre-cache `Usage` JSON deserializable.
    #[serde(default)]
    pub cache_read_tokens: u32,
    /// Prompt tokens written to the provider's cache. Billed at a premium tier
    /// (Anthropic ~1.25x the input rate); 0 when caching is absent or unused.
    #[serde(default)]
    pub cache_creation_tokens: u32,
}

impl Usage {
    /// Prefer this over struct literals outside `klieo-core` —
    /// `#[non_exhaustive]` forbids external struct-literal construction so that
    /// future field additions remain additive. Cache token counts default to 0;
    /// set them with [`Usage::with_cache_tokens`].
    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
        Self {
            prompt_tokens,
            completion_tokens,
            cache_read_tokens: 0,
            cache_creation_tokens: 0,
        }
    }

    /// Attach prompt-cache token counts (read = served from cache, creation =
    /// written to cache). Both bill at provider-specific tiers distinct from the
    /// base prompt rate, so they are tracked apart for correct cost pricing.
    #[must_use]
    pub fn with_cache_tokens(mut self, cache_read_tokens: u32, cache_creation_tokens: u32) -> Self {
        self.cache_read_tokens = cache_read_tokens;
        self.cache_creation_tokens = cache_creation_tokens;
        self
    }
}

/// Reason the provider stopped generating.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum FinishReason {
    /// Model emitted a stop token or end-of-turn.
    Stop,
    /// Model emitted tool calls; runtime must dispatch them.
    ToolCalls,
    /// Hit `max_tokens`.
    Length,
    /// Provider applied a content filter.
    ContentFilter,
    /// Provider error mid-stream.
    Error,
}

/// Provider capability declaration.
#[derive(Debug, Clone, Default)]
pub struct Capabilities {
    /// Supports tool calls.
    pub tool_calling: bool,
    /// Supports streaming responses.
    pub streaming: bool,
    /// Supports schema-validated structured output.
    pub structured_output: bool,
    /// Supports embeddings.
    pub embeddings: bool,
    /// Maximum context window in tokens.
    pub max_context_tokens: u32,
    /// Supports vision input. (Not used in foundation MVP.)
    pub vision: bool,
}

impl Capabilities {
    /// Start a fluent builder. Equivalent to `Capabilities::default()`
    /// followed by chained setters.
    ///
    /// ```
    /// use klieo_core::Capabilities;
    /// let caps = Capabilities::builder()
    ///     .tool_calling(true)
    ///     .streaming(true)
    ///     .max_context_tokens(8000)
    ///     .build();
    /// assert!(caps.tool_calling);
    /// assert!(caps.streaming);
    /// assert_eq!(caps.max_context_tokens, 8000);
    /// ```
    pub fn builder() -> CapabilitiesBuilder {
        CapabilitiesBuilder(Capabilities::default())
    }
}

/// Fluent builder for [`Capabilities`]. Build via [`Capabilities::builder`].
#[derive(Debug, Clone, Default)]
pub struct CapabilitiesBuilder(Capabilities);

impl CapabilitiesBuilder {
    /// Set the `tool_calling` flag.
    pub fn tool_calling(mut self, v: bool) -> Self {
        self.0.tool_calling = v;
        self
    }
    /// Set the `streaming` flag.
    pub fn streaming(mut self, v: bool) -> Self {
        self.0.streaming = v;
        self
    }
    /// Set the `structured_output` flag.
    pub fn structured_output(mut self, v: bool) -> Self {
        self.0.structured_output = v;
        self
    }
    /// Set the `embeddings` flag.
    pub fn embeddings(mut self, v: bool) -> Self {
        self.0.embeddings = v;
        self
    }
    /// Set the maximum context window in tokens.
    pub fn max_context_tokens(mut self, v: u32) -> Self {
        self.0.max_context_tokens = v;
        self
    }
    /// Set the `vision` flag.
    pub fn vision(mut self, v: bool) -> Self {
        self.0.vision = v;
        self
    }
    /// Consume the builder and return the configured [`Capabilities`].
    pub fn build(self) -> Capabilities {
        self.0
    }
}

/// One chunk of a streaming response.
///
/// Marked `#[non_exhaustive]` — future field additions are
/// additive (no SemVer-major bump required). Code outside `klieo-core`
/// that constructs a `ChatChunk` must use [`ChatChunk::new`] instead of
/// a struct literal; pattern matching must use the `..` rest pattern.
/// In-crate construction is unaffected.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct ChatChunk {
    /// Incremental content delta.
    pub delta: String,
    /// Tool calls emitted in this chunk (rare; usually only at end).
    pub tool_calls: Vec<ToolCall>,
    /// `Some` once the provider signals completion.
    pub finish_reason: Option<FinishReason>,
    /// Token usage. `Some` only on the final chunk for providers
    /// that surface usage in stream mode; `None` otherwise.
    /// Consumers should treat absence as "unknown" rather than zero.
    pub usage: Option<Usage>,
}

impl ChatChunk {
    /// Construct a [`ChatChunk`] from its constituent fields.
    ///
    /// Prefer this constructor over struct literals in code outside
    /// `klieo-core` — `#[non_exhaustive]` forbids external struct-literal
    /// construction so that future field additions remain additive.
    pub fn new(
        delta: String,
        tool_calls: Vec<ToolCall>,
        finish_reason: Option<FinishReason>,
        usage: Option<Usage>,
    ) -> Self {
        Self {
            delta,
            tool_calls,
            finish_reason,
            usage,
        }
    }
}

/// Streaming response handle.
pub type ChunkStream = Pin<Box<dyn Stream<Item = Result<ChatChunk, LlmError>> + Send + 'static>>;

/// Vector embedding for one input text.
pub type Embedding = Vec<f32>;

/// LLM provider trait.
///
/// Implementors live in their own crates (`klieo-llm-ollama`, etc.).
/// `Capabilities` are inspected by the runtime before it sends an
/// unsupported request.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep};
/// use klieo_core::{ChatRequest, FinishReason, LlmClient};
/// let llm = FakeLlmClient::new("fake")
///     .with_steps(vec![FakeLlmStep::Text("hello".into())]);
/// assert_eq!(llm.name(), "fake");
/// assert!(llm.capabilities().tool_calling);
/// let resp = llm.complete(ChatRequest::new(vec![])).await.unwrap();
/// assert_eq!(resp.message.content, "hello");
/// assert_eq!(resp.finish_reason, FinishReason::Stop);
/// # });
/// ```
#[async_trait]
pub trait LlmClient: Send + Sync {
    /// Stable identifier for this client (e.g. `"ollama:qwen2.5:14b"`).
    fn name(&self) -> &str;

    /// Capabilities declared by this client.
    fn capabilities(&self) -> &Capabilities;

    /// One-shot completion.
    ///
    /// ```
    /// # tokio_test::block_on(async {
    /// use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep};
    /// use klieo_core::{ChatRequest, FinishReason, LlmClient};
    /// let llm = FakeLlmClient::new("fake")
    ///     .with_steps(vec![FakeLlmStep::Text("hello".into())]);
    /// let resp = llm.complete(ChatRequest::new(vec![])).await.unwrap();
    /// assert_eq!(resp.message.content, "hello");
    /// assert_eq!(resp.finish_reason, FinishReason::Stop);
    /// # });
    /// ```
    async fn complete(&self, req: ChatRequest) -> Result<ChatResponse, LlmError>;

    /// Streaming completion.
    ///
    /// ```
    /// # tokio_test::block_on(async {
    /// use klieo_core::test_utils::FakeLlmClient;
    /// use klieo_core::{ChatRequest, LlmClient, LlmError};
    /// let llm = FakeLlmClient::new("fake");
    /// match llm.stream(ChatRequest::new(vec![])).await {
    ///     Ok(_) => panic!("expected Unsupported"),
    ///     Err(e) => assert!(matches!(e, LlmError::Unsupported(_))),
    /// }
    /// # });
    /// ```
    async fn stream(&self, req: ChatRequest) -> Result<ChunkStream, LlmError>;

    /// Compute embeddings for the supplied texts.
    ///
    /// ```
    /// # tokio_test::block_on(async {
    /// use klieo_core::test_utils::FakeLlmClient;
    /// use klieo_core::{LlmClient, LlmError};
    /// let llm = FakeLlmClient::new("fake");
    /// let err = llm.embed(&["hello".into()]).await.unwrap_err();
    /// assert!(matches!(err, LlmError::Unsupported(_)));
    /// # });
    /// ```
    async fn embed(&self, texts: &[String]) -> Result<Vec<Embedding>, LlmError>;
}

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

    /// Compile-time check that LlmClient is dyn-compatible.
    #[allow(dead_code)]
    fn _assert_dyn_compatible(_: &dyn LlmClient) {}

    #[test]
    fn chat_request_default_has_no_tools() {
        let req = ChatRequest::new(vec![]);
        assert!(req.tools.is_empty());
        assert!(matches!(req.response_format, ResponseFormat::Text));
    }

    #[test]
    fn message_constructors_set_role_and_clear_tool_fields() {
        let m = Message::system("be terse");
        assert_eq!(m.role, Role::System);
        assert_eq!(m.content, "be terse");
        assert!(m.tool_calls.is_empty());
        assert!(m.tool_call_id.is_none());
        let user = Message::user("hi");
        assert_eq!((user.role, user.content.as_str()), (Role::User, "hi"));
        let assistant = Message::assistant("ok");
        assert_eq!(
            (assistant.role, assistant.content.as_str()),
            (Role::Assistant, "ok")
        );
    }

    #[test]
    fn chat_request_builder_sets_every_option_and_appends_raw_messages() {
        let timeout = Duration::from_secs(7);
        let schema = serde_json::json!({"type": "object"});
        let req = ChatRequest::builder()
            .assistant("a")
            .message(Message::user("raw"))
            .tools(vec![ToolDef::new("t", "desc", serde_json::json!({}))])
            .max_tokens(42)
            .response_format(ResponseFormat::Json {
                schema: schema.clone(),
            })
            .timeout(timeout)
            .build();
        let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
        assert_eq!(roles, vec![Role::Assistant, Role::User]);
        assert_eq!(req.messages[1].content, "raw");
        assert_eq!(req.tools.len(), 1);
        assert_eq!(req.max_tokens, Some(42));
        assert_eq!(req.timeout, Some(timeout));
        assert!(matches!(req.response_format, ResponseFormat::Json { .. }));
    }

    #[test]
    fn chat_request_builder_orders_messages_and_keeps_defaults() {
        let req = ChatRequest::builder()
            .system("sys")
            .user("u")
            .temperature(0.5)
            .build();
        let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
        assert_eq!(roles, vec![Role::System, Role::User]);
        assert_eq!(req.temperature, Some(0.5));
        assert!(req.tools.is_empty());
        assert!(matches!(req.response_format, ResponseFormat::Text));
    }

    #[test]
    fn capabilities_builder_default_matches_struct_default() {
        let built = Capabilities::builder().build();
        let direct = Capabilities::default();
        assert_eq!(built.tool_calling, direct.tool_calling);
        assert_eq!(built.streaming, direct.streaming);
        assert_eq!(built.structured_output, direct.structured_output);
        assert_eq!(built.embeddings, direct.embeddings);
        assert_eq!(built.max_context_tokens, direct.max_context_tokens);
        assert_eq!(built.vision, direct.vision);
    }

    #[test]
    fn capabilities_builder_sets_tool_calling() {
        let c = Capabilities::builder().tool_calling(true).build();
        assert!(c.tool_calling);
    }

    #[test]
    fn capabilities_builder_sets_streaming() {
        let c = Capabilities::builder().streaming(true).build();
        assert!(c.streaming);
    }

    #[test]
    fn capabilities_builder_sets_structured_output() {
        let c = Capabilities::builder().structured_output(true).build();
        assert!(c.structured_output);
    }

    #[test]
    fn capabilities_builder_sets_embeddings() {
        let c = Capabilities::builder().embeddings(true).build();
        assert!(c.embeddings);
    }

    #[test]
    fn capabilities_builder_sets_max_context_tokens() {
        let c = Capabilities::builder().max_context_tokens(8000).build();
        assert_eq!(c.max_context_tokens, 8000);
    }

    #[test]
    fn capabilities_builder_sets_vision() {
        let c = Capabilities::builder().vision(true).build();
        assert!(c.vision);
    }

    #[test]
    fn capabilities_builder_chains_all_setters() {
        let c = Capabilities::builder()
            .tool_calling(true)
            .streaming(true)
            .structured_output(true)
            .embeddings(true)
            .max_context_tokens(32_000)
            .vision(false)
            .build();
        assert!(c.tool_calling && c.streaming && c.structured_output && c.embeddings);
        assert_eq!(c.max_context_tokens, 32_000);
        assert!(!c.vision);
    }

    #[test]
    fn chat_chunk_usage_defaults_to_none_in_struct_literal() {
        let chunk = ChatChunk {
            delta: String::new(),
            tool_calls: vec![],
            finish_reason: None,
            usage: None,
        };
        assert!(chunk.usage.is_none());
    }

    #[test]
    fn chat_chunk_with_usage_round_trips() {
        let chunk = ChatChunk {
            delta: "done".into(),
            tool_calls: vec![],
            finish_reason: Some(FinishReason::Stop),
            usage: Some(Usage::new(10, 32)),
        };
        let u = chunk.usage.as_ref().expect("usage set");
        assert_eq!(u.prompt_tokens, 10);
        assert_eq!(u.completion_tokens, 32);
    }

    #[test]
    fn usage_new_sets_token_counts() {
        let u = Usage::new(12, 34);
        assert_eq!(u.prompt_tokens, 12);
        assert_eq!(u.completion_tokens, 34);
    }

    #[test]
    fn usage_cache_tokens_default_to_zero() {
        let u = Usage::new(12, 34);
        assert_eq!(u.cache_read_tokens, 0);
        assert_eq!(u.cache_creation_tokens, 0);
    }

    #[test]
    fn usage_deserializes_legacy_json_without_cache_fields() {
        // Pre-cache `Usage` JSON lacks the cache fields; `#[serde(default)]`
        // must keep it deserializable (cache counts read back as 0).
        let legacy = r#"{"prompt_tokens":10,"completion_tokens":5}"#;
        let u: Usage = serde_json::from_str(legacy).unwrap();
        assert_eq!(u.cache_read_tokens, 0);
        assert_eq!(u.cache_creation_tokens, 0);
    }

    #[test]
    fn usage_with_cache_tokens_sets_cache_fields() {
        // Cache tokens are billed at provider-specific tiers (Anthropic reads
        // ~0.1x input, creation ~1.25x), so they are tracked apart from the
        // base prompt/completion counts for correct cost pricing.
        let u = Usage::new(12, 34).with_cache_tokens(100, 8);
        assert_eq!(u.prompt_tokens, 12);
        assert_eq!(u.completion_tokens, 34);
        assert_eq!(u.cache_read_tokens, 100);
        assert_eq!(u.cache_creation_tokens, 8);
    }

    #[test]
    fn tooldef_new_sets_fields() {
        let d = ToolDef::new("echo", "echoes input", serde_json::json!({"type":"object"}));
        assert_eq!(d.name, "echo");
        assert_eq!(d.description, "echoes input");
        assert_eq!(d.json_schema, serde_json::json!({"type":"object"}));
    }

    #[test]
    fn toolcall_new_sets_fields() {
        let c = ToolCall::new("id-1", "search", serde_json::json!({"q":"x"}));
        assert_eq!(c.id, "id-1");
        assert_eq!(c.name, "search");
        assert_eq!(c.args, serde_json::json!({"q":"x"}));
    }

    #[test]
    fn chatresponse_new_carries_model() {
        let msg = Message {
            role: Role::Assistant,
            content: "hi".into(),
            tool_calls: vec![],
            tool_call_id: None,
        };
        let r = ChatResponse::new(
            msg,
            Usage::new(1, 2),
            FinishReason::Stop,
            "anthropic:claude-sonnet-4-6",
        );
        assert_eq!(r.model, "anthropic:claude-sonnet-4-6");
        assert_eq!(r.finish_reason, FinishReason::Stop);
        assert_eq!(r.message.content, "hi");
        assert_eq!(r.usage.prompt_tokens, 1);
    }
}