Skip to content
Merged
1 change: 1 addition & 0 deletions lib/bindings/python/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

289 changes: 273 additions & 16 deletions lib/llm/src/protocols/openai/chat_completions/aggregator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,10 +181,13 @@ impl DeltaAggregator {
.collect();

// Initialize and push the converted tool calls to state_choice.tool_calls
if let Some(existing_tool_calls) = &mut state_choice.tool_calls {
existing_tool_calls.extend(converted_tool_calls);
} else {
state_choice.tool_calls = Some(converted_tool_calls);
// Only set tool_calls to Some if there are actual tool calls
if !converted_tool_calls.is_empty() {
if let Some(existing_tool_calls) = &mut state_choice.tool_calls {
existing_tool_calls.extend(converted_tool_calls);
} else {
state_choice.tool_calls = Some(converted_tool_calls);
}
}
}

Expand Down Expand Up @@ -257,6 +260,17 @@ impl From<DeltaChoice> for dynamo_async_openai::types::ChatChoice {
/// # Note
/// The `function_call` field is deprecated.
fn from(delta: DeltaChoice) -> Self {
// If tool calls are present and non-empty, finish reason should be ToolCalls
let finish_reason = if delta
.tool_calls
.as_ref()
.is_some_and(|calls| !calls.is_empty())
{
Some(dynamo_async_openai::types::FinishReason::ToolCalls)
} else {
delta.finish_reason
};

dynamo_async_openai::types::ChatChoice {
message: dynamo_async_openai::types::ChatCompletionResponseMessage {
role: delta.role.expect("delta should have a Role"),
Expand All @@ -272,7 +286,7 @@ impl From<DeltaChoice> for dynamo_async_openai::types::ChatChoice {
reasoning_content: delta.reasoning_content,
},
index: delta.index,
finish_reason: delta.finish_reason,
finish_reason,
logprobs: delta.logprobs,
}
}
Expand Down Expand Up @@ -347,7 +361,7 @@ mod tests {
tool_calls.map(|tool_calls| serde_json::from_str(tool_calls).unwrap());

let tool_call_chunks = if let Some(tool_calls) = tool_calls {
vec![
Some(vec![
dynamo_async_openai::types::ChatCompletionMessageToolCallChunk {
index: 0,
id: Some("test_id".to_string()),
Expand All @@ -357,22 +371,15 @@ mod tests {
arguments: Some(serde_json::to_string(&tool_calls["arguments"]).unwrap()),
}),
},
]
])
} else {
vec![
dynamo_async_openai::types::ChatCompletionMessageToolCallChunk {
index: 0,
id: None,
r#type: None,
function: None,
},
]
None
};

let delta = dynamo_async_openai::types::ChatCompletionStreamResponseDelta {
content: Some(text.to_string()),
function_call: None,
tool_calls: Some(tool_call_chunks),
tool_calls: tool_call_chunks,
role,
refusal: None,
reasoning_content: None,
Expand Down Expand Up @@ -625,6 +632,215 @@ mod tests {
);
}

#[tokio::test]
async fn test_tool_calling_finish_reason_override_from_stop() {
// Test that when tool calls are present but finish reason is Stop, it gets overridden to ToolCalls
let tool_call_json =
r#"{"name": "get_weather", "arguments": {"location": "New York", "unit": "celsius"}}"#;

let annotated_delta = create_test_delta(
0,
"I'll check the weather for you.",
Some(dynamo_async_openai::types::Role::Assistant),
Some(dynamo_async_openai::types::FinishReason::Stop), // Original finish reason is Stop
None,
Some(tool_call_json),
);

let data = annotated_delta.data.unwrap();
let annotated_delta = Annotated {
data: Some(data),
id: Some("test_id".to_string()),
event: None,
comment: None,
};
let stream = Box::pin(stream::iter(vec![annotated_delta]));

let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// Verify tool calls are present
assert!(choice.message.tool_calls.is_some());
let tool_calls = choice.message.tool_calls.as_ref().unwrap();
assert_eq!(tool_calls.len(), 1);

// Most importantly, verify that finish reason was overridden to ToolCalls despite original being Stop
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::ToolCalls)
);
}

#[tokio::test]
async fn test_tool_calling_finish_reason_override_from_length() {
// Test that when tool calls are present but finish reason is Length, it gets overridden to ToolCalls
let tool_call_json = r#"{"name": "search", "arguments": {"query": "rust programming"}}"#;

let annotated_delta = create_test_delta(
0,
"Let me search for that.",
Some(dynamo_async_openai::types::Role::Assistant),
Some(dynamo_async_openai::types::FinishReason::Length), // Original finish reason is Length
None,
Some(tool_call_json),
);

let data = annotated_delta.data.unwrap();
let annotated_delta = Annotated {
data: Some(data),
id: Some("test_id".to_string()),
event: None,
comment: None,
};
let stream = Box::pin(stream::iter(vec![annotated_delta]));

let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// Verify tool calls are present
assert!(choice.message.tool_calls.is_some());
let tool_calls = choice.message.tool_calls.as_ref().unwrap();
assert_eq!(tool_calls.len(), 1);

// Verify that finish reason was overridden to ToolCalls despite original being Length
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::ToolCalls)
);
}

