klieo_core/
llm.rs

1//! LLM client trait and request/response types.
2//!
3//! Providers implement [`LlmClient`]. The runtime never depends on a
4//! specific provider crate — selection is explicit at app startup.
5
6use crate::error::LlmError;
7use async_trait::async_trait;
8use futures_core::Stream;
9use serde::{Deserialize, Serialize};
10use std::pin::Pin;
11use std::time::Duration;
12
13/// One message in a chat conversation.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Message {
16    /// Speaker role.
17    pub role: Role,
18    /// Body text. May be empty when `tool_calls` carries the payload.
19    pub content: String,
20    /// Tool calls the assistant requested in this message.
21    #[serde(default)]
22    pub tool_calls: Vec<ToolCall>,
23    /// When this message answers a tool call, the id of that call.
24    #[serde(default)]
25    pub tool_call_id: Option<String>,
26}
27
28impl Message {
29    /// [`Role::System`] message; `tool_calls` is empty and `tool_call_id` is `None`.
30    pub fn system(content: impl Into<String>) -> Self {
31        Self::text(Role::System, content)
32    }
33
34    /// [`Role::User`] message; `tool_calls` is empty and `tool_call_id` is `None`.
35    pub fn user(content: impl Into<String>) -> Self {
36        Self::text(Role::User, content)
37    }
38
39    /// [`Role::Assistant`] message; `tool_calls` is empty and `tool_call_id` is `None`.
40    pub fn assistant(content: impl Into<String>) -> Self {
41        Self::text(Role::Assistant, content)
42    }
43
44    fn text(role: Role, content: impl Into<String>) -> Self {
45        Self {
46            role,
47            content: content.into(),
48            tool_calls: Vec::new(),
49            tool_call_id: None,
50        }
51    }
52}
53
54/// Speaker role.
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
56#[serde(rename_all = "lowercase")]
57#[non_exhaustive]
58pub enum Role {
59    /// System prompt.
60    System,
61    /// User-supplied input.
62    User,
63    /// Model output.
64    Assistant,
65    /// Tool call result fed back into the conversation.
66    Tool,
67}
68
69/// One tool call requested by the assistant.
70///
71/// Marked `#[non_exhaustive]` — future field additions are
72/// additive (no SemVer-major bump required). Code outside `klieo-core`
73/// that constructs a `ToolCall` must use [`ToolCall::new`] instead of
74/// a struct literal; pattern matching must use the `..` rest pattern.
75/// In-crate construction is unaffected.
76#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
77#[non_exhaustive]
78pub struct ToolCall {
79    /// Provider-stable id; echoed back in the matching tool response.
80    pub id: String,
81    /// Tool name.
82    pub name: String,
83    /// JSON arguments.
84    pub args: serde_json::Value,
85}
86
87impl ToolCall {
88    /// Prefer this over struct literals outside `klieo-core` —
89    /// `#[non_exhaustive]` forbids external struct-literal construction so that
90    /// future field additions remain additive.
91    pub fn new(id: impl Into<String>, name: impl Into<String>, args: serde_json::Value) -> Self {
92        Self {
93            id: id.into(),
94            name: name.into(),
95            args,
96        }
97    }
98}
99
100/// Tool catalogue entry shown to the LLM.
101#[derive(Debug, Clone, Serialize, Deserialize)]
102#[non_exhaustive]
103pub struct ToolDef {
104    /// Tool name (must be unique within the catalogue).
105    pub name: String,
106    /// Human-readable description for the LLM.
107    pub description: String,
108    /// JSON-schema for arguments.
109    pub json_schema: serde_json::Value,
110}
111
112impl ToolDef {
113    /// One entry in the tool catalogue offered to the model.
114    ///
115    /// Prefer this constructor over struct literals in code outside
116    /// `klieo-core` — `#[non_exhaustive]` forbids external struct-literal
117    /// construction so that future field additions remain additive.
118    pub fn new(
119        name: impl Into<String>,
120        description: impl Into<String>,
121        json_schema: serde_json::Value,
122    ) -> Self {
123        Self {
124            name: name.into(),
125            description: description.into(),
126            json_schema,
127        }
128    }
129}
130
131/// Chat completion request.
132#[derive(Debug, Clone)]
133pub struct ChatRequest {
134    /// Conversation history.
135    pub messages: Vec<Message>,
136    /// Tools available for the model to call.
137    pub tools: Vec<ToolDef>,
138    /// Sampling temperature.
139    pub temperature: Option<f32>,
140    /// Maximum response tokens.
141    pub max_tokens: Option<u32>,
142    /// Output format hint.
143    pub response_format: ResponseFormat,
144    /// Stop sequences.
145    pub stop: Vec<String>,
146    /// Per-request deadline. Provider should abort if exceeded.
147    pub timeout: Option<Duration>,
148}
149
150impl ChatRequest {
151    /// Build a request with the supplied messages and no tools.
152    pub fn new(messages: Vec<Message>) -> Self {
153        Self {
154            messages,
155            tools: Vec::new(),
156            temperature: None,
157            max_tokens: None,
158            response_format: ResponseFormat::Text,
159            stop: Vec::new(),
160            timeout: None,
161        }
162    }
163
164    /// Start building a request; messages are sent in the order they are added.
165    ///
166    /// ```
167    /// use klieo_core::llm::ChatRequest;
168    /// let req = ChatRequest::builder().system("be terse").user("hi").build();
169    /// assert_eq!(req.messages.len(), 2);
170    /// ```
171    pub fn builder() -> ChatRequestBuilder {
172        ChatRequestBuilder::default()
173    }
174}
175
176/// Builder for [`ChatRequest`]: messages accumulate in call order, and any
177/// option left unset falls back to the [`ChatRequest::new`] default on `build`.
178#[derive(Default)]
179pub struct ChatRequestBuilder {
180    messages: Vec<Message>,
181    tools: Vec<ToolDef>,
182    temperature: Option<f32>,
183    max_tokens: Option<u32>,
184    response_format: Option<ResponseFormat>,
185    timeout: Option<Duration>,
186}
187
188impl ChatRequestBuilder {
189    /// Appends a [`Role::System`] message after any already added.
190    pub fn system(mut self, content: impl Into<String>) -> Self {
191        self.messages.push(Message::system(content));
192        self
193    }
194
195    /// Appends a [`Role::User`] message after any already added.
196    pub fn user(mut self, content: impl Into<String>) -> Self {
197        self.messages.push(Message::user(content));
198        self
199    }
200
201    /// Appends a [`Role::Assistant`] message after any already added.
202    pub fn assistant(mut self, content: impl Into<String>) -> Self {
203        self.messages.push(Message::assistant(content));
204        self
205    }
206
207    /// Appends a pre-built message — the only way to add a tool-result turn.
208    pub fn message(mut self, message: Message) -> Self {
209        self.messages.push(message);
210        self
211    }
212
213    /// Replaces the tool catalogue; empty (no tools) when never called.
214    pub fn tools(mut self, tools: Vec<ToolDef>) -> Self {
215        self.tools = tools;
216        self
217    }
218
219    /// Sampling temperature; left to the provider default when never called.
220    pub fn temperature(mut self, temperature: f32) -> Self {
221        self.temperature = Some(temperature);
222        self
223    }
224
225    /// Response-token cap; provider default when never called.
226    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
227        self.max_tokens = Some(max_tokens);
228        self
229    }
230
231    /// Output format; falls back to [`ResponseFormat::Text`] when never called.
232    pub fn response_format(mut self, response_format: ResponseFormat) -> Self {
233        self.response_format = Some(response_format);
234        self
235    }
236
237    /// Per-request deadline; none (provider transport governs) when never called.
238    pub fn timeout(mut self, timeout: Duration) -> Self {
239        self.timeout = Some(timeout);
240        self
241    }
242
243    /// Builds the request; the `stop` list is always empty and other unset
244    /// options take their [`ChatRequest::new`] defaults.
245    pub fn build(self) -> ChatRequest {
246        ChatRequest {
247            messages: self.messages,
248            tools: self.tools,
249            temperature: self.temperature,
250            max_tokens: self.max_tokens,
251            response_format: self.response_format.unwrap_or(ResponseFormat::Text),
252            stop: Vec::new(),
253            timeout: self.timeout,
254        }
255    }
256}
257
258/// Output format hint passed to the provider.
259#[derive(Debug, Clone)]
260#[non_exhaustive]
261pub enum ResponseFormat {
262    /// Plain text.
263    Text,
264    /// Best-effort JSON output.
265    Json {
266        /// Schema describing expected JSON shape.
267        schema: serde_json::Value,
268    },
269    /// Strict structured output validated against the schema.
270    StructuredOutput {
271        /// Schema describing expected shape.
272        schema: serde_json::Value,
273    },
274}
275
276/// Chat completion response.
277///
278/// Marked `#[non_exhaustive]` — future field additions are additive.
279/// All construction outside `klieo-core` must use [`ChatResponse::new`].
280#[derive(Debug, Clone)]
281#[non_exhaustive]
282pub struct ChatResponse {
283    /// Assistant message returned by the provider.
284    pub message: Message,
285    /// Token usage.
286    pub usage: Usage,
287    /// Why the provider stopped generating.
288    pub finish_reason: FinishReason,
289    /// Model that served this response — provider-reported when available
290    /// (Ollama), else the client-configured model the request used. Recorded
291    /// into `Capture.model_version` for reproducibility; best-effort metadata,
292    /// not an authorization or correctness input.
293    pub model: String,
294}
295
296impl ChatResponse {
297    /// Provider response for one completion.
298    ///
299    /// `model` is the provider-reported model name when the response body
300    /// carries one (e.g. Ollama), or the configured model used for the
301    /// request otherwise.
302    pub fn new(
303        message: Message,
304        usage: Usage,
305        finish_reason: FinishReason,
306        model: impl Into<String>,
307    ) -> Self {
308        Self {
309            message,
310            usage,
311            finish_reason,
312            model: model.into(),
313        }
314    }
315}
316
317/// Token usage report.
318///
319/// Marked `#[non_exhaustive]` — future field additions are
320/// additive (no SemVer-major bump required). Code outside `klieo-core`
321/// that constructs a `Usage` must use [`Usage::new`] instead of
322/// a struct literal; pattern matching must use the `..` rest pattern.
323/// In-crate construction is unaffected.
324#[derive(Debug, Clone, Default, Serialize, Deserialize)]
325#[non_exhaustive]
326pub struct Usage {
327    /// Prompt tokens. For providers with prompt caching (e.g. Anthropic) this
328    /// is the **uncached** prompt count; cached tokens are reported separately
329    /// in [`Usage::cache_read_tokens`] / [`Usage::cache_creation_tokens`].
330    pub prompt_tokens: u32,
331    /// Completion tokens. Includes reasoning/thinking tokens where the provider
332    /// bills them as output (e.g. Gemini `thoughtsTokenCount`).
333    pub completion_tokens: u32,
334    /// Prompt tokens served from the provider's cache. Billed at a reduced tier
335    /// (Anthropic ~0.1x the input rate); 0 when caching is absent or unused.
336    /// `#[serde(default)]` keeps pre-cache `Usage` JSON deserializable.
337    #[serde(default)]
338    pub cache_read_tokens: u32,
339    /// Prompt tokens written to the provider's cache. Billed at a premium tier
340    /// (Anthropic ~1.25x the input rate); 0 when caching is absent or unused.
341    #[serde(default)]
342    pub cache_creation_tokens: u32,
343}
344
345impl Usage {
346    /// Prefer this over struct literals outside `klieo-core` —
347    /// `#[non_exhaustive]` forbids external struct-literal construction so that
348    /// future field additions remain additive. Cache token counts default to 0;
349    /// set them with [`Usage::with_cache_tokens`].
350    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
351        Self {
352            prompt_tokens,
353            completion_tokens,
354            cache_read_tokens: 0,
355            cache_creation_tokens: 0,
356        }
357    }
358
359    /// Attach prompt-cache token counts (read = served from cache, creation =
360    /// written to cache). Both bill at provider-specific tiers distinct from the
361    /// base prompt rate, so they are tracked apart for correct cost pricing.
362    #[must_use]
363    pub fn with_cache_tokens(mut self, cache_read_tokens: u32, cache_creation_tokens: u32) -> Self {
364        self.cache_read_tokens = cache_read_tokens;
365        self.cache_creation_tokens = cache_creation_tokens;
366        self
367    }
368}
369
370/// Reason the provider stopped generating.
371#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
372#[non_exhaustive]
373pub enum FinishReason {
374    /// Model emitted a stop token or end-of-turn.
375    Stop,
376    /// Model emitted tool calls; runtime must dispatch them.
377    ToolCalls,
378    /// Hit `max_tokens`.
379    Length,
380    /// Provider applied a content filter.
381    ContentFilter,
382    /// Provider error mid-stream.
383    Error,
384}
385
386/// Provider capability declaration.
387#[derive(Debug, Clone, Default)]
388pub struct Capabilities {
389    /// Supports tool calls.
390    pub tool_calling: bool,
391    /// Supports streaming responses.
392    pub streaming: bool,
393    /// Supports schema-validated structured output.
394    pub structured_output: bool,
395    /// Supports embeddings.
396    pub embeddings: bool,
397    /// Maximum context window in tokens.
398    pub max_context_tokens: u32,
399    /// Supports vision input. (Not used in foundation MVP.)
400    pub vision: bool,
401}
402
403impl Capabilities {
404    /// Start a fluent builder. Equivalent to `Capabilities::default()`
405    /// followed by chained setters.
406    ///
407    /// ```
408    /// use klieo_core::Capabilities;
409    /// let caps = Capabilities::builder()
410    ///     .tool_calling(true)
411    ///     .streaming(true)
412    ///     .max_context_tokens(8000)
413    ///     .build();
414    /// assert!(caps.tool_calling);
415    /// assert!(caps.streaming);
416    /// assert_eq!(caps.max_context_tokens, 8000);
417    /// ```
418    pub fn builder() -> CapabilitiesBuilder {
419        CapabilitiesBuilder(Capabilities::default())
420    }
421}
422
423/// Fluent builder for [`Capabilities`]. Build via [`Capabilities::builder`].
424#[derive(Debug, Clone, Default)]
425pub struct CapabilitiesBuilder(Capabilities);
426
427impl CapabilitiesBuilder {
428    /// Set the `tool_calling` flag.
429    pub fn tool_calling(mut self, v: bool) -> Self {
430        self.0.tool_calling = v;
431        self
432    }
433    /// Set the `streaming` flag.
434    pub fn streaming(mut self, v: bool) -> Self {
435        self.0.streaming = v;
436        self
437    }
438    /// Set the `structured_output` flag.
439    pub fn structured_output(mut self, v: bool) -> Self {
440        self.0.structured_output = v;
441        self
442    }
443    /// Set the `embeddings` flag.
444    pub fn embeddings(mut self, v: bool) -> Self {
445        self.0.embeddings = v;
446        self
447    }
448    /// Set the maximum context window in tokens.
449    pub fn max_context_tokens(mut self, v: u32) -> Self {
450        self.0.max_context_tokens = v;
451        self
452    }
453    /// Set the `vision` flag.
454    pub fn vision(mut self, v: bool) -> Self {
455        self.0.vision = v;
456        self
457    }
458    /// Consume the builder and return the configured [`Capabilities`].
459    pub fn build(self) -> Capabilities {
460        self.0
461    }
462}
463
464/// One chunk of a streaming response.
465///
466/// Marked `#[non_exhaustive]` — future field additions are
467/// additive (no SemVer-major bump required). Code outside `klieo-core`
468/// that constructs a `ChatChunk` must use [`ChatChunk::new`] instead of
469/// a struct literal; pattern matching must use the `..` rest pattern.
470/// In-crate construction is unaffected.
471#[derive(Debug, Clone, Default)]
472#[non_exhaustive]
473pub struct ChatChunk {
474    /// Incremental content delta.
475    pub delta: String,
476    /// Tool calls emitted in this chunk (rare; usually only at end).
477    pub tool_calls: Vec<ToolCall>,
478    /// `Some` once the provider signals completion.
479    pub finish_reason: Option<FinishReason>,
480    /// Token usage. `Some` only on the final chunk for providers
481    /// that surface usage in stream mode; `None` otherwise.
482    /// Consumers should treat absence as "unknown" rather than zero.
483    pub usage: Option<Usage>,
484}
485
486impl ChatChunk {
487    /// Construct a [`ChatChunk`] from its constituent fields.
488    ///
489    /// Prefer this constructor over struct literals in code outside
490    /// `klieo-core` — `#[non_exhaustive]` forbids external struct-literal
491    /// construction so that future field additions remain additive.
492    pub fn new(
493        delta: String,
494        tool_calls: Vec<ToolCall>,
495        finish_reason: Option<FinishReason>,
496        usage: Option<Usage>,
497    ) -> Self {
498        Self {
499            delta,
500            tool_calls,
501            finish_reason,
502            usage,
503        }
504    }
505}
506
507/// Streaming response handle.
508pub type ChunkStream = Pin<Box<dyn Stream<Item = Result<ChatChunk, LlmError>> + Send + 'static>>;
509
510/// Vector embedding for one input text.
511pub type Embedding = Vec<f32>;
512
513/// LLM provider trait.
514///
515/// Implementors live in their own crates (`klieo-llm-ollama`, etc.).
516/// `Capabilities` are inspected by the runtime before it sends an
517/// unsupported request.
518///
519/// ```
520/// # tokio_test::block_on(async {
521/// use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep};
522/// use klieo_core::{ChatRequest, FinishReason, LlmClient};
523/// let llm = FakeLlmClient::new("fake")
524///     .with_steps(vec![FakeLlmStep::Text("hello".into())]);
525/// assert_eq!(llm.name(), "fake");
526/// assert!(llm.capabilities().tool_calling);
527/// let resp = llm.complete(ChatRequest::new(vec![])).await.unwrap();
528/// assert_eq!(resp.message.content, "hello");
529/// assert_eq!(resp.finish_reason, FinishReason::Stop);
530/// # });
531/// ```
532#[async_trait]
533pub trait LlmClient: Send + Sync {
534    /// Stable identifier for this client (e.g. `"ollama:qwen2.5:14b"`).
535    fn name(&self) -> &str;
536
537    /// Capabilities declared by this client.
538    fn capabilities(&self) -> &Capabilities;
539
540    /// One-shot completion.
541    ///
542    /// ```
543    /// # tokio_test::block_on(async {
544    /// use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep};
545    /// use klieo_core::{ChatRequest, FinishReason, LlmClient};
546    /// let llm = FakeLlmClient::new("fake")
547    ///     .with_steps(vec![FakeLlmStep::Text("hello".into())]);
548    /// let resp = llm.complete(ChatRequest::new(vec![])).await.unwrap();
549    /// assert_eq!(resp.message.content, "hello");
550    /// assert_eq!(resp.finish_reason, FinishReason::Stop);
551    /// # });
552    /// ```
553    async fn complete(&self, req: ChatRequest) -> Result<ChatResponse, LlmError>;
554
555    /// Streaming completion.
556    ///
557    /// ```
558    /// # tokio_test::block_on(async {
559    /// use klieo_core::test_utils::FakeLlmClient;
560    /// use klieo_core::{ChatRequest, LlmClient, LlmError};
561    /// let llm = FakeLlmClient::new("fake");
562    /// match llm.stream(ChatRequest::new(vec![])).await {
563    ///     Ok(_) => panic!("expected Unsupported"),
564    ///     Err(e) => assert!(matches!(e, LlmError::Unsupported(_))),
565    /// }
566    /// # });
567    /// ```
568    async fn stream(&self, req: ChatRequest) -> Result<ChunkStream, LlmError>;
569
570    /// Compute embeddings for the supplied texts.
571    ///
572    /// ```
573    /// # tokio_test::block_on(async {
574    /// use klieo_core::test_utils::FakeLlmClient;
575    /// use klieo_core::{LlmClient, LlmError};
576    /// let llm = FakeLlmClient::new("fake");
577    /// let err = llm.embed(&["hello".into()]).await.unwrap_err();
578    /// assert!(matches!(err, LlmError::Unsupported(_)));
579    /// # });
580    /// ```
581    async fn embed(&self, texts: &[String]) -> Result<Vec<Embedding>, LlmError>;
582}
583
584#[cfg(test)]
585mod tests {
586    use super::*;
587
588    /// Compile-time check that LlmClient is dyn-compatible.
589    #[allow(dead_code)]
590    fn _assert_dyn_compatible(_: &dyn LlmClient) {}
591
592    #[test]
593    fn chat_request_default_has_no_tools() {
594        let req = ChatRequest::new(vec![]);
595        assert!(req.tools.is_empty());
596        assert!(matches!(req.response_format, ResponseFormat::Text));
597    }
598
599    #[test]
600    fn message_constructors_set_role_and_clear_tool_fields() {
601        let m = Message::system("be terse");
602        assert_eq!(m.role, Role::System);
603        assert_eq!(m.content, "be terse");
604        assert!(m.tool_calls.is_empty());
605        assert!(m.tool_call_id.is_none());
606        let user = Message::user("hi");
607        assert_eq!((user.role, user.content.as_str()), (Role::User, "hi"));
608        let assistant = Message::assistant("ok");
609        assert_eq!(
610            (assistant.role, assistant.content.as_str()),
611            (Role::Assistant, "ok")
612        );
613    }
614
615    #[test]
616    fn chat_request_builder_sets_every_option_and_appends_raw_messages() {
617        let timeout = Duration::from_secs(7);
618        let schema = serde_json::json!({"type": "object"});
619        let req = ChatRequest::builder()
620            .assistant("a")
621            .message(Message::user("raw"))
622            .tools(vec![ToolDef::new("t", "desc", serde_json::json!({}))])
623            .max_tokens(42)
624            .response_format(ResponseFormat::Json {
625                schema: schema.clone(),
626            })
627            .timeout(timeout)
628            .build();
629        let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
630        assert_eq!(roles, vec![Role::Assistant, Role::User]);
631        assert_eq!(req.messages[1].content, "raw");
632        assert_eq!(req.tools.len(), 1);
633        assert_eq!(req.max_tokens, Some(42));
634        assert_eq!(req.timeout, Some(timeout));
635        assert!(matches!(req.response_format, ResponseFormat::Json { .. }));
636    }
637
638    #[test]
639    fn chat_request_builder_orders_messages_and_keeps_defaults() {
640        let req = ChatRequest::builder()
641            .system("sys")
642            .user("u")
643            .temperature(0.5)
644            .build();
645        let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
646        assert_eq!(roles, vec![Role::System, Role::User]);
647        assert_eq!(req.temperature, Some(0.5));
648        assert!(req.tools.is_empty());
649        assert!(matches!(req.response_format, ResponseFormat::Text));
650    }
651
652    #[test]
653    fn capabilities_builder_default_matches_struct_default() {
654        let built = Capabilities::builder().build();
655        let direct = Capabilities::default();
656        assert_eq!(built.tool_calling, direct.tool_calling);
657        assert_eq!(built.streaming, direct.streaming);
658        assert_eq!(built.structured_output, direct.structured_output);
659        assert_eq!(built.embeddings, direct.embeddings);
660        assert_eq!(built.max_context_tokens, direct.max_context_tokens);
661        assert_eq!(built.vision, direct.vision);
662    }
663
664    #[test]
665    fn capabilities_builder_sets_tool_calling() {
666        let c = Capabilities::builder().tool_calling(true).build();
667        assert!(c.tool_calling);
668    }
669
670    #[test]
671    fn capabilities_builder_sets_streaming() {
672        let c = Capabilities::builder().streaming(true).build();
673        assert!(c.streaming);
674    }
675
676    #[test]
677    fn capabilities_builder_sets_structured_output() {
678        let c = Capabilities::builder().structured_output(true).build();
679        assert!(c.structured_output);
680    }
681
682    #[test]
683    fn capabilities_builder_sets_embeddings() {
684        let c = Capabilities::builder().embeddings(true).build();
685        assert!(c.embeddings);
686    }
687
688    #[test]
689    fn capabilities_builder_sets_max_context_tokens() {
690        let c = Capabilities::builder().max_context_tokens(8000).build();
691        assert_eq!(c.max_context_tokens, 8000);
692    }
693
694    #[test]
695    fn capabilities_builder_sets_vision() {
696        let c = Capabilities::builder().vision(true).build();
697        assert!(c.vision);
698    }
699
700    #[test]
701    fn capabilities_builder_chains_all_setters() {
702        let c = Capabilities::builder()
703            .tool_calling(true)
704            .streaming(true)
705            .structured_output(true)
706            .embeddings(true)
707            .max_context_tokens(32_000)
708            .vision(false)
709            .build();
710        assert!(c.tool_calling && c.streaming && c.structured_output && c.embeddings);
711        assert_eq!(c.max_context_tokens, 32_000);
712        assert!(!c.vision);
713    }
714
715    #[test]
716    fn chat_chunk_usage_defaults_to_none_in_struct_literal() {
717        let chunk = ChatChunk {
718            delta: String::new(),
719            tool_calls: vec![],
720            finish_reason: None,
721            usage: None,
722        };
723        assert!(chunk.usage.is_none());
724    }
725
726    #[test]
727    fn chat_chunk_with_usage_round_trips() {
728        let chunk = ChatChunk {
729            delta: "done".into(),
730            tool_calls: vec![],
731            finish_reason: Some(FinishReason::Stop),
732            usage: Some(Usage::new(10, 32)),
733        };
734        let u = chunk.usage.as_ref().expect("usage set");
735        assert_eq!(u.prompt_tokens, 10);
736        assert_eq!(u.completion_tokens, 32);
737    }
738
739    #[test]
740    fn usage_new_sets_token_counts() {
741        let u = Usage::new(12, 34);
742        assert_eq!(u.prompt_tokens, 12);
743        assert_eq!(u.completion_tokens, 34);
744    }
745
746    #[test]
747    fn usage_cache_tokens_default_to_zero() {
748        let u = Usage::new(12, 34);
749        assert_eq!(u.cache_read_tokens, 0);
750        assert_eq!(u.cache_creation_tokens, 0);
751    }
752
753    #[test]
754    fn usage_deserializes_legacy_json_without_cache_fields() {
755        // Pre-cache `Usage` JSON lacks the cache fields; `#[serde(default)]`
756        // must keep it deserializable (cache counts read back as 0).
757        let legacy = r#"{"prompt_tokens":10,"completion_tokens":5}"#;
758        let u: Usage = serde_json::from_str(legacy).unwrap();
759        assert_eq!(u.cache_read_tokens, 0);
760        assert_eq!(u.cache_creation_tokens, 0);
761    }
762
763    #[test]
764    fn usage_with_cache_tokens_sets_cache_fields() {
765        // Cache tokens are billed at provider-specific tiers (Anthropic reads
766        // ~0.1x input, creation ~1.25x), so they are tracked apart from the
767        // base prompt/completion counts for correct cost pricing.
768        let u = Usage::new(12, 34).with_cache_tokens(100, 8);
769        assert_eq!(u.prompt_tokens, 12);
770        assert_eq!(u.completion_tokens, 34);
771        assert_eq!(u.cache_read_tokens, 100);
772        assert_eq!(u.cache_creation_tokens, 8);
773    }
774
775    #[test]
776    fn tooldef_new_sets_fields() {
777        let d = ToolDef::new("echo", "echoes input", serde_json::json!({"type":"object"}));
778        assert_eq!(d.name, "echo");
779        assert_eq!(d.description, "echoes input");
780        assert_eq!(d.json_schema, serde_json::json!({"type":"object"}));
781    }
782
783    #[test]
784    fn toolcall_new_sets_fields() {
785        let c = ToolCall::new("id-1", "search", serde_json::json!({"q":"x"}));
786        assert_eq!(c.id, "id-1");
787        assert_eq!(c.name, "search");
788        assert_eq!(c.args, serde_json::json!({"q":"x"}));
789    }
790
791    #[test]
792    fn chatresponse_new_carries_model() {
793        let msg = Message {
794            role: Role::Assistant,
795            content: "hi".into(),
796            tool_calls: vec![],
797            tool_call_id: None,
798        };
799        let r = ChatResponse::new(
800            msg,
801            Usage::new(1, 2),
802            FinishReason::Stop,
803            "anthropic:claude-sonnet-4-6",
804        );
805        assert_eq!(r.model, "anthropic:claude-sonnet-4-6");
806        assert_eq!(r.finish_reason, FinishReason::Stop);
807        assert_eq!(r.message.content, "hi");
808        assert_eq!(r.usage.prompt_tokens, 1);
809    }
810}
klieo_core/llm.rs

klieo_core/
llm.rs
