use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LlmRole {
System,
User,
Assistant,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmMessage {
pub role: LlmRole,
pub content: String,
}
impl LlmMessage {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: LlmRole::System,
content: content.into(),
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: LlmRole::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: LlmRole::Assistant,
content: content.into(),
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct ChatCompletionRequest {
pub model: String,
pub messages: Vec<LlmMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(flatten)]
pub extra: serde_json::Value,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ChatCompletionResponse {
pub choices: Vec<Choice>,
#[serde(default)]
pub usage: Option<Usage>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct Choice {
pub message: ResponseMessage,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ResponseMessage {
pub role: String,
#[serde(default)]
pub content: Option<String>,
#[serde(default)]
pub reasoning_content: Option<String>,
}
impl ResponseMessage {
pub fn text(&self) -> String {
self.content
.as_deref()
.filter(|s| !s.trim().is_empty())
.or(self.reasoning_content.as_deref())
.unwrap_or_default()
.trim()
.to_string()
}
}
#[derive(Debug, Clone, Deserialize)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_response_message_text_prefers_content() {
let msg = ResponseMessage {
role: "assistant".to_string(),
content: Some("final answer".to_string()),
reasoning_content: Some("thinking...".to_string()),
};
assert_eq!(msg.text(), "final answer");
}
#[test]
fn test_response_message_text_falls_back_to_reasoning() {
let msg = ResponseMessage {
role: "assistant".to_string(),
content: Some(" ".to_string()),
reasoning_content: Some("only reasoning".to_string()),
};
assert_eq!(msg.text(), "only reasoning");
}
#[test]
fn test_response_message_text_empty() {
let msg = ResponseMessage {
role: "assistant".to_string(),
content: None,
reasoning_content: None,
};
assert_eq!(msg.text(), "");
}
#[test]
fn test_chat_completion_response_parse_deepseek() {
let json = r#"{
"choices": [{
"message": {
"role": "assistant",
"content": "",
"reasoning_content": "deep thinking"
}
}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
}"#;
let resp: ChatCompletionResponse = serde_json::from_str(json).unwrap();
assert_eq!(resp.choices[0].message.text(), "deep thinking");
assert_eq!(resp.usage.as_ref().unwrap().total_tokens, 15);
}
#[test]
fn test_chat_completion_request_serialize() {
let req = ChatCompletionRequest {
model: "deepseek-chat".to_string(),
messages: vec![LlmMessage::user("hi")],
max_tokens: Some(100),
temperature: Some(0.3),
extra: serde_json::json!({"thinking": {"type": "enabled", "effort": "high"}}),
};
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["model"], "deepseek-chat");
assert_eq!(json["thinking"]["effort"], "high");
assert_eq!(json["messages"][0]["role"], "user");
}
}