use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ChatMessage {
pub role: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
impl ChatMessage {
pub fn new(role: impl Into<String>, content: impl Into<String>) -> Self {
Self {
role: role.into(),
content: Some(content.into()),
tool_call_id: None,
tool_calls: None,
}
}
pub fn system(content: impl Into<String>) -> Self {
Self::new("system", content)
}
pub fn user(content: impl Into<String>) -> Self {
Self::new("user", content)
}
pub fn assistant(content: impl Into<String>) -> Self {
Self::new("assistant", content)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub call_type: String,
pub function: FunctionCall,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct FunctionCall {
pub name: String,
pub arguments: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct ChatRequest {
pub model: String,
pub messages: Vec<ChatMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub tools: Vec<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
}
impl ChatRequest {
pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
Self {
model: model.into(),
messages,
max_tokens: None,
temperature: None,
tools: Vec::new(),
stream: None,
}
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ChatResponse {
pub id: String,
pub choices: Vec<Choice>,
pub usage: Option<Usage>,
pub model: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Choice {
pub index: i32,
pub message: ChatMessage,
pub finish_reason: Option<String>,
}
pub use clawft_types::provider::Usage;
#[derive(Debug, Clone, PartialEq)]
pub enum StreamChunk {
TextDelta {
text: String,
},
ToolCallDelta {
index: usize,
id: Option<String>,
name: Option<String>,
arguments: Option<String>,
},
Done {
finish_reason: Option<String>,
usage: Option<Usage>,
},
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct StreamDelta {
#[serde(default)]
pub choices: Vec<StreamDeltaChoice>,
#[serde(default)]
pub usage: Option<StreamDeltaUsage>,
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct StreamDeltaChoice {
#[serde(default)]
pub delta: StreamDeltaContent,
#[serde(default)]
pub finish_reason: Option<String>,
}
#[derive(Debug, Clone, Default, Deserialize)]
pub(crate) struct StreamDeltaContent {
#[serde(default)]
pub content: Option<String>,
#[serde(default)]
pub tool_calls: Option<Vec<StreamDeltaToolCall>>,
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct StreamDeltaToolCall {
pub index: usize,
#[serde(default)]
pub id: Option<String>,
#[serde(default)]
pub function: Option<StreamDeltaFunction>,
}
#[derive(Debug, Clone, Default, Deserialize)]
pub(crate) struct StreamDeltaFunction {
#[serde(default)]
pub name: Option<String>,
#[serde(default)]
pub arguments: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct StreamDeltaUsage {
pub prompt_tokens: Option<i32>,
pub completion_tokens: Option<i32>,
pub total_tokens: Option<i32>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chat_message_new_helpers() {
let sys = ChatMessage::system("You are helpful.");
assert_eq!(sys.role, "system");
assert_eq!(sys.content.as_deref(), Some("You are helpful."));
assert!(sys.tool_call_id.is_none());
assert!(sys.tool_calls.is_none());
let user = ChatMessage::user("Hello");
assert_eq!(user.role, "user");
let asst = ChatMessage::assistant("Hi there");
assert_eq!(asst.role, "assistant");
}
#[test]
fn chat_message_serde_roundtrip() {
let msg = ChatMessage::user("Hello, world!");
let json = serde_json::to_string(&msg).unwrap();
let parsed: ChatMessage = serde_json::from_str(&json).unwrap();
assert_eq!(msg, parsed);
}
#[test]
fn chat_message_skips_none_fields() {
let msg = ChatMessage::user("Hi");
let json = serde_json::to_string(&msg).unwrap();
assert!(!json.contains("tool_call_id"));
assert!(!json.contains("tool_calls"));
}
#[test]
fn chat_message_with_tool_calls_roundtrip() {
let msg = ChatMessage {
role: "assistant".into(),
content: None,
tool_call_id: None,
tool_calls: Some(vec![ToolCall {
id: "call_abc123".into(),
call_type: "function".into(),
function: FunctionCall {
name: "get_weather".into(),
arguments: r#"{"city":"London"}"#.into(),
},
}]),
};
let json = serde_json::to_string(&msg).unwrap();
assert!(json.contains("tool_calls"));
assert!(json.contains("call_abc123"));
let parsed: ChatMessage = serde_json::from_str(&json).unwrap();
assert_eq!(msg, parsed);
}
#[test]
fn tool_call_type_field_renamed() {
let tc = ToolCall {
id: "tc1".into(),
call_type: "function".into(),
function: FunctionCall {
name: "search".into(),
arguments: "{}".into(),
},
};
let json = serde_json::to_string(&tc).unwrap();
assert!(json.contains(r#""type":"function""#));
assert!(!json.contains("call_type"));
}
#[test]
fn chat_request_serialization() {
let req = ChatRequest::new("gpt-4o", vec![ChatMessage::user("Hi")]);
let json = serde_json::to_string(&req).unwrap();
assert!(json.contains(r#""model":"gpt-4o""#));
assert!(json.contains(r#""role":"user""#));
assert!(!json.contains("tools"));
assert!(!json.contains("stream"));
assert!(!json.contains("max_tokens"));
assert!(!json.contains("temperature"));
}
#[test]
fn chat_request_with_all_fields() {
let req = ChatRequest {
model: "gpt-4o".into(),
messages: vec![ChatMessage::user("test")],
max_tokens: Some(100),
temperature: Some(0.7),
tools: vec![serde_json::json!({"type": "function", "function": {"name": "test"}})],
stream: Some(true),
};
let json = serde_json::to_string(&req).unwrap();
assert!(json.contains("max_tokens"));
assert!(json.contains("temperature"));
assert!(json.contains("tools"));
assert!(json.contains("stream"));
}
#[test]
fn chat_response_deserialization() {
let json = r#"{
"id": "chatcmpl-abc123",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello!"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15
},
"model": "gpt-4o-2024-05-13"
}"#;
let resp: ChatResponse = serde_json::from_str(json).unwrap();
assert_eq!(resp.id, "chatcmpl-abc123");
assert_eq!(resp.choices.len(), 1);
assert_eq!(resp.choices[0].message.content.as_deref(), Some("Hello!"));
assert_eq!(resp.choices[0].finish_reason.as_deref(), Some("stop"));
let usage = resp.usage.unwrap();
assert_eq!(usage.input_tokens, 10);
assert_eq!(usage.output_tokens, 5);
assert_eq!(usage.total_tokens, 15);
assert_eq!(resp.model, "gpt-4o-2024-05-13");
}
#[test]
fn chat_response_without_usage() {
let json = r#"{
"id": "resp-1",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "Ok"},
"finish_reason": null
}],
"usage": null,
"model": "test-model"
}"#;
let resp: ChatResponse = serde_json::from_str(json).unwrap();
assert!(resp.usage.is_none());
assert!(resp.choices[0].finish_reason.is_none());
}
#[test]
fn usage_serde_roundtrip() {
let usage = Usage {
input_tokens: 100,
output_tokens: 50,
total_tokens: 150,
};
let json = serde_json::to_string(&usage).unwrap();
let parsed: Usage = serde_json::from_str(&json).unwrap();
assert_eq!(usage, parsed);
}
#[test]
fn usage_deserializes_from_openai_format() {
let json = r#"{"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}"#;
let usage: Usage = serde_json::from_str(json).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.total_tokens, 150);
}
}