use crate::ai::provider::Decoder;
use crate::ai::sse::SseEvent;
use crate::ai::types::{AiError, Block, Event, Message, Request, Role, StopReason, Usage};
use serde_json::{Value, json};
fn messages(message: &Message) -> Vec<Value> {
let mut out = Vec::new();
let text: String = message
.blocks
.iter()
.filter_map(|block| if let Block::Text(text) = block { Some(text.as_str()) } else { None })
.collect::<Vec<_>>()
.join("\n");
let calls: Vec<Value> = message
.blocks
.iter()
.filter_map(|block| match block {
Block::ToolUse { id, name, input } => Some(json!({
"id": id,
"type": "function",
"function": { "name": name, "arguments": input.to_string() },
})),
_ => None,
})
.collect();
let images: Vec<Value> = message
.blocks
.iter()
.filter_map(|block| match block {
Block::Image { media_type, data_base64, .. } => Some(json!({
"type": "image_url",
"image_url": { "url": format!("data:{media_type};base64,{data_base64}") },
})),
_ => None,
})
.collect();
for block in &message.blocks {
if let Block::ToolResult { id, content, .. } = block {
out.push(json!({ "role": "tool", "tool_call_id": id, "content": content }));
}
}
if message.role == Role::Assistant {
let mut assistant =
json!({ "role": "assistant", "content": if text.is_empty() { Value::Null } else { json!(text) } });
if !calls.is_empty() {
assistant["tool_calls"] = json!(calls);
}
out.push(assistant);
} else if !images.is_empty() {
let mut content = images;
if !text.is_empty() {
content.push(json!({ "type": "text", "text": text }));
}
out.push(json!({ "role": "user", "content": content }));
} else if !text.is_empty() {
out.push(json!({ "role": "user", "content": text }));
}
out
}
pub fn body(request: &Request) -> Value {
let mut all: Vec<Value> =
request.system.iter().map(|system| json!({ "role": "system", "content": system })).collect();
all.extend(request.messages.iter().flat_map(messages));
let mut body = json!({
"model": request.model,
"max_tokens": request.max_tokens,
"stream": true,
"stream_options": { "include_usage": true },
"messages": all,
});
if !request.tools.is_empty() {
body["tools"] = request
.tools
.iter()
.map(|tool| {
json!({ "type": "function", "function": { "name": tool.name, "description": tool.description, "parameters": tool.schema } })
})
.collect();
}
body
}
fn stop_reason(reason: &str) -> StopReason {
match reason {
"stop" => StopReason::EndTurn,
"tool_calls" | "function_call" => StopReason::ToolUse,
"length" => StopReason::MaxTokens,
other => StopReason::Other(other.to_owned()),
}
}
#[derive(Default)]
pub struct OpenAiDecoder {
calls: Vec<(String, String, String)>,
tool_calls: bool,
usage: Usage,
stop: Option<StopReason>,
finished: bool,
complete: bool,
}
impl OpenAiDecoder {
fn flush_calls(&mut self, out: &mut Vec<Event>) -> Result<(), AiError> {
for (index, (id, name, arguments)) in std::mem::take(&mut self.calls).into_iter().enumerate() {
if name.is_empty() {
continue;
}
let id = if id.is_empty() { format!("call_{index}") } else { id };
self.tool_calls = true;
out.push(Event::ToolUse { id, name, input: crate::ai::wire::tool_input(&arguments) });
}
Ok(())
}
fn slot(&mut self, call: &Value) -> usize {
let index = match call["index"].as_u64() {
Some(index) => index as usize,
None if call["id"].is_string() || self.calls.is_empty() => self.calls.len(),
None => self.calls.len() - 1,
};
if self.calls.len() <= index {
self.calls.resize(index + 1, Default::default());
}
index
}
}
impl Decoder for OpenAiDecoder {
fn feed(&mut self, event: &SseEvent, out: &mut Vec<Event>) -> Result<(), AiError> {
if event.data == "[DONE]" {
self.complete = true;
self.flush_calls(out)?;
self.finish(out);
return Ok(());
}
let data: Value = serde_json::from_str(&event.data).map_err(|error| AiError::Decode(error.to_string()))?;
if let Some(message) = data["error"]["message"].as_str() {
return Err(AiError::Api(message.to_owned()));
}
if let Some(usage) = data.get("usage").filter(|usage| !usage.is_null()) {
self.usage = Usage {
input_tokens: usage["prompt_tokens"].as_u64().unwrap_or(0) as u32,
output_tokens: usage["completion_tokens"].as_u64().unwrap_or(0) as u32,
};
}
let choice = &data["choices"][0];
if let Some(text) = choice["delta"]["content"].as_str().filter(|text| !text.is_empty()) {
out.push(Event::Text(text.to_owned()));
}
for call in choice["delta"]["tool_calls"].as_array().into_iter().flatten() {
let index = self.slot(call);
let slot = &mut self.calls[index];
if let Some(id) = call["id"].as_str() {
slot.0 = id.to_owned();
}
if let Some(name) = call["function"]["name"].as_str() {
slot.1.push_str(name);
}
slot.2.push_str(call["function"]["arguments"].as_str().unwrap_or_default());
}
if let Some(reason) = choice["finish_reason"].as_str() {
self.complete = true;
self.flush_calls(out)?;
self.stop = Some(if self.tool_calls { StopReason::ToolUse } else { stop_reason(reason) });
}
Ok(())
}
fn finish(&mut self, out: &mut Vec<Event>) {
if std::mem::replace(&mut self.finished, true) {
return;
}
out.push(Event::Usage(self.usage));
out.push(Event::Done(self.stop.take().unwrap_or(StopReason::EndTurn)));
}
fn is_complete(&self) -> bool {
self.complete
}
}
#[cfg(test)]
mod tests;