use serde::{Deserialize, Serialize};
use crate::MessageId;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SmartComposeRequest {
pub prompt: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub message_id: Option<MessageId>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
}
impl SmartComposeRequest {
pub fn new(prompt: impl Into<String>) -> Self {
Self {
prompt: prompt.into(),
message_id: None,
temperature: None,
}
}
pub fn reply_to(prompt: impl Into<String>, message_id: MessageId) -> Self {
Self {
prompt: prompt.into(),
message_id: Some(message_id),
temperature: None,
}
}
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.temperature = Some(temperature);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SmartComposeResponse {
pub suggestion: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SmartComposeStreamEvent {
#[serde(rename = "type")]
pub event_type: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_smart_compose_request_new() {
let request = SmartComposeRequest::new("Write a professional email");
assert_eq!(request.prompt, "Write a professional email");
assert_eq!(request.message_id, None);
assert_eq!(request.temperature, None);
}
#[test]
fn test_smart_compose_request_reply_to() {
let message_id = MessageId::new("msg_123");
let request = SmartComposeRequest::reply_to("Write a polite decline", message_id.clone());
assert_eq!(request.prompt, "Write a polite decline");
assert_eq!(request.message_id, Some(message_id));
assert_eq!(request.temperature, None);
}
#[test]
fn test_smart_compose_request_with_temperature() {
let request = SmartComposeRequest::new("Generate email").with_temperature(0.7);
assert_eq!(request.temperature, Some(0.7));
}
#[test]
fn test_smart_compose_request_serialization() {
let request = SmartComposeRequest::new("Test prompt").with_temperature(0.5);
let json = serde_json::to_string(&request).unwrap();
assert!(json.contains("Test prompt"));
assert!(json.contains("0.5"));
let deserialized: SmartComposeRequest = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.prompt, "Test prompt");
assert_eq!(deserialized.temperature, Some(0.5));
}
#[test]
fn test_smart_compose_response_deserialization() {
let json = r#"{"suggestion":"Dear John,\n\nI hope this finds you well."}"#;
let response: SmartComposeResponse = serde_json::from_str(json).unwrap();
assert!(response.suggestion.contains("Dear John"));
}
#[test]
fn test_smart_compose_stream_event_content() {
let json = r#"{"type":"content","content":"Hello "}"#;
let event: SmartComposeStreamEvent = serde_json::from_str(json).unwrap();
assert_eq!(event.event_type, "content");
assert_eq!(event.content, Some("Hello ".to_string()));
assert_eq!(event.error, None);
}
#[test]
fn test_smart_compose_stream_event_done() {
let json = r#"{"type":"done"}"#;
let event: SmartComposeStreamEvent = serde_json::from_str(json).unwrap();
assert_eq!(event.event_type, "done");
assert_eq!(event.content, None);
assert_eq!(event.error, None);
}
#[test]
fn test_smart_compose_stream_event_error() {
let json = r#"{"type":"error","error":"Rate limit exceeded"}"#;
let event: SmartComposeStreamEvent = serde_json::from_str(json).unwrap();
assert_eq!(event.event_type, "error");
assert_eq!(event.error, Some("Rate limit exceeded".to_string()));
assert_eq!(event.content, None);
}
}