autoagents_llm/chat/
mod.rs

1use std::collections::HashMap;
2use std::fmt;
3
4use async_trait::async_trait;
5use futures::stream::{Stream, StreamExt};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9use crate::{error::LLMError, ToolCall};
10
11/// Role of a participant in a chat conversation.
12#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
13pub enum ChatRole {
14    /// System
15    System,
16    /// The user/human participant in the conversation
17    User,
18    /// The AI assistant participant in the conversation
19    Assistant,
20    /// Tool/function response
21    Tool,
22}
23
24/// The supported MIME type of an image.
25#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
26#[non_exhaustive]
27pub enum ImageMime {
28    /// JPEG image
29    JPEG,
30    /// PNG image
31    PNG,
32    /// GIF image
33    GIF,
34    /// WebP image
35    WEBP,
36}
37
38impl ImageMime {
39    pub fn mime_type(&self) -> &'static str {
40        match self {
41            ImageMime::JPEG => "image/jpeg",
42            ImageMime::PNG => "image/png",
43            ImageMime::GIF => "image/gif",
44            ImageMime::WEBP => "image/webp",
45        }
46    }
47}
48
49/// The type of a message in a chat conversation.
50#[derive(Debug, Clone, PartialEq, Eq, Default, Deserialize, Serialize)]
51pub enum MessageType {
52    /// A text message
53    #[default]
54    Text,
55    /// An image message
56    Image((ImageMime, Vec<u8>)),
57    /// PDF message
58    Pdf(Vec<u8>),
59    /// An image URL message
60    ImageURL(String),
61    /// A tool use
62    ToolUse(Vec<ToolCall>),
63    /// Tool result
64    ToolResult(Vec<ToolCall>),
65}
66
67/// The type of reasoning effort for a message in a chat conversation.
68pub enum ReasoningEffort {
69    /// Low reasoning effort
70    Low,
71    /// Medium reasoning effort
72    Medium,
73    /// High reasoning effort
74    High,
75}
76
77/// A single message in a chat conversation.
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct ChatMessage {
80    /// The role of who sent this message (user or assistant)
81    pub role: ChatRole,
82    /// The type of the message (text, image, audio, video, etc)
83    pub message_type: MessageType,
84    /// The text content of the message
85    pub content: String,
86}
87
88/// Represents a parameter in a function tool
89#[derive(Debug, Clone, Serialize)]
90pub struct ParameterProperty {
91    /// The type of the parameter (e.g. "string", "number", "array", etc)
92    #[serde(rename = "type")]
93    pub property_type: String,
94    /// Description of what the parameter does
95    pub description: String,
96    /// When type is "array", this defines the type of the array items
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub items: Option<Box<ParameterProperty>>,
99    /// When type is "enum", this defines the possible values for the parameter
100    #[serde(skip_serializing_if = "Option::is_none", rename = "enum")]
101    pub enum_list: Option<Vec<String>>,
102}
103
104/// Represents the parameters schema for a function tool
105#[derive(Debug, Clone, Serialize)]
106pub struct ParametersSchema {
107    /// The type of the parameters object (usually "object")
108    #[serde(rename = "type")]
109    pub schema_type: String,
110    /// Map of parameter names to their properties
111    pub properties: HashMap<String, ParameterProperty>,
112    /// List of required parameter names
113    pub required: Vec<String>,
114}
115
116/// Represents a function definition for a tool.
117///
118/// The `parameters` field stores the JSON Schema describing the function
119/// arguments.  It is kept as a raw `serde_json::Value` to allow arbitrary
120/// complexity (nested arrays/objects, `oneOf`, etc.) without requiring a
121/// bespoke Rust structure.
122///
123/// Builder helpers can still generate simple schemas automatically, but the
124/// user may also provide any valid schema directly.
125#[derive(Debug, Clone, Serialize)]
126pub struct FunctionTool {
127    /// Name of the function
128    pub name: String,
129    /// Human-readable description
130    pub description: String,
131    /// JSON Schema describing the parameters
132    pub parameters: Value,
133}
134
135/// 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).
136/// Individual providers may have additional requirements or restrictions, but these should be handled by each provider's backend implementation.
137///
138/// 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.
139///
140/// ## Example
141///
142/// ```
143/// use autoagents_llm::chat::StructuredOutputFormat;
144/// use serde_json::json;
145///
146/// let response_format = r#"
147///     {
148///         "name": "Student",
149///         "description": "A student object",
150///         "schema": {
151///             "type": "object",
152///             "properties": {
153///                 "name": {
154///                     "type": "string"
155///                 },
156///                 "age": {
157///                     "type": "integer"
158///                 },
159///                 "is_student": {
160///                     "type": "boolean"
161///                 }
162///             },
163///             "required": ["name", "age", "is_student"]
164///         }
165///     }
166/// "#;
167/// let structured_output: StructuredOutputFormat = serde_json::from_str(response_format).unwrap();
168/// assert_eq!(structured_output.name, "Student");
169/// assert_eq!(structured_output.description, Some("A student object".to_string()));
170/// ```
171#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
172
173pub struct StructuredOutputFormat {
174    /// Name of the schema
175    pub name: String,
176    /// The description of the schema
177    pub description: Option<String>,
178    /// The JSON schema for the structured output
179    pub schema: Option<Value>,
180    /// Whether to enable strict schema adherence
181    pub strict: Option<bool>,
182}
183
184/// Represents a tool that can be used in chat
185#[derive(Debug, Clone, Serialize)]
186pub struct Tool {
187    /// The type of tool (e.g. "function")
188    #[serde(rename = "type")]
189    pub tool_type: String,
190    /// The function definition if this is a function tool
191    pub function: FunctionTool,
192}
193
194/// Tool choice determines how the LLM uses available tools.
195/// The behavior is standardized across different LLM providers.
196#[derive(Debug, Clone, Default)]
197pub enum ToolChoice {
198    /// Model can use any tool, but it must use at least one.
199    /// This is useful when you want to force the model to use tools.
200    Any,
201
202    /// Model can use any tool, and may elect to use none.
203    /// This is the default behavior and gives the model flexibility.
204    #[default]
205    Auto,
206
207    /// Model must use the specified tool and only the specified tool.
208    /// The string parameter is the name of the required tool.
209    /// This is useful when you want the model to call a specific function.
210    Tool(String),
211
212    /// Explicitly disables the use of tools.
213    /// The model will not use any tools even if they are provided.
214    None,
215}
216
217impl Serialize for ToolChoice {
218    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
219    where
220        S: serde::Serializer,
221    {
222        match self {
223            ToolChoice::Any => serializer.serialize_str("required"),
224            ToolChoice::Auto => serializer.serialize_str("auto"),
225            ToolChoice::None => serializer.serialize_str("none"),
226            ToolChoice::Tool(name) => {
227                use serde::ser::SerializeMap;
228
229                // For tool_choice: {"type": "function", "function": {"name": "function_name"}}
230                let mut map = serializer.serialize_map(Some(2))?;
231                map.serialize_entry("type", "function")?;
232
233                // Inner function object
234                let mut function_obj = std::collections::HashMap::new();
235                function_obj.insert("name", name.as_str());
236
237                map.serialize_entry("function", &function_obj)?;
238                map.end()
239            }
240        }
241    }
242}
243
244pub trait ChatResponse: std::fmt::Debug + std::fmt::Display + Send + Sync {
245    fn text(&self) -> Option<String>;
246    fn tool_calls(&self) -> Option<Vec<ToolCall>>;
247    fn thinking(&self) -> Option<String> {
248        None
249    }
250}
251
252/// Trait for providers that support chat-style interactions.
253#[async_trait]
254pub trait ChatProvider: Sync + Send {
255    /// Sends a chat request to the provider with a sequence of messages.
256    ///
257    /// # Arguments
258    ///
259    /// * `messages` - The conversation history as a slice of chat messages
260    ///
261    /// # Returns
262    ///
263    /// The provider's response text or an error
264    async fn chat(
265        &self,
266        messages: &[ChatMessage],
267        json_schema: Option<StructuredOutputFormat>,
268    ) -> Result<Box<dyn ChatResponse>, LLMError> {
269        self.chat_with_tools(messages, None, json_schema).await
270    }
271
272    /// Sends a chat request to the provider with a sequence of messages and tools.
273    ///
274    /// # Arguments
275    ///
276    /// * `messages` - The conversation history as a slice of chat messages
277    /// * `tools` - Optional slice of tools to use in the chat
278    ///
279    /// # Returns
280    ///
281    /// The provider's response text or an error
282    async fn chat_with_tools(
283        &self,
284        messages: &[ChatMessage],
285        tools: Option<&[Tool]>,
286        json_schema: Option<StructuredOutputFormat>,
287    ) -> Result<Box<dyn ChatResponse>, LLMError>;
288
289    /// Sends a streaming chat request to the provider with a sequence of messages.
290    ///
291    /// # Arguments
292    ///
293    /// * `messages` - The conversation history as a slice of chat messages
294    ///
295    /// # Returns
296    ///
297    /// A stream of text tokens or an error
298    async fn chat_stream(
299        &self,
300        _messages: &[ChatMessage],
301    ) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>, LLMError>
302    {
303        Err(LLMError::Generic(
304            "Streaming not supported for this provider".to_string(),
305        ))
306    }
307
308    /// Get current memory contents if provider supports memory
309    async fn memory_contents(&self) -> Option<Vec<ChatMessage>> {
310        None
311    }
312
313    /// Summarizes a conversation history into a concise 2-3 sentence summary
314    ///
315    /// # Arguments
316    /// * `msgs` - The conversation messages to summarize
317    ///
318    /// # Returns
319    /// A string containing the summary or an error if summarization fails
320    async fn summarize_history(&self, msgs: &[ChatMessage]) -> Result<String, LLMError> {
321        let prompt = format!(
322            "Summarize in 2-3 sentences:\n{}",
323            msgs.iter()
324                .map(|m| format!("{:?}: {}", m.role, m.content))
325                .collect::<Vec<_>>()
326                .join("\n"),
327        );
328        let req = [ChatMessage::user().content(prompt).build()];
329        self.chat(&req, None)
330            .await?
331            .text()
332            .ok_or(LLMError::Generic("no text in summary response".into()))
333    }
334}
335
336impl fmt::Display for ReasoningEffort {
337    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
338        match self {
339            ReasoningEffort::Low => write!(f, "low"),
340            ReasoningEffort::Medium => write!(f, "medium"),
341            ReasoningEffort::High => write!(f, "high"),
342        }
343    }
344}
345
346impl ChatMessage {
347    /// Create a new builder for a user message
348    pub fn user() -> ChatMessageBuilder {
349        ChatMessageBuilder::new(ChatRole::User)
350    }
351
352    /// Create a new builder for an assistant message
353    pub fn assistant() -> ChatMessageBuilder {
354        ChatMessageBuilder::new(ChatRole::Assistant)
355    }
356}
357
358/// Builder for ChatMessage
359#[derive(Debug)]
360pub struct ChatMessageBuilder {
361    role: ChatRole,
362    message_type: MessageType,
363    content: String,
364}
365
366impl ChatMessageBuilder {
367    /// Create a new ChatMessageBuilder with specified role
368    pub fn new(role: ChatRole) -> Self {
369        Self {
370            role,
371            message_type: MessageType::default(),
372            content: String::new(),
373        }
374    }
375
376    /// Set the message content
377    pub fn content<S: Into<String>>(mut self, content: S) -> Self {
378        self.content = content.into();
379        self
380    }
381
382    /// Set the message type as Image
383    pub fn image(mut self, image_mime: ImageMime, raw_bytes: Vec<u8>) -> Self {
384        self.message_type = MessageType::Image((image_mime, raw_bytes));
385        self
386    }
387
388    /// Set the message type as Image
389    pub fn pdf(mut self, raw_bytes: Vec<u8>) -> Self {
390        self.message_type = MessageType::Pdf(raw_bytes);
391        self
392    }
393
394    /// Set the message type as ImageURL
395    pub fn image_url(mut self, url: impl Into<String>) -> Self {
396        self.message_type = MessageType::ImageURL(url.into());
397        self
398    }
399
400    /// Set the message type as ToolUse
401    pub fn tool_use(mut self, tools: Vec<ToolCall>) -> Self {
402        self.message_type = MessageType::ToolUse(tools);
403        self
404    }
405
406    /// Set the message type as ToolResult
407    pub fn tool_result(mut self, tools: Vec<ToolCall>) -> Self {
408        self.message_type = MessageType::ToolResult(tools);
409        self
410    }
411
412    /// Build the ChatMessage
413    pub fn build(self) -> ChatMessage {
414        ChatMessage {
415            role: self.role,
416            message_type: self.message_type,
417            content: self.content,
418        }
419    }
420}
421
422/// Creates a Server-Sent Events (SSE) stream from an HTTP response.
423///
424/// # Arguments
425///
426/// * `response` - The HTTP response from the streaming API
427/// * `parser` - Function to parse each SSE chunk into optional text content
428///
429/// # Returns
430///
431/// A pinned stream of text tokens or an error
432#[allow(dead_code)]
433pub(crate) fn create_sse_stream<F>(
434    response: reqwest::Response,
435    parser: F,
436) -> std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>
437where
438    F: Fn(&str) -> Result<Option<String>, LLMError> + Send + 'static,
439{
440    let stream = response
441        .bytes_stream()
442        .map(move |chunk| match chunk {
443            Ok(bytes) => {
444                let text = String::from_utf8_lossy(&bytes);
445                parser(&text)
446            }
447            Err(e) => Err(LLMError::HttpError(e.to_string())),
448        })
449        .filter_map(|result| async move {
450            match result {
451                Ok(Some(content)) => Some(Ok(content)),
452                Ok(None) => None,
453                Err(e) => Some(Err(e)),
454            }
455        });
456
457    Box::pin(stream)
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463    use crate::error::LLMError;
464    use serde_json::json;
465
466    #[test]
467    fn test_chat_role_serialization() {
468        let user_role = ChatRole::User;
469        let serialized = serde_json::to_string(&user_role).unwrap();
470        assert_eq!(serialized, "\"User\"");
471
472        let assistant_role = ChatRole::Assistant;
473        let serialized = serde_json::to_string(&assistant_role).unwrap();
474        assert_eq!(serialized, "\"Assistant\"");
475
476        let system_role = ChatRole::System;
477        let serialized = serde_json::to_string(&system_role).unwrap();
478        assert_eq!(serialized, "\"System\"");
479
480        let tool_role = ChatRole::Tool;
481        let serialized = serde_json::to_string(&tool_role).unwrap();
482        assert_eq!(serialized, "\"Tool\"");
483    }
484
485    #[test]
486    fn test_chat_role_deserialization() {
487        let deserialized: ChatRole = serde_json::from_str("\"User\"").unwrap();
488        assert_eq!(deserialized, ChatRole::User);
489
490        let deserialized: ChatRole = serde_json::from_str("\"Assistant\"").unwrap();
491        assert_eq!(deserialized, ChatRole::Assistant);
492
493        let deserialized: ChatRole = serde_json::from_str("\"System\"").unwrap();
494        assert_eq!(deserialized, ChatRole::System);
495
496        let deserialized: ChatRole = serde_json::from_str("\"Tool\"").unwrap();
497        assert_eq!(deserialized, ChatRole::Tool);
498    }
499
500    #[test]
501    fn test_image_mime_type() {
502        assert_eq!(ImageMime::JPEG.mime_type(), "image/jpeg");
503        assert_eq!(ImageMime::PNG.mime_type(), "image/png");
504        assert_eq!(ImageMime::GIF.mime_type(), "image/gif");
505        assert_eq!(ImageMime::WEBP.mime_type(), "image/webp");
506    }
507
508    #[test]
509    fn test_message_type_default() {
510        let default_type = MessageType::default();
511        assert_eq!(default_type, MessageType::Text);
512    }
513
514    #[test]
515    fn test_message_type_serialization() {
516        let text_type = MessageType::Text;
517        let serialized = serde_json::to_string(&text_type).unwrap();
518        assert_eq!(serialized, "\"Text\"");
519
520        let image_type = MessageType::Image((ImageMime::JPEG, vec![1, 2, 3]));
521        let serialized = serde_json::to_string(&image_type).unwrap();
522        assert!(serialized.contains("Image"));
523    }
524
525    #[test]
526    fn test_reasoning_effort_display() {
527        assert_eq!(ReasoningEffort::Low.to_string(), "low");
528        assert_eq!(ReasoningEffort::Medium.to_string(), "medium");
529        assert_eq!(ReasoningEffort::High.to_string(), "high");
530    }
531
532    #[test]
533    fn test_chat_message_builder_user() {
534        let message = ChatMessage::user().content("Hello, world!").build();
535
536        assert_eq!(message.role, ChatRole::User);
537        assert_eq!(message.content, "Hello, world!");
538        assert_eq!(message.message_type, MessageType::Text);
539    }
540
541    #[test]
542    fn test_chat_message_builder_assistant() {
543        let message = ChatMessage::assistant().content("Hi there!").build();
544
545        assert_eq!(message.role, ChatRole::Assistant);
546        assert_eq!(message.content, "Hi there!");
547        assert_eq!(message.message_type, MessageType::Text);
548    }
549
550    #[test]
551    fn test_chat_message_builder_image() {
552        let image_data = vec![1, 2, 3, 4, 5];
553        let message = ChatMessage::user()
554            .content("Check this image")
555            .image(ImageMime::PNG, image_data.clone())
556            .build();
557
558        assert_eq!(message.role, ChatRole::User);
559        assert_eq!(message.content, "Check this image");
560        assert_eq!(
561            message.message_type,
562            MessageType::Image((ImageMime::PNG, image_data))
563        );
564    }
565
566    #[test]
567    fn test_chat_message_builder_pdf() {
568        let pdf_data = vec![0x25, 0x50, 0x44, 0x46]; // PDF header
569        let message = ChatMessage::user()
570            .content("Review this PDF")
571            .pdf(pdf_data.clone())
572            .build();
573
574        assert_eq!(message.role, ChatRole::User);
575        assert_eq!(message.content, "Review this PDF");
576        assert_eq!(message.message_type, MessageType::Pdf(pdf_data));
577    }
578
579    #[test]
580    fn test_chat_message_builder_image_url() {
581        let image_url = "https://example.com/image.jpg";
582        let message = ChatMessage::user()
583            .content("See this image")
584            .image_url(image_url)
585            .build();
586
587        assert_eq!(message.role, ChatRole::User);
588        assert_eq!(message.content, "See this image");
589        assert_eq!(
590            message.message_type,
591            MessageType::ImageURL(image_url.to_string())
592        );
593    }
594
595    #[test]
596    fn test_chat_message_builder_tool_use() {
597        let tool_calls = vec![crate::ToolCall {
598            id: "call_1".to_string(),
599            call_type: "function".to_string(),
600            function: crate::FunctionCall {
601                name: "get_weather".to_string(),
602                arguments: "{\"location\": \"New York\"}".to_string(),
603            },
604        }];
605
606        let message = ChatMessage::assistant()
607            .content("Using weather tool")
608            .tool_use(tool_calls.clone())
609            .build();
610
611        assert_eq!(message.role, ChatRole::Assistant);
612        assert_eq!(message.content, "Using weather tool");
613        assert_eq!(message.message_type, MessageType::ToolUse(tool_calls));
614    }
615
616    #[test]
617    fn test_chat_message_builder_tool_result() {
618        let tool_results = vec![crate::ToolCall {
619            id: "call_1".to_string(),
620            call_type: "function".to_string(),
621            function: crate::FunctionCall {
622                name: "get_weather".to_string(),
623                arguments: "{\"temperature\": 75, \"condition\": \"sunny\"}".to_string(),
624            },
625        }];
626
627        let message = ChatMessage::user()
628            .content("Weather result")
629            .tool_result(tool_results.clone())
630            .build();
631
632        assert_eq!(message.role, ChatRole::User);
633        assert_eq!(message.content, "Weather result");
634        assert_eq!(message.message_type, MessageType::ToolResult(tool_results));
635    }
636
637    #[test]
638    fn test_structured_output_format_serialization() {
639        let format = StructuredOutputFormat {
640            name: "Person".to_string(),
641            description: Some("A person object".to_string()),
642            schema: Some(json!({
643                "type": "object",
644                "properties": {
645                    "name": {"type": "string"},
646                    "age": {"type": "integer"}
647                },
648                "required": ["name", "age"]
649            })),
650            strict: Some(true),
651        };
652
653        let serialized = serde_json::to_string(&format).unwrap();
654        let deserialized: StructuredOutputFormat = serde_json::from_str(&serialized).unwrap();
655
656        assert_eq!(deserialized.name, "Person");
657        assert_eq!(
658            deserialized.description,
659            Some("A person object".to_string())
660        );
661        assert_eq!(deserialized.strict, Some(true));
662        assert!(deserialized.schema.is_some());
663    }
664
665    #[test]
666    fn test_structured_output_format_equality() {
667        let format1 = StructuredOutputFormat {
668            name: "Test".to_string(),
669            description: None,
670            schema: None,
671            strict: None,
672        };
673
674        let format2 = StructuredOutputFormat {
675            name: "Test".to_string(),
676            description: None,
677            schema: None,
678            strict: None,
679        };
680
681        assert_eq!(format1, format2);
682    }
683
684    #[test]
685    fn test_tool_choice_serialization() {
686        // Test Auto
687        let choice = ToolChoice::Auto;
688        let serialized = serde_json::to_string(&choice).unwrap();
689        assert_eq!(serialized, "\"auto\"");
690
691        // Test Any
692        let choice = ToolChoice::Any;
693        let serialized = serde_json::to_string(&choice).unwrap();
694        assert_eq!(serialized, "\"required\"");
695
696        // Test None
697        let choice = ToolChoice::None;
698        let serialized = serde_json::to_string(&choice).unwrap();
699        assert_eq!(serialized, "\"none\"");
700
701        // Test specific tool
702        let choice = ToolChoice::Tool("my_function".to_string());
703        let serialized = serde_json::to_string(&choice).unwrap();
704        // Parse back to verify structure
705        let parsed: serde_json::Value = serde_json::from_str(&serialized).unwrap();
706        assert_eq!(parsed["type"], "function");
707        assert_eq!(parsed["function"]["name"], "my_function");
708    }
709
710    #[test]
711    fn test_tool_choice_default() {
712        let default_choice = ToolChoice::default();
713        assert!(matches!(default_choice, ToolChoice::Auto));
714    }
715
716    #[test]
717    fn test_parameter_property_serialization() {
718        let property = ParameterProperty {
719            property_type: "string".to_string(),
720            description: "A test parameter".to_string(),
721            items: None,
722            enum_list: Some(vec!["option1".to_string(), "option2".to_string()]),
723        };
724
725        let serialized = serde_json::to_string(&property).unwrap();
726        assert!(serialized.contains("\"type\":\"string\""));
727        assert!(serialized.contains("\"description\":\"A test parameter\""));
728        assert!(serialized.contains("\"enum\":[\"option1\",\"option2\"]"));
729    }
730
731    #[test]
732    fn test_parameter_property_with_items() {
733        let item_property = ParameterProperty {
734            property_type: "string".to_string(),
735            description: "Array item".to_string(),
736            items: None,
737            enum_list: None,
738        };
739
740        let array_property = ParameterProperty {
741            property_type: "array".to_string(),
742            description: "An array parameter".to_string(),
743            items: Some(Box::new(item_property)),
744            enum_list: None,
745        };
746
747        let serialized = serde_json::to_string(&array_property).unwrap();
748        assert!(serialized.contains("\"type\":\"array\""));
749        assert!(serialized.contains("\"items\""));
750    }
751
752    #[test]
753    fn test_function_tool_serialization() {
754        let mut properties = HashMap::new();
755        properties.insert(
756            "name".to_string(),
757            ParameterProperty {
758                property_type: "string".to_string(),
759                description: "Name parameter".to_string(),
760                items: None,
761                enum_list: None,
762            },
763        );
764
765        let schema = ParametersSchema {
766            schema_type: "object".to_string(),
767            properties,
768            required: vec!["name".to_string()],
769        };
770
771        let function = FunctionTool {
772            name: "test_function".to_string(),
773            description: "A test function".to_string(),
774            parameters: serde_json::to_value(schema).unwrap(),
775        };
776
777        let serialized = serde_json::to_string(&function).unwrap();
778        assert!(serialized.contains("\"name\":\"test_function\""));
779        assert!(serialized.contains("\"description\":\"A test function\""));
780        assert!(serialized.contains("\"parameters\""));
781    }
782
783    #[test]
784    fn test_tool_serialization() {
785        let function = FunctionTool {
786            name: "test_tool".to_string(),
787            description: "A test tool".to_string(),
788            parameters: json!({"type": "object", "properties": {}}),
789        };
790
791        let tool = Tool {
792            tool_type: "function".to_string(),
793            function,
794        };
795
796        let serialized = serde_json::to_string(&tool).unwrap();
797        assert!(serialized.contains("\"type\":\"function\""));
798        assert!(serialized.contains("\"function\""));
799    }
800
801    #[test]
802    fn test_chat_message_serialization() {
803        let message = ChatMessage {
804            role: ChatRole::User,
805            message_type: MessageType::Text,
806            content: "Hello, world!".to_string(),
807        };
808
809        let serialized = serde_json::to_string(&message).unwrap();
810        let deserialized: ChatMessage = serde_json::from_str(&serialized).unwrap();
811
812        assert_eq!(deserialized.role, ChatRole::User);
813        assert_eq!(deserialized.message_type, MessageType::Text);
814        assert_eq!(deserialized.content, "Hello, world!");
815    }
816
817    #[tokio::test]
818    async fn test_chat_provider_summarize_history() {
819        struct MockChatProvider;
820
821        #[async_trait]
822        impl ChatProvider for MockChatProvider {
823            async fn chat_with_tools(
824                &self,
825                messages: &[ChatMessage],
826                _tools: Option<&[Tool]>,
827                _json_schema: Option<StructuredOutputFormat>,
828            ) -> Result<Box<dyn ChatResponse>, LLMError> {
829                // Mock implementation that returns the prompt as summary
830                let prompt = messages.first().unwrap().content.clone();
831                Ok(Box::new(MockChatResponse {
832                    text: Some(format!("Summary: {prompt}")),
833                }))
834            }
835        }
836
837        struct MockChatResponse {
838            text: Option<String>,
839        }
840
841        impl ChatResponse for MockChatResponse {
842            fn text(&self) -> Option<String> {
843                self.text.clone()
844            }
845
846            fn tool_calls(&self) -> Option<Vec<crate::ToolCall>> {
847                None
848            }
849        }
850
851        impl std::fmt::Debug for MockChatResponse {
852            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
853                write!(f, "MockChatResponse")
854            }
855        }
856
857        impl std::fmt::Display for MockChatResponse {
858            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
859                write!(f, "{}", self.text.as_deref().unwrap_or(""))
860            }
861        }
862
863        let provider = MockChatProvider;
864        let messages = vec![
865            ChatMessage::user().content("First message").build(),
866            ChatMessage::assistant().content("Response").build(),
867        ];
868
869        let summary = provider.summarize_history(&messages).await.unwrap();
870        assert!(summary.contains("Summary:"));
871        assert!(summary.contains("First message"));
872        assert!(summary.contains("Response"));
873    }
874
875    #[tokio::test]
876    async fn test_chat_provider_default_chat_implementation() {
877        struct MockChatProvider;
878
879        #[async_trait]
880        impl ChatProvider for MockChatProvider {
881            async fn chat_with_tools(
882                &self,
883                _messages: &[ChatMessage],
884                _tools: Option<&[Tool]>,
885                _json_schema: Option<StructuredOutputFormat>,
886            ) -> Result<Box<dyn ChatResponse>, LLMError> {
887                Ok(Box::new(MockChatResponse {
888                    text: Some("Default response".to_string()),
889                }))
890            }
891        }
892
893        struct MockChatResponse {
894            text: Option<String>,
895        }
896
897        impl ChatResponse for MockChatResponse {
898            fn text(&self) -> Option<String> {
899                self.text.clone()
900            }
901
902            fn tool_calls(&self) -> Option<Vec<crate::ToolCall>> {
903                None
904            }
905        }
906
907        impl std::fmt::Debug for MockChatResponse {
908            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
909                write!(f, "MockChatResponse")
910            }
911        }
912
913        impl std::fmt::Display for MockChatResponse {
914            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
915                write!(f, "{}", self.text.as_deref().unwrap_or(""))
916            }
917        }
918
919        let provider = MockChatProvider;
920        let messages = vec![ChatMessage::user().content("Test").build()];
921
922        let response = provider.chat(&messages, None).await.unwrap();
923        assert_eq!(response.text(), Some("Default response".to_string()));
924    }
925
926    #[tokio::test]
927    async fn test_chat_provider_default_chat_stream_implementation() {
928        struct MockChatProvider;
929
930        #[async_trait]
931        impl ChatProvider for MockChatProvider {
932            async fn chat_with_tools(
933                &self,
934                _messages: &[ChatMessage],
935                _tools: Option<&[Tool]>,
936                _json_schema: Option<StructuredOutputFormat>,
937            ) -> Result<Box<dyn ChatResponse>, LLMError> {
938                unreachable!()
939            }
940        }
941
942        let provider = MockChatProvider;
943        let messages = vec![ChatMessage::user().content("Test").build()];
944
945        let result = provider.chat_stream(&messages).await;
946        assert!(result.is_err());
947        if let Err(error) = result {
948            assert!(error.to_string().contains("Streaming not supported"));
949        }
950    }
951
952    #[tokio::test]
953    async fn test_chat_provider_default_memory_contents() {
954        struct MockChatProvider;
955
956        #[async_trait]
957        impl ChatProvider for MockChatProvider {
958            async fn chat_with_tools(
959                &self,
960                _messages: &[ChatMessage],
961                _tools: Option<&[Tool]>,
962                _json_schema: Option<StructuredOutputFormat>,
963            ) -> Result<Box<dyn ChatResponse>, LLMError> {
964                unreachable!()
965            }
966        }
967
968        let provider = MockChatProvider;
969        let memory_contents = provider.memory_contents().await;
970        assert!(memory_contents.is_none());
971    }
972}
autoagents_llm/chat/mod.rs

autoagents_llm/chat/
mod.rs
