autoagents_llm/chat/
mod.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::pin::Pin;
4
5use async_trait::async_trait;
6use futures::stream::Stream;
7#[cfg(not(target_arch = "wasm32"))]
8use futures::stream::StreamExt;
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12use crate::{ToolCall, error::LLMError};
13
14/// Usage metadata for a chat response.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct Usage {
17    /// Number of tokens in the prompt
18    #[serde(alias = "input_tokens")]
19    pub prompt_tokens: u32,
20    /// Number of tokens in the completion
21    #[serde(alias = "output_tokens")]
22    pub completion_tokens: u32,
23    /// Total number of tokens used
24    pub total_tokens: u32,
25    /// Breakdown of completion tokens, if available
26    #[serde(
27        skip_serializing_if = "Option::is_none",
28        alias = "output_tokens_details"
29    )]
30    pub completion_tokens_details: Option<CompletionTokensDetails>,
31    /// Breakdown of prompt tokens, if available
32    #[serde(
33        skip_serializing_if = "Option::is_none",
34        alias = "input_tokens_details"
35    )]
36    pub prompt_tokens_details: Option<PromptTokensDetails>,
37}
38
39/// Stream response chunk that mimics OpenAI's streaming response format
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct StreamResponse {
42    /// Array of choices in the response
43    pub choices: Vec<StreamChoice>,
44    /// Usage metadata, typically present in the final chunk
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub usage: Option<Usage>,
47}
48
49/// Individual choice in a streaming response
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct StreamChoice {
52    /// Delta containing the incremental content
53    pub delta: StreamDelta,
54}
55
56/// Delta content in a streaming response
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct StreamDelta {
59    /// The incremental content, if any
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub content: Option<String>,
62    /// The incremental reasoning content, if any
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub reasoning_content: Option<String>,
65    /// The incremental tool calls, if any
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub tool_calls: Option<Vec<ToolCall>>,
68}
69
70/// A streaming chunk that can be either text or a tool call event.
71///
72/// This enum provides a unified representation of streaming events
73/// when using `chat_stream_with_tools`. It allows callers to receive
74/// text deltas as they arrive while also handling tool use blocks.
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub enum StreamChunk {
77    /// Text content delta
78    Text(String),
79    /// Reasoning content delta
80    ReasoningContent(String),
81
82    /// Tool use block started (contains tool id and name)
83    ToolUseStart {
84        /// The index of this content block in the response
85        index: usize,
86        /// The unique ID for this tool use
87        id: String,
88        /// The name of the tool being called
89        name: String,
90    },
91
92    /// Tool use input JSON delta (partial JSON string)
93    ToolUseInputDelta {
94        /// The index of this content block
95        index: usize,
96        /// Partial JSON string for the tool input
97        partial_json: String,
98    },
99
100    /// Tool use block complete with assembled ToolCall
101    ToolUseComplete {
102        /// The index of this content block
103        index: usize,
104        /// The complete tool call with id, name, and parsed arguments
105        tool_call: ToolCall,
106    },
107
108    /// Stream ended with stop reason
109    Done {
110        /// The reason the stream stopped (e.g., "end_turn", "tool_use")
111        stop_reason: String,
112    },
113    Usage(Usage),
114}
115
116/// Breakdown of completion tokens.
117#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
118pub struct CompletionTokensDetails {
119    /// Tokens used for reasoning (for reasoning models)
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub reasoning_tokens: Option<u32>,
122    /// Tokens used for audio output
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub audio_tokens: Option<u32>,
125}
126
127/// Breakdown of prompt tokens.
128#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
129pub struct PromptTokensDetails {
130    /// Tokens used for cached content
131    #[serde(skip_serializing_if = "Option::is_none")]
132    pub cached_tokens: Option<u32>,
133    /// Tokens used for audio input
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub audio_tokens: Option<u32>,
136}
137
138/// Role of a participant in a chat conversation.
139#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
140pub enum ChatRole {
141    // The system Prompt
142    System,
143    /// The user/human participant in the conversation
144    User,
145    /// The AI assistant participant in the conversation
146    Assistant,
147    /// Tool/function response
148    Tool,
149}
150
151impl fmt::Display for ChatRole {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        let value = match self {
154            ChatRole::System => "system",
155            ChatRole::User => "user",
156            ChatRole::Assistant => "assistant",
157            ChatRole::Tool => "tool",
158        };
159        f.write_str(value)
160    }
161}
162
163/// The supported MIME type of an image.
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
165#[non_exhaustive]
166pub enum ImageMime {
167    /// JPEG image
168    JPEG,
169    /// PNG image
170    PNG,
171    /// GIF image
172    GIF,
173    /// WebP image
174    WEBP,
175}
176
177impl ImageMime {
178    pub fn mime_type(&self) -> &'static str {
179        match self {
180            ImageMime::JPEG => "image/jpeg",
181            ImageMime::PNG => "image/png",
182            ImageMime::GIF => "image/gif",
183            ImageMime::WEBP => "image/webp",
184        }
185    }
186}
187
188/// The type of a message in a chat conversation.
189#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
190pub enum MessageType {
191    /// A text message
192    #[default]
193    Text,
194    /// An image message
195    Image((ImageMime, Vec<u8>)),
196    /// PDF message
197    Pdf(Vec<u8>),
198    /// An image URL message
199    ImageURL(String),
200    /// A tool use
201    ToolUse(Vec<ToolCall>),
202    /// Tool result
203    ToolResult(Vec<ToolCall>),
204}
205
206/// The type of reasoning effort for a message in a chat conversation.
207pub enum ReasoningEffort {
208    /// Low reasoning effort
209    Low,
210    /// Medium reasoning effort
211    Medium,
212    /// High reasoning effort
213    High,
214}
215
216/// A single message in a chat conversation.
217#[derive(Debug, Clone, Serialize, Deserialize)]
218pub struct ChatMessage {
219    /// The role of who sent this message (user or assistant)
220    pub role: ChatRole,
221    /// The type of the message (text, image, audio, video, etc)
222    pub message_type: MessageType,
223    /// The text content of the message
224    pub content: String,
225}
226
227/// Represents a parameter in a function tool
228#[derive(Debug, Clone, Serialize)]
229pub struct ParameterProperty {
230    /// The type of the parameter (e.g. "string", "number", "array", etc)
231    #[serde(rename = "type")]
232    pub property_type: String,
233    /// Description of what the parameter does
234    pub description: String,
235    /// When type is "array", this defines the type of the array items
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub items: Option<Box<ParameterProperty>>,
238    /// When type is "enum", this defines the possible values for the parameter
239    #[serde(skip_serializing_if = "Option::is_none", rename = "enum")]
240    pub enum_list: Option<Vec<String>>,
241}
242
243/// Represents the parameters schema for a function tool
244#[derive(Debug, Clone, Serialize)]
245pub struct ParametersSchema {
246    /// The type of the parameters object (usually "object")
247    #[serde(rename = "type")]
248    pub schema_type: String,
249    /// Map of parameter names to their properties
250    pub properties: HashMap<String, ParameterProperty>,
251    /// List of required parameter names
252    pub required: Vec<String>,
253}
254
255/// Represents a function definition for a tool.
256///
257/// The `parameters` field stores the JSON Schema describing the function
258/// arguments.  It is kept as a raw `serde_json::Value` to allow arbitrary
259/// complexity (nested arrays/objects, `oneOf`, etc.) without requiring a
260/// bespoke Rust structure.
261///
262/// Builder helpers can still generate simple schemas automatically, but the
263/// user may also provide any valid schema directly.
264#[derive(Debug, Clone, Serialize)]
265pub struct FunctionTool {
266    /// Name of the function
267    pub name: String,
268    /// Human-readable description
269    pub description: String,
270    /// JSON Schema describing the parameters
271    pub parameters: Value,
272}
273
274/// Defines rules for structured output responses based on [OpenAI's structured output requirements](https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format).
275/// Individual providers may have additional requirements or restrictions, but these should be handled by each provider's backend implementation.
276///
277/// If you plan on deserializing into this struct, make sure the source text has a `"name"` field, since that's technically the only thing required by OpenAI.
278///
279/// ## Example
280///
281/// ```
282/// use autoagents_llm::chat::StructuredOutputFormat;
283/// use serde_json::json;
284///
285/// let response_format = r#"
286///     {
287///         "name": "Student",
288///         "description": "A student object",
289///         "schema": {
290///             "type": "object",
291///             "properties": {
292///                 "name": {
293///                     "type": "string"
294///                 },
295///                 "age": {
296///                     "type": "integer"
297///                 },
298///                 "is_student": {
299///                     "type": "boolean"
300///                 }
301///             },
302///             "required": ["name", "age", "is_student"]
303///         }
304///     }
305/// "#;
306/// let structured_output: StructuredOutputFormat = serde_json::from_str(response_format).unwrap();
307/// assert_eq!(structured_output.name, "Student");
308/// assert_eq!(structured_output.description, Some("A student object".to_string()));
309/// ```
310#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
311
312pub struct StructuredOutputFormat {
313    /// Name of the schema
314    pub name: String,
315    /// The description of the schema
316    pub description: Option<String>,
317    /// The JSON schema for the structured output
318    pub schema: Option<Value>,
319    /// Whether to enable strict schema adherence
320    pub strict: Option<bool>,
321}
322
323/// Represents a tool that can be used in chat
324#[derive(Debug, Clone, Serialize)]
325pub struct Tool {
326    /// The type of tool (e.g. "function")
327    #[serde(rename = "type")]
328    pub tool_type: String,
329    /// The function definition if this is a function tool
330    pub function: FunctionTool,
331}
332
333/// Tool choice determines how the LLM uses available tools.
334/// The behavior is standardized across different LLM providers.
335#[derive(Debug, Clone, Default)]
336pub enum ToolChoice {
337    /// Model can use any tool, but it must use at least one.
338    /// This is useful when you want to force the model to use tools.
339    Any,
340
341    /// Model can use any tool, and may elect to use none.
342    /// This is the default behavior and gives the model flexibility.
343    #[default]
344    Auto,
345
346    /// Model must use the specified tool and only the specified tool.
347    /// The string parameter is the name of the required tool.
348    /// This is useful when you want the model to call a specific function.
349    Tool(String),
350
351    /// Explicitly disables the use of tools.
352    /// The model will not use any tools even if they are provided.
353    None,
354}
355
356impl Serialize for ToolChoice {
357    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
358    where
359        S: serde::Serializer,
360    {
361        match self {
362            ToolChoice::Any => serializer.serialize_str("required"),
363            ToolChoice::Auto => serializer.serialize_str("auto"),
364            ToolChoice::None => serializer.serialize_str("none"),
365            ToolChoice::Tool(name) => {
366                use serde::ser::SerializeMap;
367
368                // For tool_choice: {"type": "function", "function": {"name": "function_name"}}
369                let mut map = serializer.serialize_map(Some(2))?;
370                map.serialize_entry("type", "function")?;
371
372                // Inner function object
373                let mut function_obj = std::collections::HashMap::new();
374                function_obj.insert("name", name.as_str());
375
376                map.serialize_entry("function", &function_obj)?;
377                map.end()
378            }
379        }
380    }
381}
382
383pub trait ChatResponse: std::fmt::Debug + std::fmt::Display + Send + Sync {
384    fn text(&self) -> Option<String>;
385    fn tool_calls(&self) -> Option<Vec<ToolCall>>;
386    fn thinking(&self) -> Option<String> {
387        None
388    }
389    fn usage(&self) -> Option<Usage> {
390        None
391    }
392}
393
394/// Per-call sampling overrides for [`ChatProvider`] methods.
395///
396/// Backends that support per-call sampling (e.g. `LlamaCppProvider`) apply
397/// these overrides on top of the defaults configured at provider construction.
398/// Backends that do not support per-call overrides (the default trait impl)
399/// silently ignore overrides — passing `Some(SamplingOverrides::...)` is safe
400/// against any backend.
401///
402/// `None` on any field means "use the provider default" (no override). Passing
403/// `sampling: None` to the `_and_sampling` methods is equivalent to calling
404/// the non-sampling-aware method (no behaviour change).
405#[derive(Debug, Default, Clone, PartialEq)]
406pub struct SamplingOverrides {
407    /// Temperature override. `None` = use provider default.
408    pub temperature: Option<f32>,
409    /// Top-p (nucleus sampling) override. `None` = use provider default.
410    pub top_p: Option<f32>,
411    /// Max output tokens override. `None` = use provider default.
412    pub max_tokens: Option<u32>,
413}
414
415impl SamplingOverrides {
416    /// All overrides unset. Equivalent to [`SamplingOverrides::default`] —
417    /// included for call-site readability when explicitly opting out.
418    pub fn empty() -> Self {
419        Self::default()
420    }
421
422    /// Convenience constructor: override only `temperature`.
423    pub fn with_temperature(temperature: f32) -> Self {
424        Self {
425            temperature: Some(temperature),
426            ..Self::default()
427        }
428    }
429
430    /// Convenience constructor: override only `top_p`.
431    pub fn with_top_p(top_p: f32) -> Self {
432        Self {
433            top_p: Some(top_p),
434            ..Self::default()
435        }
436    }
437
438    /// Convenience constructor: override only `max_tokens`.
439    pub fn with_max_tokens(max_tokens: u32) -> Self {
440        Self {
441            max_tokens: Some(max_tokens),
442            ..Self::default()
443        }
444    }
445}
446
447/// Trait for providers that support chat-style interactions.
448#[async_trait]
449pub trait ChatProvider: Sync + Send {
450    /// Sends a chat request to the provider with a sequence of messages.
451    ///
452    /// # Arguments
453    ///
454    /// * `messages` - The conversation history as a slice of chat messages
455    /// * `json_schema` - Optional json_schema for the response format
456    ///
457    /// # Returns
458    ///
459    /// The provider's response text or an error
460    async fn chat(
461        &self,
462        messages: &[ChatMessage],
463        json_schema: Option<StructuredOutputFormat>,
464    ) -> Result<Box<dyn ChatResponse>, LLMError> {
465        self.chat_with_tools(messages, None, json_schema).await
466    }
467
468    /// Sends a chat request to the provider with a sequence of messages and tools.
469    ///
470    /// # Arguments
471    ///
472    /// * `messages` - The conversation history as a slice of chat messages
473    /// * `tools` - Optional slice of tools to use in the chat
474    /// * `json_schema` - Optional json_schema for the response format
475    ///
476    /// # Returns
477    ///
478    /// The provider's response text or an error
479    async fn chat_with_tools(
480        &self,
481        messages: &[ChatMessage],
482        tools: Option<&[Tool]>,
483        json_schema: Option<StructuredOutputFormat>,
484    ) -> Result<Box<dyn ChatResponse>, LLMError>;
485
486    /// Sends a chat request with optional per-call sampling overrides.
487    ///
488    /// Equivalent to [`ChatProvider::chat`] when `sampling` is `None`. When
489    /// `sampling` is `Some(...)`, backends that support per-call sampling
490    /// apply the overrides on top of provider-construction defaults; backends
491    /// that do not (the default impl) silently ignore them.
492    ///
493    /// Backwards compatible: callers that don't need per-call sampling should
494    /// continue to use [`ChatProvider::chat`].
495    async fn chat_and_sampling(
496        &self,
497        messages: &[ChatMessage],
498        json_schema: Option<StructuredOutputFormat>,
499        sampling: Option<&SamplingOverrides>,
500    ) -> Result<Box<dyn ChatResponse>, LLMError> {
501        self.chat_with_tools_and_sampling(messages, None, json_schema, sampling)
502            .await
503    }
504
505    /// Sends a chat request with tools and optional per-call sampling overrides.
506    ///
507    /// Equivalent to [`ChatProvider::chat_with_tools`] when `sampling` is
508    /// `None`. When `sampling` is `Some(...)`, backends that support per-call
509    /// sampling apply the overrides on top of provider-construction defaults;
510    /// backends that do not (the default impl) silently ignore them.
511    ///
512    /// Backwards compatible: the default implementation delegates to
513    /// [`ChatProvider::chat_with_tools`], dropping `sampling` on the floor.
514    /// Backends that wish to honour per-call sampling override this method.
515    async fn chat_with_tools_and_sampling(
516        &self,
517        messages: &[ChatMessage],
518        tools: Option<&[Tool]>,
519        json_schema: Option<StructuredOutputFormat>,
520        sampling: Option<&SamplingOverrides>,
521    ) -> Result<Box<dyn ChatResponse>, LLMError> {
522        // Default impl: ignore sampling, delegate to existing chat_with_tools.
523        let _ = sampling;
524        self.chat_with_tools(messages, tools, json_schema).await
525    }
526
527    /// Sends a chat with web search request to the provider
528    ///
529    /// # Arguments
530    ///
531    /// * `input` - The input message
532    ///
533    /// # Returns
534    ///
535    /// The provider's response text or an error
536    async fn chat_with_web_search(
537        &self,
538        _input: String,
539    ) -> Result<Box<dyn ChatResponse>, LLMError> {
540        Err(LLMError::Generic(
541            "Web search not supported for this provider".to_string(),
542        ))
543    }
544
545    /// Sends a streaming chat request to the provider with a sequence of messages.
546    ///
547    /// # Arguments
548    ///
549    /// * `messages` - The conversation history as a slice of chat messages
550    /// * `json_schema` - Optional json_schema for the response format
551    ///
552    /// # Returns
553    ///
554    /// A stream of text tokens or an error
555    async fn chat_stream(
556        &self,
557        _messages: &[ChatMessage],
558        _json_schema: Option<StructuredOutputFormat>,
559    ) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>, LLMError>
560    {
561        Err(LLMError::Generic(
562            "Streaming not supported for this provider".to_string(),
563        ))
564    }
565
566    /// Sends a streaming chat request that returns structured response chunks.
567    ///
568    /// ⚠️ Getting usage metadata while streaming have been noticed to be a unstable depending on the provider
569    /// (it can be missing).
570    ///
571    /// This method returns a stream of `StreamResponse` objects that mimic OpenAI's
572    /// streaming response format with `.choices[0].delta.content` and `.usage`.
573    ///
574    /// # Arguments
575    ///
576    /// * `messages` - The conversation history as a slice of chat messages
577    /// * `tools` - Optional slice of tools to use in the chat
578    /// * `json_schema` - Optional json_schema for the response format
579    ///
580    /// # Returns
581    ///
582    /// A stream of `StreamResponse` objects or an error
583    async fn chat_stream_struct(
584        &self,
585        _messages: &[ChatMessage],
586        _tools: Option<&[Tool]>,
587        _json_schema: Option<StructuredOutputFormat>,
588    ) -> Result<
589        std::pin::Pin<Box<dyn Stream<Item = Result<StreamResponse, LLMError>> + Send>>,
590        LLMError,
591    > {
592        Err(LLMError::Generic(
593            "Structured streaming not supported for this provider".to_string(),
594        ))
595    }
596
597    /// Sends a streaming chat request with tool support.
598    ///
599    /// Returns a stream of `StreamChunk` which can be text deltas or tool call events.
600    /// When `stop_reason` is "tool_use", the caller should execute the tool(s)
601    /// and continue the conversation.
602    ///
603    /// This method is ideal for agentic workflows where you want to stream text
604    /// output to the user while still receiving tool call requests.
605    ///
606    /// # Arguments
607    ///
608    /// * `messages` - The conversation history as a slice of chat messages
609    /// * `tools` - Optional slice of tools available for the model to use
610    /// * `json_schema` - Optional json_schema for the response format
611    ///
612    /// # Returns
613    ///
614    /// A stream of `StreamChunk` items or an error
615    ///
616    /// # Example
617    ///
618    /// ```ignore
619    /// use futures::StreamExt;
620    ///
621    /// let mut stream = client
622    ///     .chat_stream_with_tools(&messages, Some(&tools))
623    ///     .await?;
624    ///
625    /// let mut tool_calls = Vec::new();
626    /// while let Some(chunk) = stream.next().await {
627    ///     match chunk? {
628    ///         StreamChunk::Text(text) => print!("{}", text),
629    ///         StreamChunk::ToolUseComplete { tool_call, .. } => {
630    ///             tool_calls.push(tool_call);
631    ///         }
632    ///         StreamChunk::Done { stop_reason } => {
633    ///             if stop_reason == "tool_use" {
634    ///                 // Execute tool_calls and continue conversation
635    ///             }
636    ///         }
637    ///         _ => {}
638    ///     }
639    /// }
640    /// ```
641    async fn chat_stream_with_tools(
642        &self,
643        _messages: &[ChatMessage],
644        _tools: Option<&[Tool]>,
645        _json_schema: Option<StructuredOutputFormat>,
646    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send>>, LLMError> {
647        Err(LLMError::Generic(
648            "Streaming with tools not supported for this provider".to_string(),
649        ))
650    }
651
652    /// Streaming variant of [`ChatProvider::chat_stream`] with per-call
653    /// sampling overrides. Default impl ignores `sampling` and delegates to
654    /// [`ChatProvider::chat_stream`]. Backends that honour per-call sampling
655    /// override this method.
656    async fn chat_stream_and_sampling(
657        &self,
658        messages: &[ChatMessage],
659        json_schema: Option<StructuredOutputFormat>,
660        sampling: Option<&SamplingOverrides>,
661    ) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>, LLMError>
662    {
663        let _ = sampling;
664        self.chat_stream(messages, json_schema).await
665    }
666
667    /// Streaming variant of [`ChatProvider::chat_stream_struct`] with per-call
668    /// sampling overrides. Default impl ignores `sampling` and delegates to
669    /// [`ChatProvider::chat_stream_struct`]. Backends that honour per-call
670    /// sampling override this method.
671    async fn chat_stream_struct_and_sampling(
672        &self,
673        messages: &[ChatMessage],
674        tools: Option<&[Tool]>,
675        json_schema: Option<StructuredOutputFormat>,
676        sampling: Option<&SamplingOverrides>,
677    ) -> Result<
678        std::pin::Pin<Box<dyn Stream<Item = Result<StreamResponse, LLMError>> + Send>>,
679        LLMError,
680    > {
681        let _ = sampling;
682        self.chat_stream_struct(messages, tools, json_schema).await
683    }
684
685    /// Returns the model identifier this provider was configured with.
686    ///
687    /// Default returns an empty string for backwards compatibility with impls
688    /// that predate this trait method. Concrete backends should override to
689    /// return their configured model string so consumers can route requests
690    /// based on backend model identity (e.g. selecting grammar-based vs
691    /// prompt-based structured output, or capability-aware fallback ladders).
692    fn model(&self) -> &str {
693        ""
694    }
695}
696
697impl fmt::Display for ReasoningEffort {
698    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
699        match self {
700            ReasoningEffort::Low => write!(f, "low"),
701            ReasoningEffort::Medium => write!(f, "medium"),
702            ReasoningEffort::High => write!(f, "high"),
703        }
704    }
705}
706
707impl ChatMessage {
708    /// Create a new builder for a user message
709    pub fn user() -> ChatMessageBuilder {
710        ChatMessageBuilder::new(ChatRole::User)
711    }
712
713    /// Create a new builder for an assistant message
714    pub fn assistant() -> ChatMessageBuilder {
715        ChatMessageBuilder::new(ChatRole::Assistant)
716    }
717}
718
719/// Builder for ChatMessage
720#[derive(Debug)]
721pub struct ChatMessageBuilder {
722    role: ChatRole,
723    message_type: MessageType,
724    content: String,
725}
726
727impl ChatMessageBuilder {
728    /// Create a new ChatMessageBuilder with specified role
729    pub fn new(role: ChatRole) -> Self {
730        Self {
731            role,
732            message_type: MessageType::default(),
733            content: String::default(),
734        }
735    }
736
737    /// Set the message content
738    pub fn content<S: Into<String>>(mut self, content: S) -> Self {
739        self.content = content.into();
740        self
741    }
742
743    /// Set the message type as Image
744    pub fn image(mut self, image_mime: ImageMime, raw_bytes: Vec<u8>) -> Self {
745        self.message_type = MessageType::Image((image_mime, raw_bytes));
746        self
747    }
748
749    /// Set the message type as Image
750    pub fn pdf(mut self, raw_bytes: Vec<u8>) -> Self {
751        self.message_type = MessageType::Pdf(raw_bytes);
752        self
753    }
754
755    /// Set the message type as ImageURL
756    pub fn image_url(mut self, url: impl Into<String>) -> Self {
757        self.message_type = MessageType::ImageURL(url.into());
758        self
759    }
760
761    /// Set the message type as ToolUse
762    pub fn tool_use(mut self, tools: Vec<ToolCall>) -> Self {
763        self.message_type = MessageType::ToolUse(tools);
764        self
765    }
766
767    /// Set the message type as ToolResult
768    pub fn tool_result(mut self, tools: Vec<ToolCall>) -> Self {
769        self.message_type = MessageType::ToolResult(tools);
770        self
771    }
772
773    /// Build the ChatMessage
774    pub fn build(self) -> ChatMessage {
775        ChatMessage {
776            role: self.role,
777            message_type: self.message_type,
778            content: self.content,
779        }
780    }
781}
782
783/// Creates a Server-Sent Events (SSE) stream from an HTTP response.
784///
785/// # Arguments
786///
787/// * `response` - The HTTP response from the streaming API
788/// * `parser` - Function to parse each SSE chunk into optional text content
789///
790/// # Returns
791///
792/// A pinned stream of text tokens or an error
793#[cfg(not(target_arch = "wasm32"))]
794#[allow(dead_code)]
795pub(crate) fn create_sse_stream<F>(
796    response: reqwest::Response,
797    parser: F,
798) -> std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>
799where
800    F: Fn(&str) -> Result<Option<String>, LLMError> + Send + 'static,
801{
802    let stream = response
803        .bytes_stream()
804        .scan(
805            (String::default(), Vec::default()),
806            move |(buffer, utf8_buffer): &mut (String, Vec<u8>),
807                  chunk: Result<bytes::Bytes, reqwest::Error>| {
808                let result = match chunk {
809                    Ok(bytes) => {
810                        utf8_buffer.extend_from_slice(&bytes);
811
812                        match String::from_utf8(utf8_buffer.clone()) {
813                            Ok(text) => {
814                                buffer.push_str(&text);
815                                utf8_buffer.clear();
816                            }
817                            Err(e) => {
818                                let valid_up_to = e.utf8_error().valid_up_to();
819                                if valid_up_to > 0 {
820                                    // Safe to use from_utf8_lossy here since valid_up_to points to
821                                    // a valid UTF-8 boundary - no replacement characters will be introduced
822                                    let valid =
823                                        String::from_utf8_lossy(&utf8_buffer[..valid_up_to]);
824                                    buffer.push_str(&valid);
825                                    utf8_buffer.drain(..valid_up_to);
826                                }
827                            }
828                        }
829
830                        let mut results = Vec::default();
831
832                        while let Some(pos) = buffer.find("\n\n") {
833                            let event = buffer[..pos + 2].to_string();
834                            buffer.drain(..pos + 2);
835
836                            match parser(&event) {
837                                Ok(Some(content)) => results.push(Ok(content)),
838                                Ok(None) => {}
839                                Err(e) => results.push(Err(e)),
840                            }
841                        }
842
843                        Some(results)
844                    }
845                    Err(e) => Some(vec![Err(LLMError::HttpError(e.to_string()))]),
846                };
847
848                async move { result }
849            },
850        )
851        .flat_map(futures::stream::iter);
852
853    Box::pin(stream)
854}
855
856#[cfg(test)]
857mod tests {
858    use super::*;
859    use bytes::Bytes;
860    use futures::stream::StreamExt;
861
862    #[test]
863    fn test_chat_message_builder_user() {
864        let msg = ChatMessage::user().content("hello").build();
865        assert_eq!(msg.role, ChatRole::User);
866        assert_eq!(msg.content, "hello");
867        assert!(matches!(msg.message_type, MessageType::Text));
868    }
869
870    #[test]
871    fn test_chat_message_builder_assistant() {
872        let msg = ChatMessage::assistant().content("reply").build();
873        assert_eq!(msg.role, ChatRole::Assistant);
874        assert_eq!(msg.content, "reply");
875    }
876
877    #[test]
878    fn test_chat_message_builder_image() {
879        let msg = ChatMessage::user()
880            .content("describe")
881            .image(ImageMime::PNG, vec![1, 2, 3])
882            .build();
883        assert!(matches!(msg.message_type, MessageType::Image(_)));
884    }
885
886    #[test]
887    fn test_chat_message_builder_pdf() {
888        let msg = ChatMessage::user()
889            .content("read")
890            .pdf(vec![4, 5, 6])
891            .build();
892        assert!(matches!(msg.message_type, MessageType::Pdf(_)));
893    }
894
895    #[test]
896    fn test_chat_message_builder_tool_use() {
897        let tc = crate::ToolCall {
898            id: "t1".to_string(),
899            call_type: "function".to_string(),
900            function: crate::FunctionCall {
901                name: "tool".to_string(),
902                arguments: "{}".to_string(),
903            },
904        };
905        let msg = ChatMessage::assistant()
906            .content("calling tool")
907            .tool_use(vec![tc])
908            .build();
909        assert!(matches!(msg.message_type, MessageType::ToolUse(_)));
910    }
911
912    #[test]
913    fn test_chat_message_builder_tool_result() {
914        let tc = crate::ToolCall {
915            id: "t1".to_string(),
916            call_type: "function".to_string(),
917            function: crate::FunctionCall {
918                name: "tool".to_string(),
919                arguments: "result".to_string(),
920            },
921        };
922        let msg = ChatMessageBuilder::new(ChatRole::Tool)
923            .tool_result(vec![tc])
924            .build();
925        assert!(matches!(msg.message_type, MessageType::ToolResult(_)));
926        assert_eq!(msg.role, ChatRole::Tool);
927    }
928
929    #[test]
930    fn test_chat_role_display() {
931        assert_eq!(format!("{}", ChatRole::System), "system");
932        assert_eq!(format!("{}", ChatRole::User), "user");
933        assert_eq!(format!("{}", ChatRole::Assistant), "assistant");
934        assert_eq!(format!("{}", ChatRole::Tool), "tool");
935    }
936
937    #[test]
938    fn test_image_mime_mime_type() {
939        assert_eq!(ImageMime::JPEG.mime_type(), "image/jpeg");
940        assert_eq!(ImageMime::PNG.mime_type(), "image/png");
941        assert_eq!(ImageMime::GIF.mime_type(), "image/gif");
942        assert_eq!(ImageMime::WEBP.mime_type(), "image/webp");
943    }
944
945    #[test]
946    fn test_reasoning_effort_display() {
947        assert_eq!(format!("{}", ReasoningEffort::Low), "low");
948        assert_eq!(format!("{}", ReasoningEffort::Medium), "medium");
949        assert_eq!(format!("{}", ReasoningEffort::High), "high");
950    }
951
952    #[test]
953    fn test_tool_choice_serialization() {
954        let any_json = serde_json::to_value(&ToolChoice::Any).unwrap();
955        assert_eq!(any_json, "required");
956
957        let auto_json = serde_json::to_value(&ToolChoice::Auto).unwrap();
958        assert_eq!(auto_json, "auto");
959
960        let none_json = serde_json::to_value(&ToolChoice::None).unwrap();
961        assert_eq!(none_json, "none");
962
963        let tool_json = serde_json::to_value(ToolChoice::Tool("my_func".to_string())).unwrap();
964        assert_eq!(tool_json["type"], "function");
965        assert_eq!(tool_json["function"]["name"], "my_func");
966    }
967
968    #[test]
969    fn test_structured_output_format_roundtrip() {
970        let format = StructuredOutputFormat {
971            name: "Test".to_string(),
972            description: Some("A test".to_string()),
973            schema: Some(serde_json::json!({"type": "object"})),
974            strict: Some(true),
975        };
976        let json = serde_json::to_string(&format).unwrap();
977        let parsed: StructuredOutputFormat = serde_json::from_str(&json).unwrap();
978        assert_eq!(parsed, format);
979    }
980
981    #[test]
982    fn test_structured_output_format_minimal() {
983        let json_str = r#"{"name":"Minimal"}"#;
984        let parsed: StructuredOutputFormat = serde_json::from_str(json_str).unwrap();
985        assert_eq!(parsed.name, "Minimal");
986        assert_eq!(parsed.description, None);
987        assert_eq!(parsed.schema, None);
988        assert_eq!(parsed.strict, None);
989    }
990
991    #[test]
992    fn test_chat_message_builder_image_url() {
993        let msg = ChatMessage::user()
994            .image_url("https://example.com/img.png")
995            .content("describe this")
996            .build();
997        assert!(matches!(msg.message_type, MessageType::ImageURL(_)));
998    }
999
1000    #[tokio::test]
1001    async fn test_create_sse_stream_handles_split_utf8() {
1002        let test_data = "data: Positive reactions\n\n".as_bytes();
1003
1004        let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
1005            Ok(Bytes::from(&test_data[..10])),
1006            Ok(Bytes::from(&test_data[10..])),
1007        ];
1008
1009        let mock_response = create_mock_response(chunks);
1010
1011        let parser = |event: &str| -> Result<Option<String>, LLMError> {
1012            if let Some(content) = event.strip_prefix("data: ") {
1013                let content = content.trim();
1014                if content.is_empty() {
1015                    return Ok(None);
1016                }
1017                Ok(Some(content.to_string()))
1018            } else {
1019                Ok(None)
1020            }
1021        };
1022
1023        let mut stream = create_sse_stream(mock_response, parser);
1024
1025        let mut results = Vec::new();
1026        while let Some(result) = stream.next().await {
1027            results.push(result);
1028        }
1029
1030        assert_eq!(results.len(), 1);
1031        assert_eq!(results[0].as_ref().unwrap(), "Positive reactions");
1032    }
1033
1034    #[tokio::test]
1035    async fn test_create_sse_stream_handles_split_sse_events() {
1036        let event1 = "data: First event\n\n";
1037        let event2 = "data: Second event\n\n";
1038        let combined = format!("{}{}", event1, event2);
1039        let test_data = combined.as_bytes().to_vec();
1040
1041        let split_point = event1.len() + 5;
1042        let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
1043            Ok(Bytes::from(test_data[..split_point].to_vec())),
1044            Ok(Bytes::from(test_data[split_point..].to_vec())),
1045        ];
1046
1047        let mock_response = create_mock_response(chunks);
1048
1049        let parser = |event: &str| -> Result<Option<String>, LLMError> {
1050            if let Some(content) = event.strip_prefix("data: ") {
1051                let content = content.trim();
1052                if content.is_empty() {
1053                    return Ok(None);
1054                }
1055                Ok(Some(content.to_string()))
1056            } else {
1057                Ok(None)
1058            }
1059        };
1060
1061        let mut stream = create_sse_stream(mock_response, parser);
1062
1063        let mut results = Vec::new();
1064        while let Some(result) = stream.next().await {
1065            results.push(result);
1066        }
1067
1068        assert_eq!(results.len(), 2);
1069        assert_eq!(results[0].as_ref().unwrap(), "First event");
1070        assert_eq!(results[1].as_ref().unwrap(), "Second event");
1071    }
1072
1073    #[tokio::test]
1074    async fn test_create_sse_stream_handles_multibyte_utf8_split() {
1075        let multibyte_char = "✨";
1076        let event = format!("data: Star {}\n\n", multibyte_char);
1077        let test_data = event.as_bytes().to_vec();
1078
1079        let emoji_start = event.find(multibyte_char).unwrap();
1080        let split_in_emoji = emoji_start + 1;
1081
1082        let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
1083            Ok(Bytes::from(test_data[..split_in_emoji].to_vec())),
1084            Ok(Bytes::from(test_data[split_in_emoji..].to_vec())),
1085        ];
1086
1087        let mock_response = create_mock_response(chunks);
1088
1089        let parser = |event: &str| -> Result<Option<String>, LLMError> {
1090            if let Some(content) = event.strip_prefix("data: ") {
1091                let content = content.trim();
1092                if content.is_empty() {
1093                    return Ok(None);
1094                }
1095                Ok(Some(content.to_string()))
1096            } else {
1097                Ok(None)
1098            }
1099        };
1100
1101        let mut stream = create_sse_stream(mock_response, parser);
1102
1103        let mut results = Vec::new();
1104        while let Some(result) = stream.next().await {
1105            results.push(result);
1106        }
1107
1108        assert_eq!(results.len(), 1);
1109        assert_eq!(
1110            results[0].as_ref().unwrap(),
1111            &format!("Star {}", multibyte_char)
1112        );
1113    }
1114
1115    fn create_mock_response(chunks: Vec<Result<Bytes, reqwest::Error>>) -> reqwest::Response {
1116        use http_body_util::StreamBody;
1117        use reqwest::Body;
1118
1119        let frame_stream = futures::stream::iter(
1120            chunks
1121                .into_iter()
1122                .map(|chunk| chunk.map(hyper::body::Frame::data)),
1123        );
1124
1125        let body = StreamBody::new(frame_stream);
1126        let body = Body::wrap(body);
1127
1128        let http_response = http::Response::builder().status(200).body(body).unwrap();
1129
1130        http_response.into()
1131    }
1132}
1133
1134/// Tests for `ChatProvider::fn model(&self) -> &str`.
1135#[cfg(test)]
1136mod model_accessor_tests {
1137    use super::*;
1138
1139    /// Default impl returns empty string for impls that don't override.
1140    /// A minimal mock that only satisfies `chat_with_tools` gets `""` for free.
1141    #[test]
1142    fn default_impl_returns_empty_string() {
1143        struct MinimalMock;
1144        #[async_trait]
1145        impl ChatProvider for MinimalMock {
1146            async fn chat_with_tools(
1147                &self,
1148                _messages: &[ChatMessage],
1149                _tools: Option<&[Tool]>,
1150                _json_schema: Option<StructuredOutputFormat>,
1151            ) -> Result<Box<dyn ChatResponse>, crate::error::LLMError> {
1152                unimplemented!()
1153            }
1154        }
1155        let mock = MinimalMock;
1156        assert_eq!(mock.model(), "");
1157    }
1158
1159    /// Concrete Ollama backend exposes its configured model string.
1160    #[cfg(all(feature = "ollama", not(target_arch = "wasm32")))]
1161    #[test]
1162    fn ollama_backend_exposes_model_string() {
1163        let ollama = crate::backends::ollama::Ollama::new(
1164            "http://localhost:11434",        // base_url
1165            None,                            // api_key
1166            Some("qwen2.5:14b".to_string()), // model
1167            None,                            // max_tokens
1168            None,                            // temperature
1169            None,                            // timeout_seconds
1170            None,                            // top_p
1171            None,                            // top_k
1172            None,                            // keep_alive
1173            None,                            // system
1174            None,                            // think
1175            None,                            // stop
1176            None,                            // seed
1177            None,                            // presence_penalty
1178            None,                            // frequency_penalty
1179            None,                            // num_ctx
1180            None,                            // repeat_penalty
1181            None,                            // repeat_last_n
1182            None,                            // min_p
1183        );
1184        assert_eq!(ollama.model(), "qwen2.5:14b");
1185    }
1186
1187    /// Concrete Anthropic backend exposes its configured model string.
1188    #[cfg(all(feature = "anthropic", not(target_arch = "wasm32")))]
1189    #[test]
1190    fn anthropic_backend_exposes_model_string() {
1191        let anthropic = crate::backends::anthropic::Anthropic::new(
1192            "test-key",                                    // api_key
1193            Some("claude-haiku-4-5-20251001".to_string()), // model
1194            None,                                          // max_tokens
1195            None,                                          // temperature
1196            None,                                          // timeout_seconds
1197            None,                                          // top_p
1198            None,                                          // top_k
1199            None,                                          // tool_choice
1200            None,                                          // reasoning
1201            None,                                          // thinking_budget_tokens
1202        );
1203        assert_eq!(anthropic.model(), "claude-haiku-4-5-20251001");
1204    }
1205
1206    /// `Arc<dyn ChatProvider>` dispatches `.model()` to the inner impl via
1207    /// `Deref` coercion. Note: there is no blanket `impl ChatProvider for Arc<T>` —
1208    /// this test exercises method dispatch on `dyn ChatProvider` through `Arc`,
1209    /// not a trait impl on `Arc<T>` itself.
1210    #[cfg(all(feature = "ollama", not(target_arch = "wasm32")))]
1211    #[test]
1212    fn arc_dyn_chat_provider_dispatches_model_via_deref() {
1213        use std::sync::Arc;
1214        let ollama = crate::backends::ollama::Ollama::new(
1215            "http://localhost:11434",
1216            None,
1217            Some("qwen2.5:14b".to_string()),
1218            None,
1219            None,
1220            None,
1221            None,
1222            None,
1223            None,
1224            None,
1225            None,
1226            None,
1227            None,
1228            None,
1229            None,
1230            None,
1231            None,
1232            None,
1233            None,
1234        );
1235        let arc: Arc<dyn ChatProvider> = Arc::new(ollama);
1236        assert_eq!(arc.model(), "qwen2.5:14b");
1237    }
1238
1239    /// Integration test: ChatProvider::model() returns the configured string AND a subsequent
1240    /// chat_with_tools call observes the same model in the outgoing request body.
1241    ///
1242    /// Uses httpmock so no live Ollama server is required. The mock asserts that
1243    /// the POST body contains the model name reported by `.model()`, closing the
1244    /// loop between the accessor and the actual wire format. Gated only by
1245    /// `#[ignore]` — run with `cargo test -- --ignored`.
1246    ///
1247    /// Run manually:
1248    /// ```sh
1249    /// cargo test -p autoagents-llm --features ollama \
1250    ///   --lib -- model_accessor_tests::model_accessor_wires_to_chat_request --ignored --nocapture
1251    /// ```
1252    #[cfg(all(feature = "ollama", not(target_arch = "wasm32")))]
1253    #[tokio::test]
1254    #[ignore]
1255    async fn model_accessor_wires_to_chat_request() {
1256        use httpmock::{Method::POST, MockServer};
1257        use serde_json::json;
1258
1259        let configured_model = "qwen2.5:14b";
1260        let server = MockServer::start();
1261
1262        let provider = crate::backends::ollama::Ollama::new(
1263            server.base_url(),
1264            None,
1265            Some(configured_model.to_string()),
1266            Some(128),
1267            Some(0.0),
1268            None,
1269            None,
1270            None,
1271            None,
1272            None,
1273            None,
1274            None,
1275            None,
1276            None,
1277            None,
1278            None,
1279            None,
1280            None,
1281            None,
1282        );
1283
1284        // model() returns the configured string.
1285        assert_eq!(provider.model(), configured_model);
1286
1287        // Set up mock to verify the same model string appears in the wire request.
1288        let model_in_body = format!("\"model\":\"{configured_model}\"");
1289        let chat_mock = server.mock(|when, then| {
1290            when.method(POST)
1291                .path("/api/chat")
1292                .body_includes(model_in_body.as_str());
1293            then.status(200).json_body(json!({
1294                "message": {
1295                    "content": "mock reply",
1296                    "tool_calls": null
1297                }
1298            }));
1299        });
1300
1301        let messages = vec![ChatMessage::user().content("ping").build()];
1302        let response = provider
1303            .chat_with_tools(&messages, None, None)
1304            .await
1305            .expect("Mock-backed chat_with_tools must succeed");
1306
1307        // Response comes back correctly (proves the full call path ran).
1308        assert!(response.text().is_some(), "Response must contain text");
1309
1310        // Mock was hit exactly once — the model string reached the wire.
1311        chat_mock.assert();
1312    }
1313}
autoagents_llm/chat/mod.rs

autoagents_llm/chat/
mod.rs
