362 lines
9.9 KiB
Rust
362 lines
9.9 KiB
Rust
use serde::{Deserialize, Deserializer, Serialize};
|
|
use serde_json::Value;
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ChatCompletionRequest {
|
|
pub model: String,
|
|
pub messages: Vec<ChatMessage>,
|
|
#[serde(default = "default_max_tokens", alias = "max_completion_tokens")]
|
|
pub max_tokens: u32,
|
|
#[serde(default = "default_temperature")]
|
|
pub temperature: f32,
|
|
#[serde(default = "default_top_p")]
|
|
pub top_p: f32,
|
|
#[serde(default)]
|
|
pub stream: bool,
|
|
#[serde(default, deserialize_with = "deserialize_stop_sequences")]
|
|
pub stop: Vec<String>,
|
|
#[serde(default = "default_n")]
|
|
pub n: u32,
|
|
#[serde(default)]
|
|
pub seed: Option<u32>,
|
|
#[serde(default)]
|
|
pub user: Option<String>,
|
|
#[serde(default)]
|
|
pub response_format: Option<ResponseFormat>,
|
|
#[serde(default)]
|
|
pub tools: Vec<ToolDefinition>,
|
|
#[serde(default)]
|
|
pub tool_choice: Option<ToolChoice>,
|
|
#[serde(default)]
|
|
pub parallel_tool_calls: Option<bool>,
|
|
#[serde(default)]
|
|
pub stream_options: Option<StreamOptions>,
|
|
}
|
|
|
|
fn default_max_tokens() -> u32 {
|
|
1024
|
|
}
|
|
|
|
fn default_temperature() -> f32 {
|
|
0.15
|
|
}
|
|
|
|
fn default_top_p() -> f32 {
|
|
0.95
|
|
}
|
|
|
|
fn default_n() -> u32 {
|
|
1
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ChatMessage {
|
|
pub role: Role,
|
|
#[serde(
|
|
default,
|
|
deserialize_with = "deserialize_message_content",
|
|
skip_serializing_if = "Option::is_none"
|
|
)]
|
|
pub content: Option<String>,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub name: Option<String>,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub tool_calls: Option<Vec<ToolCall>>,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub tool_call_id: Option<String>,
|
|
}
|
|
|
|
impl ChatMessage {
|
|
pub fn text(&self) -> &str {
|
|
self.content.as_deref().unwrap_or("")
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum Role {
|
|
System,
|
|
Developer,
|
|
User,
|
|
Assistant,
|
|
Tool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct ChatCompletionResponse {
|
|
pub id: String,
|
|
pub object: &'static str,
|
|
pub created: i64,
|
|
pub model: String,
|
|
pub choices: Vec<Choice>,
|
|
pub usage: Usage,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct Choice {
|
|
pub index: u32,
|
|
pub message: ChatMessage,
|
|
pub finish_reason: FinishReason,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum FinishReason {
|
|
Stop,
|
|
Length,
|
|
ToolCalls,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct ChatCompletionChunk {
|
|
pub id: String,
|
|
pub object: &'static str,
|
|
pub created: i64,
|
|
pub model: String,
|
|
pub choices: Vec<ChunkChoice>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub usage: Option<Usage>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct ChunkChoice {
|
|
pub index: u32,
|
|
pub delta: Delta,
|
|
pub finish_reason: Option<FinishReason>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct Delta {
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub role: Option<Role>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub content: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub tool_calls: Option<Vec<DeltaToolCall>>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ToolDefinition {
|
|
#[serde(rename = "type")]
|
|
pub tool_type: ToolType,
|
|
pub function: FunctionDefinition,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum ToolType {
|
|
Function,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct FunctionDefinition {
|
|
pub name: String,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub description: Option<String>,
|
|
#[serde(default)]
|
|
pub parameters: Value,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ToolCall {
|
|
pub id: String,
|
|
#[serde(rename = "type")]
|
|
pub tool_type: ToolType,
|
|
pub function: ToolFunctionCall,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ToolFunctionCall {
|
|
pub name: String,
|
|
pub arguments: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct DeltaToolCall {
|
|
pub index: u32,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub id: Option<String>,
|
|
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
|
pub tool_type: Option<ToolType>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub function: Option<DeltaToolFunctionCall>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct DeltaToolFunctionCall {
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub name: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub arguments: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
#[serde(untagged)]
|
|
pub enum ToolChoice {
|
|
Mode(String),
|
|
Specific(ToolChoiceObject),
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ToolChoiceObject {
|
|
#[serde(rename = "type")]
|
|
pub choice_type: ToolType,
|
|
pub function: ToolChoiceFunction,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ToolChoiceFunction {
|
|
pub name: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ResponseFormat {
|
|
#[serde(rename = "type")]
|
|
pub format_type: String,
|
|
#[serde(default)]
|
|
pub json_schema: Option<Value>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct StreamOptions {
|
|
#[serde(default)]
|
|
pub include_usage: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Default)]
|
|
pub struct Usage {
|
|
pub prompt_tokens: u32,
|
|
pub completion_tokens: u32,
|
|
pub total_tokens: u32,
|
|
pub tokens_per_second: f64,
|
|
}
|
|
|
|
fn deserialize_stop_sequences<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
|
|
where
|
|
D: Deserializer<'de>,
|
|
{
|
|
let value = Option::<Value>::deserialize(deserializer)?;
|
|
let Some(value) = value else {
|
|
return Ok(Vec::new());
|
|
};
|
|
|
|
match value {
|
|
Value::Null => Ok(Vec::new()),
|
|
Value::String(s) => Ok(vec![s]),
|
|
Value::Array(items) => items
|
|
.into_iter()
|
|
.map(|item| match item {
|
|
Value::String(s) => Ok(s),
|
|
_ => Err(serde::de::Error::custom("stop items must be strings")),
|
|
})
|
|
.collect(),
|
|
_ => Err(serde::de::Error::custom(
|
|
"stop must be a string, array of strings, or null",
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn deserialize_message_content<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
|
|
where
|
|
D: Deserializer<'de>,
|
|
{
|
|
let value = Option::<Value>::deserialize(deserializer)?;
|
|
let Some(value) = value else {
|
|
return Ok(None);
|
|
};
|
|
|
|
match value {
|
|
Value::Null => Ok(None),
|
|
Value::String(s) => Ok(Some(s)),
|
|
Value::Array(items) => {
|
|
let mut text_parts = Vec::new();
|
|
|
|
for item in items {
|
|
match item {
|
|
Value::String(s) => text_parts.push(s),
|
|
Value::Object(map) => {
|
|
let item_type = map.get("type").and_then(Value::as_str).unwrap_or("text");
|
|
if matches!(item_type, "text" | "input_text" | "output_text") {
|
|
if let Some(text) = map.get("text").and_then(Value::as_str) {
|
|
text_parts.push(text.to_owned());
|
|
} else if let Some(text) = map.get("content").and_then(Value::as_str) {
|
|
text_parts.push(text.to_owned());
|
|
}
|
|
}
|
|
}
|
|
_ => {
|
|
return Err(serde::de::Error::custom(
|
|
"message content array items must be strings or text parts",
|
|
));
|
|
}
|
|
}
|
|
}
|
|
|
|
if text_parts.is_empty() {
|
|
Ok(None)
|
|
} else {
|
|
Ok(Some(text_parts.join("\n")))
|
|
}
|
|
}
|
|
_ => Err(serde::de::Error::custom(
|
|
"message content must be a string, array, or null",
|
|
)),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{ChatCompletionRequest, Role, ToolChoice};
|
|
|
|
#[test]
|
|
fn deserializes_openai_style_message_parts_and_aliases() {
|
|
let request: ChatCompletionRequest = serde_json::from_str(
|
|
r#"{
|
|
"model": "alias-model",
|
|
"messages": [
|
|
{
|
|
"role": "developer",
|
|
"content": [
|
|
{ "type": "text", "text": "Write code" },
|
|
{ "type": "input_text", "text": "Use Rust" }
|
|
]
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "Build a parser"
|
|
}
|
|
],
|
|
"max_completion_tokens": 77,
|
|
"stop": "</tool_call>",
|
|
"tool_choice": "auto"
|
|
}"#,
|
|
)
|
|
.unwrap();
|
|
|
|
assert_eq!(request.max_tokens, 77);
|
|
assert_eq!(request.stop, vec!["</tool_call>"]);
|
|
assert!(matches!(request.messages[0].role, Role::Developer));
|
|
assert_eq!(request.messages[0].text(), "Write code\nUse Rust");
|
|
assert!(matches!(request.tool_choice, Some(ToolChoice::Mode(ref mode)) if mode == "auto"));
|
|
}
|
|
|
|
#[test]
|
|
fn deserializes_stop_arrays_and_null_content() {
|
|
let request: ChatCompletionRequest = serde_json::from_str(
|
|
r#"{
|
|
"model": "alias-model",
|
|
"messages": [
|
|
{ "role": "assistant", "content": null },
|
|
{ "role": "tool", "tool_call_id": "call_1", "content": "done" }
|
|
],
|
|
"stop": ["END", "STOP"]
|
|
}"#,
|
|
)
|
|
.unwrap();
|
|
|
|
assert_eq!(request.stop, vec!["END", "STOP"]);
|
|
assert_eq!(request.messages[0].content, None);
|
|
assert_eq!(request.messages[1].tool_call_id.as_deref(), Some("call_1"));
|
|
}
|
|
}
|