use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use futures_core::Stream;
use futures_util::StreamExt;
use tokio::sync::RwLock;
use crate::beta::openai::{
bind_chat_inputs, decode_base64_image, ChatCompletionChunk,
ChatCompletionReasoningEffort, ChatInputs,
};
use crate::client::Result;
use crate::MunaError;
use crate::services::{PredictionService, PredictorService};
use crate::types::{self, Acceleration, Device, Dtype, Prediction, Signature, Value};
use super::schema::{
ContentBlock, ContentBlockDelta, ContentBlockParam, ImageSource, Message,
MessageContent, MessageCreateParams, MessageDelta, MessageParam,
OutputEffort, RawMessageStreamEvent, StopReason, ThinkingConfig,
Tool, Usage,
};
pub type MessageEventStream = Pin<Box<dyn Stream<Item = Result<RawMessageStreamEvent>> + Send>>;
type OutputStream = Pin<Box<dyn Stream<Item = Result<serde_json::Map<String, serde_json::Value>>> + Send>>;
type ChunkStream = Pin<Box<dyn Stream<Item = Result<ChatCompletionChunk>> + Send>>;
#[derive(Clone)]
struct DelegateInfo {
signature: Signature,
output_param_idx: usize,
}
#[derive(Clone)]
pub struct MessageService {
predictors: PredictorService,
predictions: PredictionService,
cache: Arc<RwLock<HashMap<String, DelegateInfo>>>,
}
impl MessageService {
pub fn new(
predictors: PredictorService,
predictions: PredictionService
) -> Self {
Self {
predictors,
predictions,
cache: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn create(&self, params: MessageCreateParams) -> Result<Message> {
let mut events = self.stream(params).await?;
let mut message: Option<Message> = None;
while let Some(event) = events.next().await {
apply_event(&mut message, &event?);
}
message.ok_or_else(|| {
MunaError::Prediction(
"Failed to create message because the model did not return any outputs".into(),
)
})
}
pub async fn stream(&self, params: MessageCreateParams) -> Result<MessageEventStream> {
let model = params.model.clone();
let (
input_map,
output_param_idx,
acceleration,
devices
) = self.prepare_prediction(params).await?;
let prediction_stream = self
.predictions
.stream(&model, input_map, Some(acceleration), devices.as_deref())
.await?;
let outputs = gather_prediction_outputs(
prediction_stream,
output_param_idx,
model
);
Ok(events_from_chunks(parse_completion_chunks(outputs)))
}
async fn prepare_prediction(
&self,
mut params: MessageCreateParams,
) -> Result<(HashMap<String, Value>, usize, Acceleration, Option<Vec<Device>>)> {
self.ensure_delegate_info(¶ms.model).await?;
let info = {
let cache = self.cache.read().await;
cache.get(¶ms.model).cloned().ok_or_else(|| {
MunaError::Prediction(format!(
"{} cannot be used with Anthropic messages API because \
the predictor metadata could not be cached.",
params.model
))
})?
};
let acceleration = params.acceleration.take().unwrap_or(Acceleration::LocalAuto);
let devices = params.devices.take();
let inputs = tokio::task::spawn_blocking(move || params.chat_inputs())
.await
.map_err(|e| MunaError::Prediction(format!("Failed to prepare chat inputs: {e}")))??;
let input_map = bind_chat_inputs(inputs, &info.signature)?;
Ok((input_map, info.output_param_idx, acceleration, devices))
}
async fn ensure_delegate_info(&self, tag: &str) -> Result<()> {
{
let cache = self.cache.read().await;
if cache.contains_key(tag) {
return Ok(());
}
}
let info = self.create_delegate_info(tag).await?;
self.cache
.write()
.await
.entry(tag.to_string())
.or_insert(info);
Ok(())
}
async fn create_delegate_info(&self, tag: &str) -> Result<DelegateInfo> {
let predictor = self.predictors.retrieve(tag).await?.ok_or_else(|| {
MunaError::Prediction(format!(
"{tag} cannot be used with Anthropic messages API because \
the predictor could not be found. Check that your access key \
is valid and that you have access to the predictor."
))
})?;
let signature = &predictor.signature;
let output_param_idx = signature
.outputs
.iter()
.position(|param| {
param.dtype == Some(Dtype::Dict)
&& param
.schema
.as_ref()
.and_then(|s| s.get("title"))
.and_then(|v| v.as_str())
.is_some_and(|title| title == "ChatCompletionChunk")
})
.ok_or_else(|| {
MunaError::Prediction(format!(
"{tag} cannot be used with Anthropic messages API because \
it does not have a valid chat completion chunk output parameter. \
Chat predictors must yield `ChatCompletionChunk` outputs."
))
})?;
Ok(DelegateInfo {
signature: signature.clone(),
output_param_idx,
})
}
}
impl MessageCreateParams {
pub fn reasoning_effort(&self) -> Option<ChatCompletionReasoningEffort> {
use ChatCompletionReasoningEffort as E;
if matches!(self.thinking, Some(ThinkingConfig::Disabled)) {
return Some(E::None);
}
let effort = self
.output_config
.as_ref()
.and_then(|config| config.effort)
.map(|effort| match effort {
OutputEffort::Low => E::Low,
OutputEffort::Medium => E::Medium,
OutputEffort::High => E::High,
OutputEffort::XHigh => E::XHigh,
OutputEffort::Max => E::XHigh,
});
match (effort, &self.thinking) {
(Some(effort), _) => Some(effort),
(None, Some(_)) => Some(E::Medium),
(None, None) => None,
}
}
pub fn chat_inputs(&self) -> Result<ChatInputs> {
if self.messages.is_empty() {
return Err(MunaError::InvalidInput(
"`messages` must contain at least one message.".into(),
));
}
if let Some(MessageContent::Blocks(blocks)) = &self.system {
if blocks.iter().any(|block| matches!(block, ContentBlockParam::Image { .. })) {
return Err(MunaError::InvalidInput(
"`system` content cannot contain `image` blocks.".into(),
));
}
}
for (index, message) in self.messages.iter().enumerate() {
validate_blocks(index, message)?;
validate_message(index, message)?;
}
let images = decode_images(&self.messages)?;
Ok(ChatInputs {
messages: to_openai_messages(self.system.as_ref(), &self.messages),
images: (!images.is_empty()).then_some(images),
tools: self
.tools
.as_deref()
.filter(|tools| !tools.is_empty())
.map(to_openai_tools),
response_format: None,
reasoning_effort: self.reasoning_effort().map(|e| e.as_str().to_string()),
max_output_tokens: Some(self.max_tokens),
temperature: self.temperature,
top_p: self.top_p,
seed: None,
frequency_penalty: None,
presence_penalty: None,
stop_sequences: self.stop_sequences.clone(),
top_k: self.top_k,
})
}
}
fn validate_blocks(index: usize, message: &MessageParam) -> Result<()> {
let MessageContent::Blocks(blocks) = &message.content else {
return Ok(());
};
for (block_index, block) in blocks.iter().enumerate() {
let location = format!("`messages[{index}].content[{block_index}]`");
let kind = match block {
ContentBlockParam::Image { .. } if message.role == "system" => {
return Err(MunaError::InvalidInput(format!(
"{location} is an `image` block in a `system` message, which \
is not supported."
)));
}
ContentBlockParam::Image { source } => {
validate_image_source(&location, source)?;
continue;
}
ContentBlockParam::ToolResult { content: Some(content), .. } => {
if let MessageContent::Blocks(inner) = content.as_ref() {
if inner.iter().any(|block| matches!(block, ContentBlockParam::Image { .. })) {
return Err(MunaError::InvalidInput(format!(
"{location} is a `tool_result` block with `image` \
content, which is not supported."
)));
}
}
continue;
}
ContentBlockParam::Document => "a `document` block",
ContentBlockParam::Unsupported => "a block of an unrecognized type",
_ => continue,
};
return Err(MunaError::InvalidInput(format!(
"{location} is {kind}, which is not supported. Supported block types \
are `text`, `image`, `tool_use`, `tool_result`, `thinking`, and \
`redacted_thinking`."
)));
}
Ok(())
}
fn validate_image_source(
location: &str,
source: &ImageSource
) -> Result<()> {
let kind = match source {
ImageSource::Base64 { .. } => return Ok(()),
ImageSource::Url { .. } => "a `url` source",
ImageSource::Unsupported => "an unsupported source",
};
Err(MunaError::InvalidInput(format!(
"{location} is an `image` block with {kind}. Only `base64` image sources \
are supported; fetch remote images and pass them inline as `base64`."
)))
}
fn decode_images(messages: &[MessageParam]) -> Result<Vec<types::Image>> {
messages
.iter()
.filter_map(|message| match &message.content {
MessageContent::Blocks(blocks) => Some(blocks.iter()),
MessageContent::Text(_) => None,
})
.flatten()
.filter_map(|block| match block {
ContentBlockParam::Image {
source: ImageSource::Base64 { media_type, data },
} => Some(decode_base64_image(data, media_type)),
_ => None,
})
.collect()
}
fn validate_message(index: usize, message: &MessageParam) -> Result<()> {
if message.role == "assistant" {
return Ok(());
}
let stripped;
let content = if message.role == "system" {
match strip_billing_header(&message.content) {
Some(content) => {
stripped = content;
&stripped
}
None => return Ok(()),
}
} else {
&message.content
};
let has_payload = match content {
MessageContent::Text(text) => !text.trim().is_empty(),
MessageContent::Blocks(blocks) => blocks.iter().any(|block| match block {
ContentBlockParam::Text { text } => !text.trim().is_empty(),
ContentBlockParam::ToolUse { .. } | ContentBlockParam::ToolResult { .. } => true,
ContentBlockParam::Thinking { .. } | ContentBlockParam::RedactedThinking { .. } => false,
ContentBlockParam::Image { .. }
| ContentBlockParam::Document
| ContentBlockParam::Unsupported => true,
}),
};
if has_payload {
return Ok(());
}
Err(MunaError::InvalidInput(format!(
"`messages[{index}]` ({} message) must have non-empty `content`. Text \
content must contain non-whitespace characters.",
message.role
)))
}
const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header";
fn strip_billing_header(content: &MessageContent) -> Option<MessageContent> {
match content {
MessageContent::Text(text) => {
(!text.starts_with(BILLING_HEADER_PREFIX)).then(|| content.clone())
}
MessageContent::Blocks(blocks) => {
let kept: Vec<ContentBlockParam> = blocks
.iter()
.filter(|block| {
!matches!(block, ContentBlockParam::Text { text } if text.starts_with(BILLING_HEADER_PREFIX))
})
.cloned()
.collect();
(!kept.is_empty()).then_some(MessageContent::Blocks(kept))
}
}
}
fn to_openai_messages(
system: Option<&MessageContent>,
messages: &[MessageParam]
) -> Vec<serde_json::Value> {
let mut result = Vec::with_capacity(messages.len() + usize::from(system.is_some()));
if let Some(system) = system.and_then(strip_billing_header) {
result.push(serde_json::json!({
"role": "system",
"content": system.flatten(),
}));
}
for message in messages {
if message.role == "system" {
if let Some(content) = strip_billing_header(&message.content) {
result.push(serde_json::json!({
"role": "system",
"content": content.flatten(),
}));
}
continue;
}
translate_message_for_openai(message, &mut result);
}
result
}
fn to_openai_tools(tools: &[Tool]) -> Vec<serde_json::Value> {
tools.iter().map(openai_tool).collect()
}
fn openai_tool(tool: &Tool) -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.input_schema,
}
})
}
enum RunPart<'a> {
Text(&'a str),
Image,
}
fn translate_message_for_openai(
message: &MessageParam,
messages: &mut Vec<serde_json::Value>
) {
let MessageContent::Blocks(blocks) = &message.content else {
messages.push(serde_json::json!({
"role": message.role,
"content": message.content.flatten(),
}));
return;
};
let mut run: Vec<RunPart> = Vec::new();
let flush = |run: &mut Vec<RunPart>, messages: &mut Vec<serde_json::Value>| {
if run.is_empty() {
return;
}
let content = if run.iter().any(|part| matches!(part, RunPart::Image)) {
serde_json::Value::Array(
run.iter()
.map(|part| match part {
RunPart::Text(text) => serde_json::json!({ "type": "text", "text": text }),
RunPart::Image => serde_json::json!({ "type": "image" }),
})
.collect()
)
} else {
let texts: Vec<&str> = run
.iter()
.filter_map(|part| match part {
RunPart::Text(text) => Some(*text),
RunPart::Image => None,
})
.collect();
serde_json::Value::String(texts.join("\n"))
};
messages.push(serde_json::json!({
"role": message.role,
"content": content,
}));
run.clear();
};
for block in blocks {
match block {
ContentBlockParam::Text { text } => {
run.push(RunPart::Text(text.as_str()));
}
ContentBlockParam::Image { .. } => {
run.push(RunPart::Image);
}
ContentBlockParam::ToolUse { id, name, input } => {
flush(&mut run, messages);
messages.push(serde_json::json!({
"role": "assistant",
"content": null,
"tool_calls": [{
"id": id,
"type": "function",
"function": { "name": name, "arguments": input.to_string() },
}],
}));
}
ContentBlockParam::ToolResult { tool_use_id, content, .. } => {
flush(&mut run, messages);
messages.push(serde_json::json!({
"role": "tool",
"tool_call_id": tool_use_id,
"content": content.as_ref().map(|c| c.flatten()),
}));
}
ContentBlockParam::Thinking { .. }
| ContentBlockParam::RedactedThinking { .. }
| ContentBlockParam::Document
| ContentBlockParam::Unsupported => {}
}
}
flush(&mut run, messages);
}
#[derive(Clone, PartialEq, Eq)]
enum BlockKind {
Thinking,
Text,
ToolUse {
id: String,
name: String,
},
}
impl BlockKind {
fn empty_block(&self) -> ContentBlock {
match self {
Self::Thinking => ContentBlock::Thinking {
thinking: String::new(),
signature: String::new(),
},
Self::Text => ContentBlock::Text {
text: String::new(),
},
Self::ToolUse { id, name } => ContentBlock::ToolUse {
id: id.clone(),
name: name.clone(),
input: serde_json::json!({}),
},
}
}
fn delta(&self, fragment: &str) -> ContentBlockDelta {
match self {
Self::Thinking => ContentBlockDelta::ThinkingDelta {
thinking: fragment.to_string(),
},
Self::Text => ContentBlockDelta::TextDelta {
text: fragment.to_string(),
},
Self::ToolUse { .. } => ContentBlockDelta::InputJsonDelta {
partial_json: fragment.to_string(),
},
}
}
}
fn gather_prediction_outputs(
mut predictions: Pin<Box<dyn Stream<Item = Result<Prediction>> + Send>>,
output_param_idx: usize,
model: String,
) -> OutputStream {
Box::pin(async_stream::try_stream! {
while let Some(prediction) = predictions.next().await {
let prediction = prediction?;
if let Some(error) = prediction.error {
Err(MunaError::from_prediction_error(error))?;
}
let results = prediction
.results
.ok_or_else(|| MunaError::Prediction(format!("{model} returned no results")))?;
let output = results.get(output_param_idx).ok_or_else(|| {
MunaError::Prediction(format!("{model} returned fewer results than expected"))
})?;
match output {
Value::Dict(map) => yield map.clone(),
_ => Err(MunaError::Prediction(format!(
"{model} returned non-dict message output"
)))?,
}
}
})
}
fn parse_completion_chunks(mut outputs: OutputStream) -> ChunkStream {
Box::pin(async_stream::try_stream! {
while let Some(output) = outputs.next().await {
yield parse_completion_chunk(output?)?;
}
})
}
fn parse_completion_chunk(
output: serde_json::Map<String, serde_json::Value>,
) -> Result<ChatCompletionChunk> {
match output.get("object").and_then(|v| v.as_str()) {
Some("chat.completion.chunk") => from_object(output),
_ => Err(MunaError::Prediction(
"Failed to parse chat completion chunk from model output. \
Chat predictors must yield `ChatCompletionChunk` outputs."
.into(),
)),
}
}
fn events_from_chunks(mut chunks: ChunkStream) -> MessageEventStream {
Box::pin(async_stream::try_stream! {
let mut started = false;
let mut block_idx = 0usize;
let mut block_kind: Option<BlockKind> = None;
let mut stop_reason = StopReason::EndTurn;
let mut usage = Usage::default();
while let Some(chunk) = chunks.next().await {
let chunk = chunk?;
if !started {
yield RawMessageStreamEvent::MessageStart {
message: Message {
id: chunk.id.clone(),
r#type: "message".to_string(),
role: "assistant".to_string(),
content: Vec::new(),
model: chunk.model.clone(),
stop_reason: None,
stop_sequence: None,
usage: Usage {
input_tokens: Some(0),
..Default::default()
},
},
};
started = true;
}
if let Some(chunk_usage) = &chunk.usage {
let details = chunk_usage.prompt_tokens_details.as_ref();
let cache_read = details.and_then(|d| d.cached_tokens);
let cache_write = details.and_then(|d| d.cache_write_tokens);
usage.input_tokens = Some(
chunk_usage
.prompt_tokens
.saturating_sub(cache_read.unwrap_or(0) + cache_write.unwrap_or(0)),
);
usage.output_tokens += chunk_usage.completion_tokens;
usage.cache_read_input_tokens = cache_read;
usage.cache_creation_input_tokens = cache_write;
}
let Some(choice) = chunk.choices.first() else {
continue;
};
if let Some(reason) = &choice.finish_reason {
stop_reason = map_finish_reason(reason);
}
let Some(delta) = &choice.delta else {
continue;
};
let fragments = [
(BlockKind::Thinking, delta.reasoning_content.clone()),
(BlockKind::Text, delta.content.clone()),
];
for (kind, fragment) in fragments {
let Some(fragment) = fragment else {
continue;
};
if fragment.is_empty() {
continue;
}
if block_kind.as_ref() != Some(&kind) {
if block_kind.is_some() {
yield RawMessageStreamEvent::ContentBlockStop { index: block_idx };
block_idx += 1;
}
yield RawMessageStreamEvent::ContentBlockStart {
index: block_idx,
content_block: kind.empty_block(),
};
block_kind = Some(kind.clone());
}
yield RawMessageStreamEvent::ContentBlockDelta {
index: block_idx,
delta: kind.delta(&fragment),
};
}
let Some(tool_fragments) = &delta.tool_calls else {
continue;
};
for fragment in tool_fragments {
if let Some(id) = &fragment.id {
if block_kind.is_some() {
yield RawMessageStreamEvent::ContentBlockStop { index: block_idx };
block_idx += 1;
}
let name = fragment
.function
.as_ref()
.and_then(|f| f.name.clone())
.unwrap_or_default();
let kind = BlockKind::ToolUse { id: id.clone(), name };
yield RawMessageStreamEvent::ContentBlockStart {
index: block_idx,
content_block: kind.empty_block(),
};
block_kind = Some(kind);
}
let arguments = fragment
.function
.as_ref()
.and_then(|f| f.arguments.as_ref());
if let Some(arguments) = arguments {
if !arguments.is_empty() &&
matches!(block_kind, Some(BlockKind::ToolUse { .. }))
{
yield RawMessageStreamEvent::ContentBlockDelta {
index: block_idx,
delta: ContentBlockDelta::InputJsonDelta {
partial_json: arguments.clone(),
},
};
}
}
}
}
if block_kind.is_some() {
yield RawMessageStreamEvent::ContentBlockStop { index: block_idx };
}
yield RawMessageStreamEvent::MessageDelta {
delta: MessageDelta {
stop_reason: Some(stop_reason),
stop_sequence: None,
},
usage,
};
yield RawMessageStreamEvent::MessageStop;
})
}
fn apply_event(message: &mut Option<Message>, event: &RawMessageStreamEvent) {
match event {
RawMessageStreamEvent::MessageStart { message: started } => {
*message = Some(started.clone());
}
RawMessageStreamEvent::ContentBlockStart { content_block, .. } => {
if let Some(message) = message {
message.content.push(content_block.clone());
}
}
RawMessageStreamEvent::ContentBlockDelta { index, delta } => {
if let Some(message) = message {
if let Some(block) = message.content.get_mut(*index) {
match (block, delta) {
(
ContentBlock::Text { text },
ContentBlockDelta::TextDelta { text: fragment },
) => text.push_str(fragment),
(
ContentBlock::Thinking { thinking, .. },
ContentBlockDelta::ThinkingDelta { thinking: fragment },
) => thinking.push_str(fragment),
(
ContentBlock::ToolUse { input, .. },
ContentBlockDelta::InputJsonDelta { partial_json },
) => match input {
serde_json::Value::String(buffer) => buffer.push_str(partial_json),
_ => *input = serde_json::Value::String(partial_json.clone()),
},
_ => {}
}
}
}
}
RawMessageStreamEvent::ContentBlockStop { index } => {
if let Some(message) = message {
if let Some(ContentBlock::ToolUse { input, .. }) = message.content.get_mut(*index) {
if let serde_json::Value::String(buffer) = input {
*input = serde_json::from_str(buffer)
.unwrap_or_else(|_| serde_json::json!({}));
}
}
}
}
RawMessageStreamEvent::MessageDelta { delta, usage } => {
if let Some(message) = message {
message.stop_reason = delta.stop_reason;
message.stop_sequence = delta.stop_sequence.clone();
if usage.input_tokens.is_some() {
message.usage.input_tokens = usage.input_tokens;
}
message.usage.output_tokens = usage.output_tokens;
message.usage.cache_creation_input_tokens = usage.cache_creation_input_tokens;
message.usage.cache_read_input_tokens = usage.cache_read_input_tokens;
}
}
_ => {}
}
}
fn map_finish_reason(finish_reason: &str) -> StopReason {
match finish_reason {
"length" => StopReason::MaxTokens,
"content_filter" => StopReason::Refusal,
"tool_calls" => StopReason::ToolUse,
_ => StopReason::EndTurn,
}
}
fn from_object<T>(output: serde_json::Map<String, serde_json::Value>) -> Result<T>
where T: serde::de::DeserializeOwned,
{
serde_json::from_value(serde_json::Value::Object(output))
.map_err(|e| MunaError::Prediction(e.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::beta::openai::{
ChatCompletionChunkChoice, ChatCompletionDelta, ChatCompletionUsage,
PromptTokensDetails,
};
use futures_util::stream;
fn chunk(
delta: ChatCompletionDelta,
finish_reason: Option<&str>,
usage: Option<ChatCompletionUsage>,
) -> ChatCompletionChunk {
ChatCompletionChunk {
object: "chat.completion.chunk".to_string(),
id: "msg-test".to_string(),
model: "test-model".to_string(),
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: Some(delta),
finish_reason: finish_reason.map(str::to_string),
logprobs: None,
}],
created: 0,
usage,
}
}
async fn collect_events(chunks: Vec<ChatCompletionChunk>) -> Vec<RawMessageStreamEvent> {
let chunk_stream = Box::pin(stream::iter(chunks.into_iter().map(Ok)));
let mut events = events_from_chunks(chunk_stream);
let mut collected = Vec::new();
while let Some(event) = events.next().await {
collected.push(event.unwrap());
}
collected
}
fn accumulate(events: &[RawMessageStreamEvent]) -> Message {
let mut message = None;
for event in events {
apply_event(&mut message, event);
}
message.unwrap()
}
#[tokio::test]
async fn chunks_become_ordered_events() {
let events = collect_events(vec![
chunk(
ChatCompletionDelta {
role: Some("assistant".to_string()),
content: None,
reasoning_content: Some("Let me think.".to_string()),
tool_calls: None,
},
None,
None,
),
chunk(
ChatCompletionDelta {
role: None,
content: Some("Paris.".to_string()),
reasoning_content: None,
tool_calls: None,
},
Some("stop"),
Some(ChatCompletionUsage {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
prompt_tokens_details: Some(PromptTokensDetails {
audio_tokens: None,
cache_write_tokens: Some(2),
cached_tokens: Some(3),
}),
completion_tokens_details: None,
}),
),
])
.await;
let types: Vec<&str> = events.iter().map(|e| e.event_type()).collect();
assert_eq!(
types,
vec![
"message_start",
"content_block_start",
"content_block_delta",
"content_block_stop",
"content_block_start",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
]
);
let message = accumulate(&events);
assert!(matches!(
&message.content[0],
ContentBlock::Thinking { thinking, .. } if thinking == "Let me think."
));
assert!(matches!(
&message.content[1],
ContentBlock::Text { text } if text == "Paris."
));
assert_eq!(message.stop_reason, Some(StopReason::EndTurn));
assert_eq!(message.usage.input_tokens, Some(5));
assert_eq!(message.usage.output_tokens, 5);
assert_eq!(message.usage.cache_read_input_tokens, Some(3));
assert_eq!(message.usage.cache_creation_input_tokens, Some(2));
}
#[tokio::test]
async fn length_finish_reason_maps_to_max_tokens() {
let events = collect_events(vec![chunk(
ChatCompletionDelta {
role: Some("assistant".to_string()),
content: Some("Hello".to_string()),
reasoning_content: None,
tool_calls: None,
},
Some("length"),
None,
)])
.await;
let message = accumulate(&events);
assert_eq!(message.stop_reason, Some(StopReason::MaxTokens));
}
#[tokio::test]
async fn tool_call_deltas_become_tool_use_blocks() {
let deltas: Vec<ChatCompletionDelta> = vec![
serde_json::from_value(serde_json::json!({
"role": "assistant",
"tool_calls": [{
"index": 0,
"id": "call_test_0",
"type": "function",
"function": { "name": "get_weather", "arguments": "" },
}],
})).unwrap(),
serde_json::from_value(serde_json::json!({
"tool_calls": [{
"index": 0,
"function": { "arguments": "{\"location\": " },
}],
})).unwrap(),
serde_json::from_value(serde_json::json!({
"tool_calls": [{
"index": 0,
"function": { "arguments": "\"Paris\"}" },
}],
})).unwrap(),
];
let mut chunks: Vec<ChatCompletionChunk> = deltas
.into_iter()
.map(|delta| chunk(delta, None, None))
.collect();
chunks.push(chunk(ChatCompletionDelta::default(), Some("tool_calls"), None));
let events = collect_events(chunks).await;
let message = accumulate(&events);
assert!(matches!(
&message.content[0],
ContentBlock::ToolUse { id, name, input }
if id == "call_test_0" &&
name == "get_weather" &&
*input == serde_json::json!({ "location": "Paris" })
));
assert_eq!(message.stop_reason, Some(StopReason::ToolUse));
let start = events.iter().find(|e| e.event_type() == "content_block_start").unwrap();
let json = serde_json::to_value(start).unwrap();
assert_eq!(json["content_block"]["type"], "tool_use");
let delta = events.iter().find(|e| e.event_type() == "content_block_delta").unwrap();
let json = serde_json::to_value(delta).unwrap();
assert_eq!(json["delta"]["type"], "input_json_delta");
}
#[test]
fn tool_blocks_translate_to_openai_messages() {
let assistant: MessageParam = serde_json::from_value(serde_json::json!({
"role": "assistant",
"content": [
{ "type": "text", "text": "Checking the weather." },
{ "type": "tool_use", "id": "call_123", "name": "get_weather",
"input": { "city": "Paris" } },
],
})).unwrap();
let user: MessageParam = serde_json::from_value(serde_json::json!({
"role": "user",
"content": [
{ "type": "tool_result", "tool_use_id": "call_123",
"content": "18C and sunny" },
],
})).unwrap();
let system = MessageContent::Text("You are a weather bot.".to_string());
let messages = to_openai_messages(Some(&system), &[assistant, user]);
assert_eq!(messages.len(), 4);
assert_eq!(messages[0], serde_json::json!({ "role": "system", "content": "You are a weather bot." }));
let messages = &messages[1..];
assert_eq!(messages[0]["role"], "assistant");
assert_eq!(messages[0]["content"], "Checking the weather.");
assert_eq!(messages[1]["tool_calls"][0]["id"], "call_123");
assert_eq!(
messages[1]["tool_calls"][0]["function"]["arguments"],
"{\"city\":\"Paris\"}"
);
assert_eq!(messages[2]["role"], "tool");
assert_eq!(messages[2]["tool_call_id"], "call_123");
assert_eq!(messages[2]["content"], "18C and sunny");
}
#[test]
fn chat_inputs_fold_system_and_translate_tools() {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"stream": true,
"system": [{ "type": "text", "text": "be brief" }],
"messages": [{ "role": "user", "content": "weather in Paris?" }],
"tools": [{
"name": "get_weather",
"description": "Weather lookup",
"input_schema": { "type": "object" },
}],
"stop_sequences": ["END"],
"top_k": 40,
"temperature": 0.5,
})).unwrap();
let inputs = params.chat_inputs().unwrap();
assert_eq!(inputs.max_output_tokens, Some(64));
assert_eq!(inputs.stop_sequences, Some(vec!["END".to_string()]));
assert_eq!(inputs.top_k, Some(40));
assert_eq!(inputs.temperature, Some(0.5));
assert_eq!(inputs.top_p, None);
assert_eq!(inputs.reasoning_effort, None);
assert_eq!(inputs.response_format, None);
assert_eq!(
inputs.messages,
vec![
serde_json::json!({ "role": "system", "content": "be brief" }),
serde_json::json!({ "role": "user", "content": "weather in Paris?" }),
]
);
assert_eq!(
inputs.tools,
Some(vec![serde_json::json!({
"type": "function",
"function": {
"name": "get_weather",
"description": "Weather lookup",
"parameters": { "type": "object" },
},
})])
);
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": [{ "role": "user", "content": "hi" }],
"tools": [],
})).unwrap();
assert_eq!(params.chat_inputs().unwrap().tools, None);
}
#[test]
fn chat_inputs_map_thinking_to_reasoning_effort() {
fn effort(extra: serde_json::Value) -> Option<String> {
let mut body = serde_json::json!({
"model": "@a/x",
"max_tokens": 4096,
"messages": [{ "role": "user", "content": "hi" }],
});
body.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone());
let params: MessageCreateParams = serde_json::from_value(body).unwrap();
params.chat_inputs().unwrap().reasoning_effort
}
assert_eq!(
effort(serde_json::json!({ "thinking": { "type": "disabled" } })),
Some("none".to_string())
);
assert_eq!(
effort(serde_json::json!({
"thinking": { "type": "enabled", "budget_tokens": 4096 },
})),
Some("medium".to_string())
);
assert_eq!(
effort(serde_json::json!({
"thinking": { "type": "adaptive", "display": "omitted" },
"output_config": { "effort": "max" },
})),
Some("xhigh".to_string())
);
assert_eq!(
effort(serde_json::json!({ "output_config": { "effort": "low" } })),
Some("low".to_string())
);
assert_eq!(effort(serde_json::json!({})), None);
assert_eq!(
effort(serde_json::json!({ "thinking": { "type": "telepathic" } })),
Some("medium".to_string())
);
}
#[test]
fn chat_inputs_reject_messages_without_payload() {
fn inputs(messages: serde_json::Value) -> Result<ChatInputs> {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": messages,
})).unwrap();
params.chat_inputs()
}
for content in [
serde_json::json!(" "),
serde_json::json!(""),
serde_json::json!([]),
serde_json::json!([{ "type": "text", "text": " \n" }]),
] {
let error = inputs(serde_json::json!([{ "role": "user", "content": content }]))
.err()
.unwrap_or_else(|| panic!("{content} was accepted"));
assert!(matches!(error, MunaError::InvalidInput(_)), "{error}");
assert!(error.to_string().contains("messages[0]"), "{error}");
}
assert!(matches!(
inputs(serde_json::json!([])),
Err(MunaError::InvalidInput(_))
));
inputs(serde_json::json!([
{ "role": "user", "content": "weather?" },
{ "role": "assistant", "content": [{
"type": "tool_use", "id": "t1", "name": "get_weather", "input": {}
}] },
{ "role": "user", "content": [{ "type": "tool_result", "tool_use_id": "t1" }] },
])).unwrap();
inputs(serde_json::json!([{ "role": "user", "content": " hi " }])).unwrap();
let inputs = inputs(serde_json::json!([
{ "role": "user", "content": "hi" },
{ "role": "assistant", "content": "" },
{ "role": "user", "content": "still there?" },
])).unwrap();
assert_eq!(inputs.messages[1]["role"], serde_json::json!("assistant"));
assert_eq!(inputs.messages[1]["content"], serde_json::json!(""));
}
#[test]
fn replayed_response_content_is_accepted() {
let reply = serde_json::to_value(vec![
ContentBlock::Thinking {
thinking: "The user said hello.".into(),
signature: String::new(),
},
ContentBlock::Text { text: "Hey! What can I help with?".into() },
]).unwrap();
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": [
{ "role": "user", "content": "hello" },
{ "role": "assistant", "content": reply },
{ "role": "assistant", "content": [
{ "type": "redacted_thinking", "data": "opaque" }
] },
{ "role": "user", "content": "Is this folder please" },
],
})).unwrap();
let inputs = params.chat_inputs().unwrap();
assert_eq!(inputs.messages, vec![
serde_json::json!({ "role": "user", "content": "hello" }),
serde_json::json!({ "role": "assistant", "content": "Hey! What can I help with?" }),
serde_json::json!({ "role": "user", "content": "Is this folder please" }),
]);
}
const RED_PNG_B64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAYAAABytg0kAAAAFUlEQVR4nGP8z8Dwn4GBgYEJRIAwAB8XAgICR7MUAAAAAElFTkSuQmCC";
fn image_block() -> serde_json::Value {
serde_json::json!({
"type": "image",
"source": { "type": "base64", "media_type": "image/png", "data": RED_PNG_B64 }
})
}
#[test]
fn base64_images_become_ordered_placeholders() {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": [
{ "role": "user", "content": [
image_block(),
{ "type": "text", "text": "what is this?" }
] },
{ "role": "assistant", "content": "A red square." },
{ "role": "user", "content": [
{ "type": "text", "text": "and this?" },
image_block(),
{ "type": "text", "text": "same?" }
] },
],
})).unwrap();
let inputs = params.chat_inputs().unwrap();
let images = inputs.images.unwrap();
assert_eq!(images.len(), 2);
assert_eq!((images[0].width, images[0].height), (2, 2));
assert_eq!(inputs.messages, vec![
serde_json::json!({ "role": "user", "content": [
{ "type": "image" },
{ "type": "text", "text": "what is this?" },
] }),
serde_json::json!({ "role": "assistant", "content": "A red square." }),
serde_json::json!({ "role": "user", "content": [
{ "type": "text", "text": "and this?" },
{ "type": "image" },
{ "type": "text", "text": "same?" },
] }),
]);
}
#[test]
fn images_outside_user_turns_are_rejected() {
let system: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"system": [image_block()],
"messages": [{ "role": "user", "content": "hi" }],
})).unwrap();
assert!(matches!(system.chat_inputs(), Err(MunaError::InvalidInput(m)) if m.contains("`system`")));
let inline_system: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": [
{ "role": "system", "content": [image_block()] },
{ "role": "user", "content": "hi" },
],
})).unwrap();
assert!(matches!(inline_system.chat_inputs(), Err(MunaError::InvalidInput(m)) if m.contains("`system`")));
}
#[test]
fn unsupported_blocks_are_rejected_by_name() {
fn error(block: serde_json::Value) -> String {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": [{ "role": "user", "content": [
{ "type": "text", "text": "what is this?" },
block
] }],
})).unwrap();
let error = params.chat_inputs().unwrap_err();
assert!(matches!(error, MunaError::InvalidInput(_)), "{error}");
error.to_string()
}
let url = error(serde_json::json!({
"type": "image",
"source": { "type": "url", "url": "https://example.com/cat.png" }
}));
assert!(url.contains("messages[0].content[1]") && url.contains("`url` source"), "{url}");
let file = error(serde_json::json!({
"type": "image",
"source": { "type": "file", "file_id": "file_123" }
}));
assert!(file.contains("unsupported source"), "{file}");
let tool_result_image = error(serde_json::json!({
"type": "tool_result",
"tool_use_id": "t1",
"content": [{
"type": "image",
"source": { "type": "base64", "media_type": "image/png", "data": RED_PNG_B64 }
}]
}));
assert!(tool_result_image.contains("`tool_result`"), "{tool_result_image}");
let document = error(serde_json::json!({
"type": "document",
"source": { "type": "text", "media_type": "text/plain", "data": "hi" }
}));
assert!(document.contains("`document`"), "{document}");
let unknown = error(serde_json::json!({ "type": "search_result", "title": "t" }));
assert!(unknown.contains("unrecognized type"), "{unknown}");
}
#[test]
fn billing_header_is_stripped_and_inline_system_kept_in_place() {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"system": [
{ "type": "text", "text": "x-anthropic-billing-header: cc_version=2.1 hash=abc" },
{ "type": "text", "text": "You are Claude Code." }
],
"messages": [
{ "role": "user", "content": "hi" },
{ "role": "system", "content": "x-anthropic-billing-header: hash=def" },
{ "role": "system", "content": [
{ "type": "text", "text": "x-anthropic-billing-header: hash=ghi" },
{ "type": "text", "text": "Conversation was compacted." }
] },
{ "role": "user", "content": "again" },
],
})).unwrap();
let inputs = params.chat_inputs().unwrap();
assert_eq!(
inputs.messages,
vec![
serde_json::json!({ "role": "system", "content": "You are Claude Code." }),
serde_json::json!({ "role": "user", "content": "hi" }),
serde_json::json!({ "role": "system", "content": "Conversation was compacted." }),
serde_json::json!({ "role": "user", "content": "again" }),
]
);
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"system": "x-anthropic-billing-header: hash=abc",
"messages": [{ "role": "user", "content": "hi" }],
})).unwrap();
let inputs = params.chat_inputs().unwrap();
assert_eq!(inputs.messages, vec![serde_json::json!({ "role": "user", "content": "hi" })]);
assert!(matches!(
inputs_for(serde_json::json!([
{ "role": "user", "content": "hi" },
{ "role": "system", "content": " " },
])),
Err(MunaError::InvalidInput(_))
));
}
fn inputs_for(messages: serde_json::Value) -> Result<ChatInputs> {
let params: MessageCreateParams = serde_json::from_value(serde_json::json!({
"model": "@a/x",
"max_tokens": 64,
"messages": messages,
})).unwrap();
params.chat_inputs()
}
#[test]
fn events_accumulate_to_final_message() {
let events = vec![
RawMessageStreamEvent::MessageStart {
message: Message {
id: "msg-1".to_string(),
r#type: "message".to_string(),
role: "assistant".to_string(),
content: Vec::new(),
model: "test-model".to_string(),
stop_reason: None,
stop_sequence: None,
usage: Usage {
input_tokens: Some(10),
..Default::default()
},
},
},
RawMessageStreamEvent::ContentBlockStart {
index: 0,
content_block: ContentBlock::Text {
text: String::new(),
},
},
RawMessageStreamEvent::ContentBlockDelta {
index: 0,
delta: ContentBlockDelta::TextDelta {
text: "Par".to_string(),
},
},
RawMessageStreamEvent::ContentBlockDelta {
index: 0,
delta: ContentBlockDelta::TextDelta {
text: "is.".to_string(),
},
},
RawMessageStreamEvent::ContentBlockStop { index: 0 },
RawMessageStreamEvent::MessageDelta {
delta: MessageDelta {
stop_reason: Some(StopReason::EndTurn),
stop_sequence: None,
},
usage: Usage {
input_tokens: Some(10),
output_tokens: 5,
..Default::default()
},
},
RawMessageStreamEvent::MessageStop,
];
let message = accumulate(&events);
assert!(matches!(
&message.content[0],
ContentBlock::Text { text } if text == "Paris."
));
assert_eq!(message.stop_reason, Some(StopReason::EndTurn));
assert_eq!(message.usage.input_tokens, Some(10));
assert_eq!(message.usage.output_tokens, 5);
}
#[test]
fn events_serialize_with_wire_tags() {
let event = RawMessageStreamEvent::ContentBlockDelta {
index: 0,
delta: ContentBlockDelta::TextDelta {
text: "Hi".to_string(),
},
};
let json = serde_json::to_value(&event).unwrap();
assert_eq!(json["type"], "content_block_delta");
assert_eq!(json["delta"]["type"], "text_delta");
assert_eq!(json["delta"]["text"], "Hi");
let round_trip: RawMessageStreamEvent = serde_json::from_value(json).unwrap();
assert_eq!(round_trip.event_type(), "content_block_delta");
}
}