use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Deserialize, Serialize)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum ChatMessage {
System {
content: String,
},
User {
content: String,
},
Assistant {
#[serde(default, skip_serializing_if = "Option::is_none")]
content: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
tool_calls: Vec<ToolCall>,
#[serde(default, skip_serializing, alias = "reasoning")]
reasoning_content: Option<String>,
},
Tool {
tool_call_id: String,
content: String,
},
}
#[derive(Clone, Deserialize, Serialize)]
pub struct ToolCall {
#[serde(default)]
pub id: String,
#[serde(default = "function_type", rename = "type")]
pub kind: String,
pub function: FunctionCall,
}
#[derive(Clone, Deserialize, Serialize)]
pub struct FunctionCall {
pub name: String,
#[serde(default)]
pub arguments: String,
}
fn function_type() -> String {
"function".to_owned()
}
#[derive(Serialize)]
pub struct ToolDef {
#[serde(rename = "type")]
pub kind: &'static str,
pub function: FunctionDef,
}
#[derive(Serialize)]
pub struct FunctionDef {
pub name: &'static str,
pub description: &'static str,
pub parameters: Value,
}
#[derive(Serialize)]
pub struct ChatRequest<'a> {
pub model: &'a str,
pub messages: &'a [ChatMessage],
pub tools: &'a [ToolDef],
pub stream: bool,
}
#[derive(Deserialize)]
pub struct ChatResponse {
pub choices: Vec<ChatChoice>,
}
#[derive(Deserialize)]
pub struct ChatChoice {
pub message: ChatMessage,
}
#[derive(Deserialize)]
pub struct ChatChunk {
#[serde(default)]
pub choices: Vec<ChunkChoice>,
}
#[derive(Deserialize)]
pub struct ChunkChoice {
#[serde(default)]
pub delta: ChunkDelta,
}
#[derive(Default, Deserialize)]
pub struct ChunkDelta {
#[serde(default)]
pub content: Option<String>,
#[serde(default)]
pub reasoning_content: Option<String>,
#[serde(default)]
pub reasoning: Option<String>,
#[serde(default)]
pub tool_calls: Vec<ToolCallDelta>,
}
#[derive(Deserialize)]
pub struct ToolCallDelta {
#[serde(default)]
pub index: usize,
#[serde(default)]
pub id: Option<String>,
#[serde(default)]
pub function: Option<FunctionCallDelta>,
}
#[derive(Deserialize)]
pub struct FunctionCallDelta {
#[serde(default)]
pub name: Option<String>,
#[serde(default)]
pub arguments: Option<String>,
}
#[derive(Default)]
pub struct ChunkAccumulator {
content: String,
reasoning: String,
tool_calls: Vec<ToolCall>,
}
impl ChunkAccumulator {
pub fn push(&mut self, chunk: ChatChunk) {
for choice in chunk.choices {
let delta = choice.delta;
if let Some(content) = delta.content {
self.content.push_str(&content);
}
if let Some(reasoning) = delta
.reasoning_content
.into_iter()
.chain(delta.reasoning)
.find(|x| !x.is_empty())
{
self.reasoning.push_str(&reasoning);
}
for fragment in delta.tool_calls {
if self.tool_calls.len() <= fragment.index {
self.tool_calls
.resize_with(fragment.index + 1, || ToolCall {
id: String::new(),
kind: function_type(),
function: FunctionCall {
name: String::new(),
arguments: String::new(),
},
});
}
let call = &mut self.tool_calls[fragment.index];
if let Some(id) = fragment.id {
call.id = id;
}
if let Some(function) = fragment.function {
if let Some(name) = function.name {
call.function.name = name;
}
if let Some(arguments) = function.arguments {
call.function.arguments.push_str(&arguments);
}
}
}
}
}
pub fn progress(&self) -> (&str, &str) {
(&self.content, &self.reasoning)
}
pub fn finish(self) -> ChatMessage {
ChatMessage::Assistant {
content: (!self.content.is_empty()).then_some(self.content),
tool_calls: self.tool_calls,
reasoning_content: (!self.reasoning.is_empty()).then_some(self.reasoning),
}
}
}