use async_trait::async_trait;
use futures::stream::{self, Stream, StreamExt};
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tracing::Instrument;
use crate::error::{Error, Result};
use crate::middleware::{FunctionInvocationContext, LiveToolList, MiddlewarePipeline, Terminal};
use crate::tools::{FunctionInvocationConfig, ToolDefinition, ToolKind};
use crate::types::{
ChatOptions, ChatResponse, ChatResponseUpdate, Content, EmbeddingGenerationOptions,
FunctionApprovalRequestContent, FunctionApprovalResponseContent, FunctionCallContent,
FunctionResultContent, GeneratedEmbeddings, Message, Role, ToolMode, UsageContent,
UsageDetails,
};
pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatResponseUpdate>> + Send>>;
#[async_trait]
pub trait ChatClient: Send + Sync {
async fn get_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatResponse>;
async fn get_streaming_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatStream>;
fn model(&self) -> Option<&str> {
None
}
}
#[async_trait]
impl<T: ChatClient + ?Sized> ChatClient for Arc<T> {
async fn get_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatResponse> {
(**self).get_response(messages, options).await
}
async fn get_streaming_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatStream> {
(**self).get_streaming_response(messages, options).await
}
fn model(&self) -> Option<&str> {
(**self).model()
}
}
#[async_trait]
pub trait EmbeddingClient: Send + Sync {
async fn get_embeddings(
&self,
values: Vec<String>,
options: Option<EmbeddingGenerationOptions>,
) -> Result<GeneratedEmbeddings>;
fn model(&self) -> Option<&str> {
None
}
}
#[async_trait]
impl<T: EmbeddingClient + ?Sized> EmbeddingClient for Arc<T> {
async fn get_embeddings(
&self,
values: Vec<String>,
options: Option<EmbeddingGenerationOptions>,
) -> Result<GeneratedEmbeddings> {
(**self).get_embeddings(values, options).await
}
fn model(&self) -> Option<&str> {
(**self).model()
}
}
pub struct FunctionInvokingChatClient<C: ChatClient> {
inner: C,
config: FunctionInvocationConfig,
observability: crate::observability::ObservabilityConfig,
function_middleware: MiddlewarePipeline<FunctionInvocationContext>,
}
impl<C: ChatClient> FunctionInvokingChatClient<C> {
pub fn new(inner: C) -> Self {
Self {
inner,
config: FunctionInvocationConfig::default(),
function_middleware: MiddlewarePipeline::default(),
observability: crate::observability::ObservabilityConfig::from_env(),
}
}
pub fn with_observability_config(
mut self,
observability: crate::observability::ObservabilityConfig,
) -> Self {
self.observability = observability;
self
}
pub fn with_config(mut self, config: FunctionInvocationConfig) -> Self {
self.config = config;
self
}
pub fn with_function_middleware(
mut self,
middleware: Vec<Arc<crate::middleware::FunctionMiddleware>>,
) -> Self {
self.function_middleware = MiddlewarePipeline::new(middleware);
self
}
pub fn inner(&self) -> &C {
&self.inner
}
async fn inner_get_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatResponse> {
self.inner.get_response(messages, options).await
}
}
fn executable_tools(options: &ChatOptions) -> Vec<ToolDefinition> {
options
.tools
.iter()
.filter(|t| t.is_executable())
.cloned()
.collect()
}
fn is_declaration_only(tool: &ToolDefinition) -> bool {
tool.kind == ToolKind::Function && tool.executor.is_none()
}
fn accumulate_usage(aggregate: &mut Option<UsageDetails>, incoming: Option<&UsageDetails>) {
let Some(incoming) = incoming else {
return;
};
match aggregate {
Some(current) => current.add_assign(incoming),
None => *aggregate = Some(incoming.clone()),
}
}
const REJECTION_MESSAGE: &str = "Error: Tool call invocation was rejected by user.";
struct ToolCallEnv<'a> {
function_middleware: &'a MiddlewarePipeline<FunctionInvocationContext>,
session: Option<&'a crate::session::AgentSession>,
live_tools: Option<&'a LiveToolList>,
observability: &'a crate::observability::ObservabilityConfig,
}
async fn execute_tool_call(
tool: Option<ToolDefinition>,
call: &FunctionCallContent,
include_detailed_errors: bool,
terminate_on_unknown: bool,
env: ToolCallEnv<'_>,
) -> Result<(bool, FunctionResultContent)> {
let ToolCallEnv {
function_middleware,
session,
live_tools,
observability,
} = env;
match tool {
None => {
if terminate_on_unknown {
return Err(Error::tool(format!("unknown tool: {}", call.name)));
}
Ok((
true,
FunctionResultContent {
call_id: call.call_id.clone(),
result: None,
exception: Some(format!("tool '{}' not found", call.name)),
},
))
}
Some(def) => {
let args = match call.parse_arguments() {
Ok(m) => Value::Object(m.into_iter().collect()),
Err(e) => {
let msg = if include_detailed_errors {
format!("invalid tool arguments: {e}")
} else {
"invalid tool arguments".to_string()
};
return Ok((
true,
FunctionResultContent {
call_id: call.call_id.clone(),
result: None,
exception: Some(msg),
},
));
}
};
let obs_config = observability.clone();
let exec = def.executor.as_ref().unwrap().clone();
let tool_name = def.name.clone();
let description = def.description.clone();
let call_id = call.call_id.clone();
let terminal: Terminal<FunctionInvocationContext> = Box::new(move |mut ctx| {
Box::pin(async move {
if ctx.terminate {
return Ok(ctx);
}
let span = crate::observability::tool_span_ex(
&tool_name,
&call_id,
Some(&description),
);
crate::observability::record_tool_arguments(&span, &ctx.arguments, &obs_config);
#[cfg(feature = "otel-metrics")]
let started = std::time::Instant::now();
let outcome = async {
let result = exec.invoke_in_context(ctx.arguments.clone(), &ctx).await;
if let Err(e) = &result {
crate::observability::record_error(&tracing::Span::current(), e);
}
result
}
.instrument(span.clone())
.await;
#[cfg(feature = "otel-metrics")]
crate::observability::metrics::record_function_invocation_duration(
&tool_name,
started.elapsed(),
outcome
.as_ref()
.err()
.map(crate::observability::error_type)
.as_deref(),
);
if let Ok(value) = &outcome {
crate::observability::record_tool_result(&span, value, &obs_config);
}
ctx.result = Some(outcome?);
Ok(ctx)
}) as crate::tools::BoxFuture<Result<FunctionInvocationContext>>
});
let ctx = FunctionInvocationContext::new(call.name.clone(), args)
.with_session(session.cloned())
.with_tools(live_tools.cloned());
match function_middleware.execute(ctx, terminal).await {
Ok(ctx) => Ok((
false,
FunctionResultContent {
call_id: call.call_id.clone(),
result: Some(ctx.result.unwrap_or(Value::Null)),
exception: None,
},
)),
Err(e) if e.is_middleware_failure() => Err(e),
Err(e) => {
let msg = if include_detailed_errors {
format!("{e}")
} else {
"tool execution failed".to_string()
};
Ok((
true,
FunctionResultContent {
call_id: call.call_id.clone(),
result: None,
exception: Some(msg),
},
))
}
}
}
}
}
fn collect_approval_responses(messages: &[Message]) -> Vec<FunctionApprovalResponseContent> {
let mut out = Vec::new();
for msg in messages {
for content in &msg.contents {
if let Content::FunctionApprovalResponse(resp) = content {
out.push(resp.clone());
}
}
}
out
}
fn replace_approval_contents_with_results(
messages: &mut [Message],
approved_results: &HashMap<String, FunctionResultContent>,
) {
struct Outstanding {
call: FunctionCallContent,
from_request: bool,
}
let mut outstanding: Vec<Outstanding> = Vec::new();
fn retire(outstanding: &mut Vec<Outstanding>, call_id: &str) {
if call_id.is_empty() {
return;
}
if let Some(position) = outstanding
.iter()
.rposition(|entry| entry.call.call_id == call_id)
{
outstanding.remove(position);
}
}
for msg in messages.iter_mut() {
let mut to_remove: Vec<usize> = Vec::new();
let mut set_role_tool = false;
for (idx, content) in msg.contents.iter_mut().enumerate() {
match content {
Content::FunctionCall(fc) => {
if fc.call_id.is_empty() {
continue;
}
if let Some(position) = outstanding
.iter()
.position(|entry| entry.from_request && entry.call == *fc)
{
outstanding[position].from_request = false;
to_remove.push(idx);
} else {
outstanding.push(Outstanding {
call: fc.clone(),
from_request: false,
});
}
}
Content::FunctionResult(fr) => {
retire(&mut outstanding, &fr.call_id);
}
Content::FunctionApprovalRequest(req) => {
if outstanding
.iter()
.any(|entry| entry.call == req.function_call)
{
to_remove.push(idx);
} else {
if !req.function_call.call_id.is_empty() {
outstanding.push(Outstanding {
call: req.function_call.clone(),
from_request: true,
});
}
*content = Content::FunctionCall(req.function_call.clone());
}
}
Content::FunctionApprovalResponse(resp) => {
let call_id = resp.function_call.call_id.clone();
let mut answered = false;
if resp.approved {
if let Some(result) = approved_results.get(&call_id) {
*content = Content::FunctionResult(result.clone());
answered = true;
}
} else {
*content = Content::FunctionResult(FunctionResultContent {
call_id: call_id.clone(),
result: Some(Value::String(REJECTION_MESSAGE.to_string())),
exception: None,
});
answered = true;
}
if answered {
retire(&mut outstanding, &call_id);
set_role_tool = true;
}
}
_ => {}
}
}
for idx in to_remove.into_iter().rev() {
msg.contents.remove(idx);
}
if set_role_tool {
msg.role = Role::tool();
}
}
}
#[async_trait]
impl<C: ChatClient> ChatClient for FunctionInvokingChatClient<C> {
async fn get_response(
&self,
messages: Vec<Message>,
mut options: ChatOptions,
) -> Result<ChatResponse> {
let response_format = options.response_format.clone();
let mut response: ChatResponse = async move {
self.config.validate()?;
let session = options.session.take();
if !options.tools.is_empty() && options.tool_choice.is_none() {
options.tool_choice = Some(ToolMode::Auto);
}
if executable_tools(&options).is_empty() || !self.config.enabled {
return self.inner_get_response(messages, options).await;
}
let live_tools = LiveToolList::new(std::mem::take(&mut options.tools));
let mut conversation = messages;
let mut carried: Vec<Message> = Vec::new();
let mut consecutive_errors = 0usize;
let mut aggregated_usage: Option<UsageDetails> = None;
for _ in 0..self.config.max_iterations {
options.tools = live_tools.snapshot();
let tools = executable_tools(&options);
let approval_responses = collect_approval_responses(&conversation);
if !approval_responses.is_empty() {
let mut approved_results: HashMap<String, FunctionResultContent> =
HashMap::new();
let mut had_error = false;
for resp in &approval_responses {
if !resp.approved {
continue;
}
let call = &resp.function_call;
let tool = tools.iter().find(|t| t.name == call.name).cloned();
let (is_error, content) = execute_tool_call(
tool,
call,
self.config.include_detailed_errors,
self.config.terminate_on_unknown_calls,
ToolCallEnv {
function_middleware: &self.function_middleware,
session: session.as_ref(),
live_tools: Some(&live_tools),
observability: &self.observability,
},
)
.await?;
had_error |= is_error;
approved_results.insert(content.call_id.clone(), content);
}
replace_approval_contents_with_results(&mut conversation, &approved_results);
if had_error {
consecutive_errors += 1;
if consecutive_errors > self.config.max_consecutive_errors_per_request {
options.tool_choice = Some(ToolMode::None);
}
}
}
let response = self
.inner_get_response(conversation.clone(), options.clone())
.await?;
accumulate_usage(&mut aggregated_usage, response.usage_details.as_ref());
let resolved_call_ids: std::collections::HashSet<&str> = response
.messages
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_result)
.map(|fr| fr.call_id.as_str())
.collect();
let calls: Vec<_> = response
.messages
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_call)
.filter(|fc| !resolved_call_ids.contains(fc.call_id.as_str()))
.cloned()
.collect();
if calls.is_empty() {
let mut final_resp = response;
let mut msgs = std::mem::take(&mut carried);
msgs.append(&mut final_resp.messages);
final_resp.messages = msgs;
final_resp.usage_details = aggregated_usage;
return Ok(final_resp);
}
let needs_approval = calls.iter().any(|c| {
tools
.iter()
.find(|t| t.name == c.name)
.map(ToolDefinition::requires_approval)
.unwrap_or(false)
});
if needs_approval {
let mut resp = response;
let approval_contents: Vec<Content> = calls
.iter()
.map(|c| {
Content::FunctionApprovalRequest(FunctionApprovalRequestContent {
id: c.call_id.clone(),
function_call: c.clone(),
})
})
.collect();
if let Some(m) = resp
.messages
.iter_mut()
.rev()
.find(|m| m.role == Role::assistant())
{
m.contents.extend(approval_contents);
} else {
resp.messages
.push(Message::with_contents(Role::assistant(), approval_contents));
}
let mut msgs = std::mem::take(&mut carried);
msgs.append(&mut resp.messages);
resp.messages = msgs;
resp.usage_details = aggregated_usage;
return Ok(resp);
}
let has_declaration_only = calls.iter().any(|c| {
options
.tools
.iter()
.any(|t| t.name == c.name && is_declaration_only(t))
});
if has_declaration_only {
let mut resp = response;
let mut msgs = std::mem::take(&mut carried);
msgs.append(&mut resp.messages);
resp.messages = msgs;
resp.usage_details = aggregated_usage;
return Ok(resp);
}
carried.extend(response.messages.iter().cloned());
let response_conversation_id = response.conversation_id.clone();
let invocations = calls.iter().map(|call| {
let tool = tools.iter().find(|t| t.name == call.name).cloned();
let call = call.clone();
let include_detailed_errors = self.config.include_detailed_errors;
let terminate_on_unknown = self.config.terminate_on_unknown_calls;
let function_middleware = self.function_middleware.clone();
let session = session.clone();
let live_tools = live_tools.clone();
let observability = self.observability.clone();
async move {
execute_tool_call(
tool,
&call,
include_detailed_errors,
terminate_on_unknown,
ToolCallEnv {
function_middleware: &function_middleware,
session: session.as_ref(),
live_tools: Some(&live_tools),
observability: &observability,
},
)
.await
}
});
let outcomes = futures::future::try_join_all(invocations).await?;
let mut result_contents: Vec<Content> = Vec::with_capacity(outcomes.len());
let mut had_error = false;
for (is_error, content) in outcomes {
had_error |= is_error;
result_contents.push(Content::FunctionResult(content));
}
if had_error {
consecutive_errors += 1;
if consecutive_errors > self.config.max_consecutive_errors_per_request {
options.tool_choice = Some(ToolMode::None);
}
} else {
consecutive_errors = 0;
}
let tool_message = Message::with_contents(Role::tool(), result_contents);
carried.push(tool_message.clone());
match response_conversation_id {
Some(cid) => {
options.conversation_id = Some(cid);
conversation = vec![tool_message];
}
None => {
conversation.extend(response.messages);
conversation.push(tool_message);
}
}
}
options.tool_choice = Some(ToolMode::None);
let mut final_resp = self.inner_get_response(conversation, options).await?;
accumulate_usage(&mut aggregated_usage, final_resp.usage_details.as_ref());
let mut msgs = std::mem::take(&mut carried);
msgs.append(&mut final_resp.messages);
final_resp.messages = msgs;
final_resp.usage_details = aggregated_usage;
Ok(final_resp)
}
.await?;
response.try_parse_value(response_format.as_ref());
Ok(response)
}
async fn get_streaming_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatStream> {
let tools = executable_tools(&options);
if tools.is_empty() || !self.config.enabled {
return self.inner.get_streaming_response(messages, options).await;
}
let response = self.get_response(messages, options).await?;
let conversation_id = response.conversation_id.clone();
let response_id = response.response_id.clone();
let finish_reason = response.finish_reason.clone();
let usage_details = response.usage_details.clone();
let last = response.messages.len().saturating_sub(1);
let keep_provider_ids = {
let mut seen = std::collections::HashSet::new();
response.messages.iter().all(|m| {
m.message_id
.as_ref()
.is_some_and(|id| !id.is_empty() && seen.insert(id.as_str()))
})
};
let mut updates: Vec<Result<ChatResponseUpdate>> = response
.messages
.into_iter()
.enumerate()
.map(|(i, m)| {
let message_id = if keep_provider_ids {
m.message_id.clone()
} else {
Some(format!("replay-{i}"))
};
let mut contents = m.contents;
let is_last = i == last;
if is_last {
if let Some(usage) = usage_details.clone() {
contents.push(Content::Usage(UsageContent { details: usage }));
}
}
Ok(ChatResponseUpdate {
contents,
role: Some(m.role),
author_name: m.author_name,
message_id,
conversation_id: conversation_id.clone(),
response_id: response_id.clone(),
finish_reason: is_last.then(|| finish_reason.clone()).flatten(),
..Default::default()
})
})
.collect();
if updates.is_empty() && (usage_details.is_some() || finish_reason.is_some()) {
let contents = usage_details
.map(|u| vec![Content::Usage(UsageContent { details: u })])
.unwrap_or_default();
updates.push(Ok(ChatResponseUpdate {
contents,
role: Some(Role::assistant()),
conversation_id,
response_id,
finish_reason,
..Default::default()
}));
}
Ok(stream::iter(updates).boxed())
}
fn model(&self) -> Option<&str> {
self.inner.model()
}
}
#[derive(Clone)]
pub enum RetryOn {
Default,
Predicate(Arc<dyn Fn(&Error) -> bool + Send + Sync>),
}
impl std::fmt::Debug for RetryOn {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RetryOn::Default => f.write_str("RetryOn::Default"),
RetryOn::Predicate(_) => f.write_str("RetryOn::Predicate(..)"),
}
}
}
impl RetryOn {
pub fn predicate<F>(f: F) -> Self
where
F: Fn(&Error) -> bool + Send + Sync + 'static,
{
RetryOn::Predicate(Arc::new(f))
}
fn should_retry(&self, err: &Error) -> bool {
match self {
RetryOn::Default => default_should_retry(err),
RetryOn::Predicate(p) => p(err),
}
}
}
fn default_should_retry(err: &Error) -> bool {
if let Some(status) = err.status() {
return status == 408 || status == 429 || (500..600).contains(&status);
}
match err {
Error::Service(msg) => {
let m = msg.to_lowercase();
m.contains("request failed")
|| m.contains("timed out")
|| m.contains("timeout")
|| m.contains("connection")
|| m.contains("stream error")
}
_ => false,
}
}
#[derive(Clone, Debug)]
pub struct RetryPolicy {
pub max_retries: usize,
pub initial_delay: Duration,
pub max_delay: Duration,
pub backoff_multiplier: f64,
pub jitter: f64,
pub retry_on: RetryOn,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
initial_delay: Duration::from_millis(500),
max_delay: Duration::from_secs(30),
backoff_multiplier: 2.0,
jitter: 0.3,
retry_on: RetryOn::Default,
}
}
}
impl RetryPolicy {
pub fn with_max_retries(max_retries: usize) -> Self {
Self {
max_retries,
..Self::default()
}
}
pub fn initial_delay(mut self, delay: Duration) -> Self {
self.initial_delay = delay;
self
}
pub fn max_delay(mut self, delay: Duration) -> Self {
self.max_delay = delay;
self
}
pub fn backoff_multiplier(mut self, multiplier: f64) -> Self {
self.backoff_multiplier = multiplier;
self
}
pub fn jitter(mut self, jitter: f64) -> Self {
self.jitter = jitter.clamp(0.0, 1.0);
self
}
pub fn retry_on(mut self, retry_on: RetryOn) -> Self {
self.retry_on = retry_on;
self
}
fn delay_for(&self, attempt: usize, err: &Error) -> Duration {
if let Some(secs) = err.retry_after() {
let capped = secs.min(self.max_delay.as_secs_f64()).max(0.0);
return Duration::from_secs_f64(capped);
}
let exp = self.backoff_multiplier.powi((attempt - 1) as i32);
let base = self.initial_delay.as_secs_f64() * exp;
let capped = base.min(self.max_delay.as_secs_f64());
let jittered = capped * jitter_factor(self.jitter);
Duration::from_secs_f64(jittered.max(0.0))
}
}
fn jitter_factor(jitter: f64) -> f64 {
let jitter = jitter.clamp(0.0, 1.0);
if jitter == 0.0 {
return 1.0;
}
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos() as u64)
.unwrap_or(0);
let mixed = nanos ^ COUNTER.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed);
let r = (mixed >> 11) as f64 / ((1u64 << 53) as f64);
1.0 - jitter * r
}
pub struct RetryingChatClient<C: ChatClient> {
inner: C,
policy: RetryPolicy,
}
impl<C: ChatClient> RetryingChatClient<C> {
pub fn new(inner: C) -> Self {
Self {
inner,
policy: RetryPolicy::default(),
}
}
pub fn with_policy(mut self, policy: RetryPolicy) -> Self {
self.policy = policy;
self
}
pub fn inner(&self) -> &C {
&self.inner
}
pub fn policy(&self) -> &RetryPolicy {
&self.policy
}
async fn backoff(&self, attempt: usize, err: &Error) {
let delay = self.policy.delay_for(attempt, err);
tracing::warn!(
attempt,
max_retries = self.policy.max_retries,
delay_ms = delay.as_millis() as u64,
retry_after = err.retry_after(),
status = err.status(),
error = %err,
"retrying chat request after transient error"
);
tokio::time::sleep(delay).await;
}
}
#[async_trait]
impl<C: ChatClient> ChatClient for RetryingChatClient<C> {
async fn get_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatResponse> {
let mut attempt = 0usize;
loop {
match self
.inner
.get_response(messages.clone(), options.clone())
.await
{
Ok(resp) => return Ok(resp),
Err(e) => {
if attempt >= self.policy.max_retries || !self.policy.retry_on.should_retry(&e)
{
return Err(e);
}
attempt += 1;
self.backoff(attempt, &e).await;
}
}
}
}
async fn get_streaming_response(
&self,
messages: Vec<Message>,
options: ChatOptions,
) -> Result<ChatStream> {
let mut attempt = 0usize;
loop {
let established = self
.inner
.get_streaming_response(messages.clone(), options.clone())
.await;
match established {
Ok(mut stream) => match stream.next().await {
Some(Err(e))
if attempt < self.policy.max_retries
&& self.policy.retry_on.should_retry(&e) =>
{
attempt += 1;
self.backoff(attempt, &e).await;
continue;
}
Some(first) => {
let head = stream::once(async move { first });
return Ok(head.chain(stream).boxed());
}
None => return Ok(stream::empty().boxed()),
},
Err(e) => {
if attempt >= self.policy.max_retries || !self.policy.retry_on.should_retry(&e)
{
return Err(e);
}
attempt += 1;
self.backoff(attempt, &e).await;
}
}
}
}
fn model(&self) -> Option<&str> {
self.inner.model()
}
}
#[cfg(test)]
mod approval_replacement_tests {
use super::*;
use crate::types::{FunctionApprovalRequestContent, FunctionArguments};
fn call(call_id: &str) -> FunctionCallContent {
FunctionCallContent::new(call_id, "get_weather", None)
}
fn approval_request(call_id: &str) -> Message {
Message::with_contents(
Role::assistant(),
vec![Content::FunctionApprovalRequest(
FunctionApprovalRequestContent {
id: format!("req_{call_id}"),
function_call: call(call_id),
},
)],
)
}
fn function_calls(messages: &[Message]) -> Vec<&str> {
messages
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_call)
.map(|fc| fc.call_id.as_str())
.collect()
}
#[test]
fn an_approval_request_in_a_separate_message_does_not_duplicate_the_call() {
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
approval_request("c1"),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn two_approval_requests_for_one_call_expand_only_once() {
let mut messages = vec![approval_request("c1"), approval_request("c1")];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn a_single_approval_request_still_expands() {
let mut messages = vec![approval_request("c1")];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn a_second_unanswered_call_reusing_an_id_still_suppresses_the_request() {
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent {
call_id: "c1".into(),
result: Some(Value::String("sunny".into())),
exception: None,
})],
),
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
approval_request("c1"),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1", "c1"]);
}
#[test]
fn a_different_invocation_reusing_an_id_is_not_suppressed() {
let older = FunctionCallContent::new(
"c1",
"get_weather",
Some(FunctionArguments::Raw("{\"city\":\"old\"}".into())),
);
let newer = FunctionCallContent::new(
"c1",
"get_weather",
Some(FunctionArguments::Raw("{\"city\":\"new\"}".into())),
);
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(older)]),
Message::with_contents(
Role::assistant(),
vec![Content::FunctionApprovalRequest(
FunctionApprovalRequestContent {
id: "req_1".into(),
function_call: newer,
},
)],
),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
let cities: Vec<String> = messages
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_call)
.filter_map(|c| {
c.parse_arguments()
.ok()?
.get("city")?
.as_str()
.map(str::to_string)
})
.collect();
assert!(
cities.contains(&"new".to_string()),
"the new invocation must survive, got {cities:?}"
);
}
#[test]
fn two_replayed_requests_for_one_outstanding_call_both_collapse() {
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
approval_request("c1"),
approval_request("c1"),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn a_completed_approval_cycle_does_not_suppress_the_next_one() {
let approved = FunctionApprovalRequestContent {
id: "req_1".into(),
function_call: call("c1"),
}
.create_response(false);
let mut messages = vec![
approval_request("c1"),
Message::with_contents(
Role::assistant(),
vec![Content::FunctionApprovalResponse(approved)],
),
approval_request("c1"),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(
function_calls(&messages),
vec!["c1", "c1"],
"the second cycle needs its own declaration"
);
}
#[test]
fn a_result_arriving_after_the_request_still_suppresses_it() {
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
approval_request("c1"),
Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent {
call_id: "c1".into(),
result: Some(Value::String("sunny".into())),
exception: None,
})],
),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn a_request_replayed_before_its_call_still_collapses_to_one() {
let mut messages = vec![
approval_request("c1"),
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1"]);
}
#[test]
fn an_already_answered_call_does_not_suppress_a_reused_call_id() {
let mut messages = vec![
Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent {
call_id: "c1".into(),
result: Some(Value::String("sunny".into())),
exception: None,
})],
),
approval_request("c1"),
];
replace_approval_contents_with_results(&mut messages, &HashMap::new());
assert_eq!(function_calls(&messages), vec!["c1", "c1"]);
}
}