llm_stack/
chat.rs

1//! Conversation primitives: messages, content blocks, and responses.
2//!
3//! A conversation is a sequence of [`ChatMessage`]s, each carrying a
4//! [`ChatRole`] and one or more [`ContentBlock`]s. Providers return a
5//! [`ChatResponse`] containing the assistant's reply along with token
6//! usage and stop-reason metadata.
7//!
8//! # Content model
9//!
10//! Every message holds a `Vec<ContentBlock>` rather than a plain string.
11//! This uniform representation handles text, images, tool calls, tool
12//! results, and reasoning traces without special-casing:
13//!
14//! ```rust
15//! use llm_stack::{ChatMessage, ContentBlock, ChatRole};
16//!
17//! // Simple text message
18//! let msg = ChatMessage::user("Hello, world!");
19//!
20//! // Mixed content (text + image) in a single message
21//! let mixed = ChatMessage {
22//!     role: ChatRole::User,
23//!     content: vec![
24//!         ContentBlock::Text("What's in this image?".into()),
25//!         ContentBlock::Image {
26//!             media_type: "image/png".into(),
27//!             data: llm_stack::ImageSource::Base64("...".into()),
28//!         },
29//!     ],
30//! };
31//! ```
32
33use std::collections::HashMap;
34use std::fmt;
35
36use serde::{Deserialize, Serialize};
37use serde_json::Value;
38
39use crate::usage::Usage;
40
41/// The role of a participant in a conversation.
42///
43/// Most conversations follow the pattern `System → User → Assistant`,
44/// with [`ChatRole::Tool`] appearing only when the assistant invokes a
45/// tool and the caller feeds back the result.
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
47#[non_exhaustive]
48pub enum ChatRole {
49    /// Sets behavioral instructions for the assistant. Typically the
50    /// first message in a conversation; not all providers surface it as
51    /// a discrete message (some use a separate `system` parameter).
52    System,
53    /// A message from the human user.
54    User,
55    /// A message generated by the model.
56    Assistant,
57    /// A tool result fed back to the model after it made a tool call.
58    Tool,
59}
60
61impl fmt::Display for ChatRole {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        match self {
64            Self::System => f.write_str("system"),
65            Self::User => f.write_str("user"),
66            Self::Assistant => f.write_str("assistant"),
67            Self::Tool => f.write_str("tool"),
68        }
69    }
70}
71
72/// A single message in a conversation.
73///
74/// Use the convenience constructors for common cases:
75///
76/// ```rust
77/// use llm_stack::ChatMessage;
78///
79/// let user  = ChatMessage::user("What is 2+2?");
80/// let asst  = ChatMessage::assistant("4");
81/// let sys   = ChatMessage::system("You are a math tutor.");
82/// let tool  = ChatMessage::tool_result("call_123", "4");
83/// ```
84///
85/// For multi-block content (images, multiple tool results), construct
86/// the `content` vec directly.
87#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
88pub struct ChatMessage {
89    /// Who sent this message.
90    pub role: ChatRole,
91    /// The content blocks that make up the message body.
92    pub content: Vec<ContentBlock>,
93}
94
95impl ChatMessage {
96    /// Creates a message with a single [`ContentBlock::Text`] block.
97    pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
98        Self {
99            role,
100            content: vec![ContentBlock::Text(text.into())],
101        }
102    }
103
104    /// Shorthand for a [`ChatRole::User`] text message.
105    pub fn user(text: impl Into<String>) -> Self {
106        Self::text(ChatRole::User, text)
107    }
108
109    /// Shorthand for a [`ChatRole::Assistant`] text message.
110    pub fn assistant(text: impl Into<String>) -> Self {
111        Self::text(ChatRole::Assistant, text)
112    }
113
114    /// Shorthand for a [`ChatRole::System`] text message.
115    pub fn system(text: impl Into<String>) -> Self {
116        Self::text(ChatRole::System, text)
117    }
118
119    /// Creates a [`ChatRole::Tool`] message with a single successful
120    /// [`ToolResult`].
121    ///
122    /// Use [`ChatMessage::tool_error`] when the tool invocation failed.
123    pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
124        Self {
125            role: ChatRole::Tool,
126            content: vec![ContentBlock::ToolResult(ToolResult {
127                tool_call_id: tool_call_id.into(),
128                content: content.into(),
129                is_error: false,
130            })],
131        }
132    }
133
134    /// Creates a [`ChatRole::Tool`] message indicating the tool call
135    /// failed. The `content` should describe the error.
136    pub fn tool_error(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
137        Self {
138            role: ChatRole::Tool,
139            content: vec![ContentBlock::ToolResult(ToolResult {
140                tool_call_id: tool_call_id.into(),
141                content: content.into(),
142                is_error: true,
143            })],
144        }
145    }
146
147    /// Returns `true` if the message has zero content blocks.
148    ///
149    /// This checks the block count only — a message containing a single
150    /// `ContentBlock::Text("")` returns `false`. Use this to detect
151    /// structurally empty messages (no blocks at all).
152    pub fn is_empty(&self) -> bool {
153        self.content.is_empty()
154    }
155
156    /// Serializes the message to a stable JSON format.
157    ///
158    /// This format is documented and versioned for persistence use cases
159    /// (conversation history, checkpointing, etc.). The format is guaranteed
160    /// to be backward compatible within the same major version.
161    ///
162    /// # Stability Guarantees
163    ///
164    /// - Field names are stable within major versions
165    /// - New optional fields may be added in minor versions
166    /// - The output is deterministic for the same input
167    ///
168    /// # Example
169    ///
170    /// ```rust
171    /// use llm_stack::ChatMessage;
172    ///
173    /// let msg = ChatMessage::user("Hello, world!");
174    /// let json = msg.to_json().expect("serialization should succeed");
175    /// assert!(json["content"][0]["text"].as_str() == Some("Hello, world!"));
176    /// ```
177    pub fn to_json(&self) -> Result<Value, serde_json::Error> {
178        serde_json::to_value(self)
179    }
180
181    /// Deserializes a message from the stable JSON format.
182    ///
183    /// Accepts JSON produced by [`to_json`](Self::to_json). This method is
184    /// the inverse of `to_json` and maintains roundtrip fidelity.
185    ///
186    /// # Example
187    ///
188    /// ```rust
189    /// use llm_stack::ChatMessage;
190    /// use serde_json::json;
191    ///
192    /// let json = json!({
193    ///     "role": "User",
194    ///     "content": [{"text": "Hello!"}]
195    /// });
196    /// let msg = ChatMessage::from_json(&json).expect("valid message");
197    /// assert_eq!(msg.role, llm_stack::ChatRole::User);
198    /// ```
199    pub fn from_json(value: &Value) -> Result<Self, serde_json::Error> {
200        serde_json::from_value(value.clone())
201    }
202
203    /// Deserializes a message from the stable JSON format, consuming the value.
204    ///
205    /// This is more efficient than [`from_json`](Self::from_json) when you have
206    /// an owned `Value` that you don't need afterward, as it avoids cloning.
207    ///
208    /// # Example
209    ///
210    /// ```rust
211    /// use llm_stack::ChatMessage;
212    /// use serde_json::json;
213    ///
214    /// let json = json!({
215    ///     "role": "User",
216    ///     "content": [{"text": "Hello!"}]
217    /// });
218    /// let msg = ChatMessage::from_json_owned(json).expect("valid message");
219    /// assert_eq!(msg.role, llm_stack::ChatRole::User);
220    /// ```
221    pub fn from_json_owned(value: Value) -> Result<Self, serde_json::Error> {
222        serde_json::from_value(value)
223    }
224}
225
226/// An individual piece of content within a [`ChatMessage`].
227///
228/// Messages are composed of one or more blocks, allowing a single
229/// message to carry text alongside images, tool invocations, or
230/// reasoning traces.
231#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
232#[serde(rename_all = "snake_case")]
233#[non_exhaustive]
234pub enum ContentBlock {
235    /// Plain text content.
236    Text(String),
237    /// An inline image, either base64-encoded or referenced by URL.
238    Image {
239        /// MIME type, e.g. `"image/png"` or `"image/jpeg"`.
240        media_type: String,
241        /// The image payload.
242        data: ImageSource,
243    },
244    /// A tool invocation requested by the assistant.
245    ToolCall(ToolCall),
246    /// The result of executing a tool, sent back to the model.
247    ToolResult(ToolResult),
248    /// An internal reasoning step (chain-of-thought) emitted by the
249    /// model. Not all providers support this.
250    Reasoning {
251        /// The reasoning text.
252        content: String,
253    },
254}
255
256/// How an image is supplied to the model.
257#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
258#[non_exhaustive]
259pub enum ImageSource {
260    /// Raw image bytes encoded as a base64 string.
261    Base64(String),
262    /// A publicly accessible URL pointing to the image.
263    ///
264    /// Stores a parsed [`url::Url`] that is guaranteed to be valid.
265    /// Use [`ImageSource::from_url`] to construct from a string.
266    Url(url::Url),
267}
268
269impl ImageSource {
270    /// Creates a [`Url`](Self::Url) variant after parsing and
271    /// normalizing the URL.
272    ///
273    /// Returns the parse error if the string is not a valid URL.
274    pub fn from_url(url: impl AsRef<str>) -> Result<Self, url::ParseError> {
275        let parsed = url::Url::parse(url.as_ref())?;
276        Ok(Self::Url(parsed))
277    }
278}
279
280/// A tool invocation requested by the assistant.
281///
282/// After receiving a `ToolCall`, the caller should execute the named
283/// tool with the given `arguments`, then feed the result back as a
284/// [`ChatMessage::tool_result`] (or [`ChatMessage::tool_error`]).
285#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
286pub struct ToolCall {
287    /// Provider-assigned identifier that links this call to its result.
288    pub id: String,
289    /// The name of the tool to invoke.
290    pub name: String,
291    /// JSON arguments to pass to the tool.
292    pub arguments: Value,
293}
294
295impl std::fmt::Display for ToolCall {
296    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
297        write!(f, "{}({})", self.name, self.id)
298    }
299}
300
301/// The result of executing a tool, returned to the model.
302#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
303pub struct ToolResult {
304    /// The [`ToolCall::id`] this result corresponds to.
305    pub tool_call_id: String,
306    /// The tool's output (or error message if `is_error` is true).
307    pub content: String,
308    /// Whether the tool invocation failed.
309    pub is_error: bool,
310}
311
312impl std::fmt::Display for ToolResult {
313    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        if self.is_error {
315            write!(f, "err:{} ({})", self.tool_call_id, self.content)
316        } else {
317            write!(f, "ok:{}", self.tool_call_id)
318        }
319    }
320}
321
322/// A complete response from a model.
323///
324/// Returned by [`Provider::generate`](crate::Provider::generate). For
325/// streaming responses, accumulate [`StreamEvent`](crate::StreamEvent)s
326/// instead.
327#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
328pub struct ChatResponse {
329    /// The content blocks produced by the model.
330    pub content: Vec<ContentBlock>,
331    /// Token counts for this request/response pair.
332    pub usage: Usage,
333    /// Why the model stopped generating.
334    pub stop_reason: StopReason,
335    /// The model identifier that actually served the request (may differ
336    /// from the requested model if the provider performed routing).
337    pub model: String,
338    /// Provider-specific metadata (e.g., request IDs, cache info).
339    ///
340    /// Contents are provider-specific. See each provider crate's docs for
341    /// the keys it populates. Common keys include `"request_id"`.
342    pub metadata: HashMap<String, Value>,
343}
344
345impl ChatResponse {
346    /// Creates an empty response with no content.
347    ///
348    /// Useful for timeout or error cases where no LLM response was received.
349    pub fn empty() -> Self {
350        Self {
351            content: Vec::new(),
352            usage: Usage::default(),
353            stop_reason: StopReason::EndTurn,
354            model: String::new(),
355            metadata: HashMap::new(),
356        }
357    }
358
359    /// Returns the text of the first [`ContentBlock::Text`] block, if any.
360    ///
361    /// This is a convenience for the common case where you only care
362    /// about the model's text output. For responses that may contain
363    /// multiple text blocks, iterate `content` directly.
364    pub fn text(&self) -> Option<&str> {
365        self.content.iter().find_map(|b| match b {
366            ContentBlock::Text(t) => Some(t.as_str()),
367            _ => None,
368        })
369    }
370
371    /// Returns references to all [`ToolCall`]s in the response.
372    ///
373    /// Returns an empty `Vec` when the response contains no tool calls.
374    /// This is the primary accessor for implementing tool-use loops.
375    ///
376    /// For allocation-free iteration, use [`tool_calls_iter`](Self::tool_calls_iter).
377    pub fn tool_calls(&self) -> Vec<&ToolCall> {
378        self.content
379            .iter()
380            .filter_map(|b| match b {
381                ContentBlock::ToolCall(tc) => Some(tc),
382                _ => None,
383            })
384            .collect()
385    }
386
387    /// Returns an iterator over all [`ToolCall`]s in the response.
388    ///
389    /// Prefer this over [`tool_calls`](Self::tool_calls) to avoid allocation.
390    pub fn tool_calls_iter(&self) -> impl Iterator<Item = &ToolCall> {
391        self.content.iter().filter_map(|b| match b {
392            ContentBlock::ToolCall(tc) => Some(tc),
393            _ => None,
394        })
395    }
396
397    /// Consumes the response and returns all [`ToolCall`]s.
398    ///
399    /// This is more efficient than [`tool_calls`](Self::tool_calls) when you
400    /// need owned `ToolCall` values and won't use the response afterward.
401    /// Non-tool-call content blocks are discarded.
402    ///
403    /// # Example
404    ///
405    /// ```rust
406    /// # use llm_stack::{ChatResponse, ToolCall, ContentBlock, StopReason, Usage};
407    /// # use serde_json::json;
408    /// let response = ChatResponse {
409    ///     content: vec![
410    ///         ContentBlock::Text("Let me help".into()),
411    ///         ContentBlock::ToolCall(ToolCall {
412    ///             id: "call_1".into(),
413    ///             name: "search".into(),
414    ///             arguments: json!({"query": "rust"}),
415    ///         }),
416    ///     ],
417    ///     stop_reason: StopReason::ToolUse,
418    ///     usage: Usage::default(),
419    ///     model: "test".into(),
420    ///     metadata: Default::default(),
421    /// };
422    ///
423    /// let calls = response.into_tool_calls();
424    /// assert_eq!(calls.len(), 1);
425    /// assert_eq!(calls[0].name, "search");
426    /// ```
427    pub fn into_tool_calls(self) -> Vec<ToolCall> {
428        self.content
429            .into_iter()
430            .filter_map(|b| match b {
431                ContentBlock::ToolCall(tc) => Some(tc),
432                _ => None,
433            })
434            .collect()
435    }
436
437    /// Consumes the response content and partitions it into tool calls and other blocks.
438    ///
439    /// Returns `(tool_calls, other_content)` where `other_content` contains all
440    /// non-tool-call blocks (`Text`, `Image`, `ToolResult`, etc.) suitable for building
441    /// an assistant message in a tool loop.
442    ///
443    /// This is more efficient than calling both [`tool_calls`](Self::tool_calls)
444    /// and filtering content separately, as it processes the content in a single pass.
445    pub fn partition_content(self) -> (Vec<ToolCall>, Vec<ContentBlock>) {
446        let mut tool_calls = Vec::new();
447        let mut other = Vec::new();
448
449        for block in self.content {
450            match block {
451                ContentBlock::ToolCall(tc) => tool_calls.push(tc),
452                // Filter out ToolResult as they shouldn't be in assistant messages
453                ContentBlock::ToolResult(_) => {}
454                other_block => other.push(other_block),
455            }
456        }
457
458        (tool_calls, other)
459    }
460}
461
462/// The reason the model stopped producing output.
463#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
464#[non_exhaustive]
465pub enum StopReason {
466    /// The model finished its response naturally.
467    EndTurn,
468    /// The model wants to invoke one or more tools.
469    ToolUse,
470    /// The response hit the `max_tokens` limit.
471    MaxTokens,
472    /// The model emitted a stop sequence.
473    StopSequence,
474}
475
476impl fmt::Display for StopReason {
477    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
478        match self {
479            Self::EndTurn => f.write_str("end_turn"),
480            Self::ToolUse => f.write_str("tool_use"),
481            Self::MaxTokens => f.write_str("max_tokens"),
482            Self::StopSequence => f.write_str("stop_sequence"),
483        }
484    }
485}
486
487#[cfg(test)]
488mod tests {
489    use super::*;
490
491    #[test]
492    fn test_chat_role_copy_hash() {
493        use std::collections::HashMap;
494        let mut map = HashMap::new();
495        let role = ChatRole::User;
496        let role_copy = role; // Copy
497        map.insert(role, "user");
498        map.insert(role_copy, "user_copy");
499        assert_eq!(map.len(), 1);
500    }
501
502    #[test]
503    fn test_chat_role_all_variants() {
504        let variants = [
505            ChatRole::System,
506            ChatRole::User,
507            ChatRole::Assistant,
508            ChatRole::Tool,
509        ];
510        for v in &variants {
511            let debug = format!("{v:?}");
512            assert!(!debug.is_empty());
513        }
514    }
515
516    #[test]
517    fn test_chat_role_serde_roundtrip() {
518        let role = ChatRole::Assistant;
519        let json = serde_json::to_string(&role).unwrap();
520        let back: ChatRole = serde_json::from_str(&json).unwrap();
521        assert_eq!(role, back);
522    }
523
524    // --- ChatMessage constructors ---
525
526    #[test]
527    fn test_user_constructor() {
528        let msg = ChatMessage::user("hello");
529        assert_eq!(msg.role, ChatRole::User);
530        assert_eq!(msg.content, vec![ContentBlock::Text("hello".into())]);
531    }
532
533    #[test]
534    fn test_assistant_constructor() {
535        let msg = ChatMessage::assistant("hi");
536        assert_eq!(msg.role, ChatRole::Assistant);
537        assert_eq!(msg.content, vec![ContentBlock::Text("hi".into())]);
538    }
539
540    #[test]
541    fn test_system_constructor() {
542        let msg = ChatMessage::system("be nice");
543        assert_eq!(msg.role, ChatRole::System);
544    }
545
546    #[test]
547    fn test_tool_result_constructor() {
548        let msg = ChatMessage::tool_result("tc_1", "42");
549        assert_eq!(msg.role, ChatRole::Tool);
550        assert!(matches!(
551            &msg.content[0],
552            ContentBlock::ToolResult(tr)
553                if tr.tool_call_id == "tc_1" && tr.content == "42" && !tr.is_error
554        ));
555    }
556
557    #[test]
558    fn test_tool_error_constructor() {
559        let msg = ChatMessage::tool_error("tc_1", "something broke");
560        assert!(matches!(
561            &msg.content[0],
562            ContentBlock::ToolResult(tr) if tr.is_error
563        ));
564    }
565
566    // --- ChatMessage clone/eq/serde ---
567
568    #[test]
569    fn test_message_text_clone_eq() {
570        let msg = ChatMessage::user("hello");
571        assert_eq!(msg, msg.clone());
572    }
573
574    #[test]
575    fn test_message_serde_roundtrip() {
576        let msg = ChatMessage::user("hello");
577        let json = serde_json::to_string(&msg).unwrap();
578        let back: ChatMessage = serde_json::from_str(&json).unwrap();
579        assert_eq!(msg, back);
580    }
581
582    #[test]
583    fn test_message_tool_use() {
584        let msg = ChatMessage {
585            role: ChatRole::Assistant,
586            content: vec![
587                ContentBlock::ToolCall(ToolCall {
588                    id: "1".into(),
589                    name: "calc".into(),
590                    arguments: serde_json::json!({"a": 1}),
591                }),
592                ContentBlock::ToolCall(ToolCall {
593                    id: "2".into(),
594                    name: "search".into(),
595                    arguments: serde_json::json!({"q": "rust"}),
596                }),
597            ],
598        };
599        assert_eq!(msg.content.len(), 2);
600        assert_eq!(msg, msg.clone());
601    }
602
603    #[test]
604    fn test_message_tool_result() {
605        let msg = ChatMessage::tool_result("1", "42");
606        assert!(matches!(
607            &msg.content[0],
608            ContentBlock::ToolResult(tr) if tr.content == "42" && !tr.is_error
609        ));
610    }
611
612    #[test]
613    fn test_message_mixed_content() {
614        let msg = ChatMessage {
615            role: ChatRole::User,
616            content: vec![
617                ContentBlock::Text("look at this".into()),
618                ContentBlock::Image {
619                    media_type: "image/png".into(),
620                    data: ImageSource::Base64("abc123".into()),
621                },
622                ContentBlock::ToolCall(ToolCall {
623                    id: "1".into(),
624                    name: "analyze".into(),
625                    arguments: serde_json::json!({}),
626                }),
627            ],
628        };
629        assert_eq!(msg.content.len(), 3);
630    }
631
632    // --- ContentBlock tests ---
633
634    #[test]
635    fn test_content_block_image_base64() {
636        let block = ContentBlock::Image {
637            media_type: "image/jpeg".into(),
638            data: ImageSource::Base64("data...".into()),
639        };
640        assert_eq!(block, block.clone());
641    }
642
643    #[test]
644    fn test_content_block_image_url() {
645        let block = ContentBlock::Image {
646            media_type: "image/png".into(),
647            data: ImageSource::from_url("https://example.com/img.png").unwrap(),
648        };
649        assert_eq!(block, block.clone());
650    }
651
652    #[test]
653    fn test_image_source_from_url_valid() {
654        let src = ImageSource::from_url("https://example.com/img.png");
655        assert!(src.is_ok());
656        let url = url::Url::parse("https://example.com/img.png").unwrap();
657        assert_eq!(src.unwrap(), ImageSource::Url(url));
658    }
659
660    #[test]
661    fn test_image_source_from_url_normalizes() {
662        // URL is parsed and normalized
663        let src = ImageSource::from_url("HTTP://EXAMPLE.COM").unwrap();
664        assert!(matches!(
665            &src,
666            ImageSource::Url(u) if u.as_str() == "http://example.com/"
667        ));
668    }
669
670    #[test]
671    fn test_image_source_from_url_invalid() {
672        let err = ImageSource::from_url("not a url");
673        assert!(err.is_err());
674        let _parse_err: url::ParseError = err.unwrap_err();
675
676        assert!(ImageSource::from_url("").is_err());
677    }
678
679    #[test]
680    fn test_content_block_reasoning() {
681        let block = ContentBlock::Reasoning {
682            content: "thinking step by step".into(),
683        };
684        assert_eq!(block, block.clone());
685    }
686
687    #[test]
688    fn test_tool_call_json_arguments() {
689        let call = ToolCall {
690            id: "tc_1".into(),
691            name: "search".into(),
692            arguments: serde_json::json!({
693                "query": "rust async",
694                "filters": {"lang": "en", "limit": 10}
695            }),
696        };
697        assert_eq!(call, call.clone());
698    }
699
700    #[test]
701    fn test_tool_result_error_flag() {
702        let ok = ToolResult {
703            tool_call_id: "1".into(),
704            content: "result".into(),
705            is_error: false,
706        };
707        let err = ToolResult {
708            tool_call_id: "1".into(),
709            content: "result".into(),
710            is_error: true,
711        };
712        assert_ne!(ok, err);
713    }
714
715    // --- ChatResponse tests ---
716
717    #[test]
718    fn test_chat_response_metadata() {
719        let mut metadata = HashMap::new();
720        metadata.insert("cost".into(), serde_json::json!({"usd": 0.01}));
721        let resp = ChatResponse {
722            content: vec![ContentBlock::Text("hi".into())],
723            usage: Usage::default(),
724            stop_reason: StopReason::EndTurn,
725            model: "test-model".into(),
726            metadata,
727        };
728        assert!(resp.metadata.contains_key("cost"));
729    }
730
731    #[test]
732    fn test_chat_response_serde_roundtrip() {
733        let resp = ChatResponse {
734            content: vec![ContentBlock::Text("hi".into())],
735            usage: Usage::default(),
736            stop_reason: StopReason::EndTurn,
737            model: "test-model".into(),
738            metadata: HashMap::new(),
739        };
740        let json = serde_json::to_string(&resp).unwrap();
741        let back: ChatResponse = serde_json::from_str(&json).unwrap();
742        assert_eq!(resp, back);
743    }
744
745    #[test]
746    fn test_chat_response_empty_content() {
747        let resp = ChatResponse {
748            content: vec![],
749            usage: Usage::default(),
750            stop_reason: StopReason::EndTurn,
751            model: "test".into(),
752            metadata: HashMap::new(),
753        };
754        assert!(resp.content.is_empty());
755    }
756
757    // --- StopReason tests ---
758
759    #[test]
760    fn test_stop_reason_all_variants() {
761        let variants = [
762            StopReason::EndTurn,
763            StopReason::ToolUse,
764            StopReason::MaxTokens,
765            StopReason::StopSequence,
766        ];
767        for v in &variants {
768            assert_eq!(*v, *v);
769        }
770    }
771
772    #[test]
773    fn test_stop_reason_serde_roundtrip() {
774        let sr = StopReason::MaxTokens;
775        let json = serde_json::to_string(&sr).unwrap();
776        let back: StopReason = serde_json::from_str(&json).unwrap();
777        assert_eq!(sr, back);
778    }
779
780    #[test]
781    fn test_stop_reason_eq_hash() {
782        use std::collections::HashMap;
783        let mut map = HashMap::new();
784        map.insert(StopReason::EndTurn, "end");
785        map.insert(StopReason::ToolUse, "tool");
786        assert_eq!(map[&StopReason::EndTurn], "end");
787        assert_eq!(map[&StopReason::ToolUse], "tool");
788    }
789
790    // --- Display impls ---
791
792    #[test]
793    fn test_chat_role_display() {
794        assert_eq!(ChatRole::System.to_string(), "system");
795        assert_eq!(ChatRole::User.to_string(), "user");
796        assert_eq!(ChatRole::Assistant.to_string(), "assistant");
797        assert_eq!(ChatRole::Tool.to_string(), "tool");
798    }
799
800    #[test]
801    fn test_stop_reason_display() {
802        assert_eq!(StopReason::EndTurn.to_string(), "end_turn");
803        assert_eq!(StopReason::ToolUse.to_string(), "tool_use");
804        assert_eq!(StopReason::MaxTokens.to_string(), "max_tokens");
805        assert_eq!(StopReason::StopSequence.to_string(), "stop_sequence");
806    }
807
808    // --- ChatResponse::text() ---
809
810    #[test]
811    fn test_chat_response_text_returns_first() {
812        let resp = ChatResponse {
813            content: vec![
814                ContentBlock::Reasoning {
815                    content: "thinking...".into(),
816                },
817                ContentBlock::Text("first".into()),
818                ContentBlock::Text("second".into()),
819            ],
820            usage: Usage::default(),
821            stop_reason: StopReason::EndTurn,
822            model: "test".into(),
823            metadata: HashMap::new(),
824        };
825        assert_eq!(resp.text(), Some("first"));
826    }
827
828    #[test]
829    fn test_chat_response_text_none_when_no_text_blocks() {
830        let resp = ChatResponse {
831            content: vec![ContentBlock::Reasoning {
832                content: "thinking".into(),
833            }],
834            usage: Usage::default(),
835            stop_reason: StopReason::EndTurn,
836            model: "test".into(),
837            metadata: HashMap::new(),
838        };
839        assert_eq!(resp.text(), None);
840    }
841
842    #[test]
843    fn test_chat_response_text_none_when_empty() {
844        let resp = ChatResponse {
845            content: vec![],
846            usage: Usage::default(),
847            stop_reason: StopReason::EndTurn,
848            model: "test".into(),
849            metadata: HashMap::new(),
850        };
851        assert_eq!(resp.text(), None);
852    }
853
854    // --- ChatResponse::tool_calls() ---
855
856    #[test]
857    fn test_chat_response_tool_calls() {
858        let resp = ChatResponse {
859            content: vec![
860                ContentBlock::Text("Let me search.".into()),
861                ContentBlock::ToolCall(ToolCall {
862                    id: "1".into(),
863                    name: "search".into(),
864                    arguments: serde_json::json!({"q": "rust"}),
865                }),
866                ContentBlock::ToolCall(ToolCall {
867                    id: "2".into(),
868                    name: "calc".into(),
869                    arguments: serde_json::json!({"expr": "2+2"}),
870                }),
871            ],
872            usage: Usage::default(),
873            stop_reason: StopReason::ToolUse,
874            model: "test".into(),
875            metadata: HashMap::new(),
876        };
877        let calls = resp.tool_calls();
878        assert_eq!(calls.len(), 2);
879        assert_eq!(calls[0].name, "search");
880        assert_eq!(calls[1].name, "calc");
881    }
882
883    #[test]
884    fn test_chat_response_tool_calls_empty_when_text_only() {
885        let resp = ChatResponse {
886            content: vec![ContentBlock::Text("hello".into())],
887            usage: Usage::default(),
888            stop_reason: StopReason::EndTurn,
889            model: "test".into(),
890            metadata: HashMap::new(),
891        };
892        assert!(resp.tool_calls().is_empty());
893    }
894
895    // --- is_empty ---
896
897    #[test]
898    fn test_message_is_empty() {
899        let empty = ChatMessage {
900            role: ChatRole::User,
901            content: vec![],
902        };
903        assert!(empty.is_empty());
904        assert!(!ChatMessage::user("hi").is_empty());
905    }
906
907    // --- ContentBlock serde wire format ---
908
909    #[test]
910    fn test_content_block_serde_text() {
911        let block = ContentBlock::Text("hello".into());
912        let val = serde_json::to_value(&block).unwrap();
913        assert_eq!(val, serde_json::json!({"text": "hello"}));
914        let back: ContentBlock = serde_json::from_value(val).unwrap();
915        assert_eq!(back, block);
916    }
917
918    #[test]
919    fn test_content_block_serde_image() {
920        let block = ContentBlock::Image {
921            media_type: "image/png".into(),
922            data: ImageSource::Base64("abc".into()),
923        };
924        let val = serde_json::to_value(&block).unwrap();
925        assert_eq!(
926            val,
927            serde_json::json!({"image": {"media_type": "image/png", "data": {"Base64": "abc"}}})
928        );
929        let back: ContentBlock = serde_json::from_value(val).unwrap();
930        assert_eq!(back, block);
931    }
932
933    #[test]
934    fn test_content_block_serde_tool_call() {
935        let block = ContentBlock::ToolCall(ToolCall {
936            id: "tc_1".into(),
937            name: "search".into(),
938            arguments: serde_json::json!({"q": "rust"}),
939        });
940        let val = serde_json::to_value(&block).unwrap();
941        assert_eq!(
942            val,
943            serde_json::json!({"tool_call": {"id": "tc_1", "name": "search", "arguments": {"q": "rust"}}})
944        );
945        let back: ContentBlock = serde_json::from_value(val).unwrap();
946        assert_eq!(back, block);
947    }
948
949    #[test]
950    fn test_content_block_serde_tool_result() {
951        let block = ContentBlock::ToolResult(ToolResult {
952            tool_call_id: "tc_1".into(),
953            content: "42".into(),
954            is_error: false,
955        });
956        let val = serde_json::to_value(&block).unwrap();
957        assert_eq!(
958            val,
959            serde_json::json!({"tool_result": {"tool_call_id": "tc_1", "content": "42", "is_error": false}})
960        );
961        let back: ContentBlock = serde_json::from_value(val).unwrap();
962        assert_eq!(back, block);
963    }
964
965    #[test]
966    fn test_content_block_serde_reasoning() {
967        let block = ContentBlock::Reasoning {
968            content: "thinking".into(),
969        };
970        let val = serde_json::to_value(&block).unwrap();
971        assert_eq!(
972            val,
973            serde_json::json!({"reasoning": {"content": "thinking"}})
974        );
975        let back: ContentBlock = serde_json::from_value(val).unwrap();
976        assert_eq!(back, block);
977    }
978
979    // --- Semantic validity: constructors produce valid role/content combinations ---
980
981    #[test]
982    fn test_user_constructor_produces_text_only() {
983        let msg = ChatMessage::user("hello");
984        assert_eq!(msg.role, ChatRole::User);
985        assert!(
986            msg.content
987                .iter()
988                .all(|b| matches!(b, ContentBlock::Text(_)))
989        );
990    }
991
992    #[test]
993    fn test_assistant_constructor_produces_text_only() {
994        let msg = ChatMessage::assistant("hi");
995        assert_eq!(msg.role, ChatRole::Assistant);
996        assert!(
997            msg.content
998                .iter()
999                .all(|b| matches!(b, ContentBlock::Text(_)))
1000        );
1001    }
1002
1003    #[test]
1004    fn test_system_constructor_produces_text_only() {
1005        let msg = ChatMessage::system("be nice");
1006        assert_eq!(msg.role, ChatRole::System);
1007        assert!(
1008            msg.content
1009                .iter()
1010                .all(|b| matches!(b, ContentBlock::Text(_)))
1011        );
1012    }
1013
1014    #[test]
1015    fn test_tool_result_constructor_produces_tool_result_only() {
1016        let msg = ChatMessage::tool_result("tc_1", "42");
1017        assert_eq!(msg.role, ChatRole::Tool);
1018        assert!(
1019            msg.content
1020                .iter()
1021                .all(|b| matches!(b, ContentBlock::ToolResult(_)))
1022        );
1023    }
1024
1025    #[test]
1026    fn test_tool_error_constructor_produces_tool_result_only() {
1027        let msg = ChatMessage::tool_error("tc_1", "boom");
1028        assert_eq!(msg.role, ChatRole::Tool);
1029        assert!(
1030            msg.content
1031                .iter()
1032                .all(|b| matches!(b, ContentBlock::ToolResult(r) if r.is_error))
1033        );
1034    }
1035
1036    #[test]
1037    fn test_assistant_tool_calls_is_valid_combination() {
1038        // Assistants can produce tool calls — this is a valid state
1039        let msg = ChatMessage {
1040            role: ChatRole::Assistant,
1041            content: vec![
1042                ContentBlock::Text("Let me search for that.".into()),
1043                ContentBlock::ToolCall(ToolCall {
1044                    id: "1".into(),
1045                    name: "search".into(),
1046                    arguments: serde_json::json!({"q": "rust"}),
1047                }),
1048            ],
1049        };
1050        assert_eq!(msg.role, ChatRole::Assistant);
1051        assert_eq!(msg.content.len(), 2);
1052    }
1053
1054    #[test]
1055    fn test_user_with_image_is_valid_combination() {
1056        // Users can send text + images — this is a valid state
1057        let msg = ChatMessage {
1058            role: ChatRole::User,
1059            content: vec![
1060                ContentBlock::Text("What's this?".into()),
1061                ContentBlock::Image {
1062                    media_type: "image/png".into(),
1063                    data: ImageSource::Base64("...".into()),
1064                },
1065            ],
1066        };
1067        assert_eq!(msg.role, ChatRole::User);
1068        assert_eq!(msg.content.len(), 2);
1069    }
1070
1071    #[test]
1072    fn test_assistant_with_reasoning_is_valid_combination() {
1073        // Assistants can produce reasoning + text — valid state
1074        let msg = ChatMessage {
1075            role: ChatRole::Assistant,
1076            content: vec![
1077                ContentBlock::Reasoning {
1078                    content: "step 1: think about it".into(),
1079                },
1080                ContentBlock::Text("The answer is 42.".into()),
1081            ],
1082        };
1083        assert_eq!(msg.role, ChatRole::Assistant);
1084        assert_eq!(msg.content.len(), 2);
1085    }
1086
1087    // --- ChatMessage serialization stability tests ---
1088
1089    #[test]
1090    fn test_chat_message_to_json() {
1091        let msg = ChatMessage::user("Hello, world!");
1092        let json = msg.to_json().unwrap();
1093        assert_eq!(json["role"], "User");
1094        assert_eq!(json["content"][0]["text"], "Hello, world!");
1095    }
1096
1097    #[test]
1098    fn test_chat_message_from_json() {
1099        let json = serde_json::json!({
1100            "role": "Assistant",
1101            "content": [{"text": "Hello!"}]
1102        });
1103        let msg = ChatMessage::from_json(&json).unwrap();
1104        assert_eq!(msg.role, ChatRole::Assistant);
1105        assert!(matches!(&msg.content[0], ContentBlock::Text(t) if t == "Hello!"));
1106    }
1107
1108    #[test]
1109    fn test_chat_message_json_roundtrip() {
1110        let original = ChatMessage {
1111            role: ChatRole::User,
1112            content: vec![
1113                ContentBlock::Text("What's this?".into()),
1114                ContentBlock::Image {
1115                    media_type: "image/png".into(),
1116                    data: ImageSource::Base64("abc123".into()),
1117                },
1118            ],
1119        };
1120        let json = original.to_json().unwrap();
1121        let restored = ChatMessage::from_json(&json).unwrap();
1122        assert_eq!(original, restored);
1123    }
1124
1125    #[test]
1126    fn test_chat_message_json_roundtrip_with_tool_result() {
1127        let original = ChatMessage::tool_result("tc_1", "success");
1128        let json = original.to_json().unwrap();
1129        let restored = ChatMessage::from_json(&json).unwrap();
1130        assert_eq!(original, restored);
1131    }
1132
1133    // --- ToolResult metadata tests ---
1134}
llm_stack/chat.rs

llm_stack/
chat.rs