#[tokio::test]
async fn test_tool_calling_finish_reason_override_from_none() {
// Test that when tool calls are present but finish reason is None, it gets set to ToolCalls
let tool_call_json = r#"{"name": "calculate", "arguments": {"expression": "2+2"}}"#;

let annotated_delta = create_test_delta(
0,
"I'll calculate that for you.",
Some(dynamo_async_openai::types::Role::Assistant),
None, // Original finish reason is None
None,
Some(tool_call_json),
);

let data = annotated_delta.data.unwrap();
let annotated_delta = Annotated {
data: Some(data),
id: Some("test_id".to_string()),
event: None,
comment: None,
};
let stream = Box::pin(stream::iter(vec![annotated_delta]));

let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// Verify tool calls are present
assert!(choice.message.tool_calls.is_some());
let tool_calls = choice.message.tool_calls.as_ref().unwrap();
assert_eq!(tool_calls.len(), 1);

// Verify that finish reason was set to ToolCalls despite original being None
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::ToolCalls)
);
}

#[tokio::test]
async fn test_no_tool_calling_preserves_original_finish_reason() {
// Test that when no tool calls are present, the original finish reason is preserved
let annotated_delta = create_test_delta(
0,
"This is a regular response without tool calls.",
Some(dynamo_async_openai::types::Role::Assistant),
Some(dynamo_async_openai::types::FinishReason::Stop),
None,
None, // No tool calls
);

let data = annotated_delta.data.unwrap();
let annotated_delta = Annotated {
data: Some(data),
id: Some("test_id".to_string()),
event: None,
comment: None,
};
let stream = Box::pin(stream::iter(vec![annotated_delta]));

let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// Verify no tool calls are present
assert!(choice.message.tool_calls.is_none());

// Verify that original finish reason (Stop) is preserved
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::Stop)
);
}

#[tokio::test]
async fn test_empty_tool_calls_preserves_original_finish_reason() {
// Test that when tool calls array is empty, the original finish reason is preserved
// Create a delta with empty tool calls by modifying the create_test_delta output
let mut annotated_delta = create_test_delta(
0,
"Response with empty tool calls array.",
Some(dynamo_async_openai::types::Role::Assistant),
Some(dynamo_async_openai::types::FinishReason::Length),
None,
None,
);

// Manually set empty tool calls array
if let Some(ref mut data) = annotated_delta.data {
data.choices[0].delta.tool_calls = Some(vec![]); // Empty tool calls array
}

let data = annotated_delta.data.unwrap();
let annotated_delta = Annotated {
data: Some(data),
id: Some("test_id".to_string()),
event: None,
comment: None,
};
let stream = Box::pin(stream::iter(vec![annotated_delta]));

let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// Verify tool calls array is empty
assert!(choice.message.tool_calls.is_none());

// Verify that original finish reason (Length) is preserved since tool calls are empty
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::Length)
);
}

#[tokio::test]
async fn test_tool_calling_output() {
// Simulate a delta with a tool call in the content
Expand Down Expand Up @@ -688,4 +904,45 @@ mod tests {
dynamo_async_openai::types::Role::Assistant
);
}

#[tokio::test]
async fn test_tool_calling_finish_reason_override_from_stop_alternative() {
// Test that when tool calls are present but finish reason is Stop, it gets overridden to ToolCalls
let tool_call_json =
r#"{"name": "get_weather", "arguments": {"location": "New York", "unit": "celsius"}}"#;

let annotated_delta = create_test_delta(
0,
"Getting weather for New York",
Some(dynamo_async_openai::types::Role::Assistant),
Some(dynamo_async_openai::types::FinishReason::Stop), // This should be overridden
None,
Some(tool_call_json),
);

let stream = Box::pin(stream::iter(vec![annotated_delta]));

// Call DeltaAggregator::apply
let result = DeltaAggregator::apply(stream, ParsingOptions::default()).await;

// Check the result
assert!(result.is_ok());
let response = result.unwrap();

// There should be one choice
assert_eq!(response.choices.len(), 1);
let choice = &response.choices[0];

// The finish_reason should be ToolCalls, not Stop, because tool calls are present
assert_eq!(
choice.finish_reason,
Some(dynamo_async_openai::types::FinishReason::ToolCalls)
);

// Verify tool calls are present
assert!(choice.message.tool_calls.is_some());
let tool_calls = choice.message.tool_calls.as_ref().unwrap();
assert_eq!(tool_calls.len(), 1);
assert_eq!(tool_calls[0].function.name, "get_weather");
}
}
Loading
Loading