108 lines
4.0 KiB
Rust
108 lines
4.0 KiB
Rust
//! Tool call correlation — matches assistant ToolUse blocks with their
|
|
//! corresponding user ToolResult blocks by ID across a message sequence.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use crate::types::{
|
|
Content, ContentBlock, RawAssistantMessage, RawMessage, RawUserMessage, ToolCall,
|
|
ToolExchange, ToolName, ToolResultData,
|
|
};
|
|
|
|
/// Extract tool calls from an assistant message's content blocks.
|
|
fn extract_tool_calls(msg: &RawAssistantMessage) -> Vec<ToolCall> {
|
|
let source_uuid = msg.uuid.clone().unwrap_or_default();
|
|
msg.message
|
|
.content
|
|
.iter()
|
|
.filter_map(|block| {
|
|
if let ContentBlock::ToolUse { id, name, input, .. } = block {
|
|
Some(ToolCall {
|
|
id: id.clone(),
|
|
name: ToolName::from(name.clone()),
|
|
input: input.clone(),
|
|
source_message_uuid: source_uuid.clone(),
|
|
})
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Extract tool results from a user message's content blocks.
|
|
fn extract_tool_results(msg: &RawUserMessage) -> Vec<ToolResultData> {
|
|
let source_uuid = msg.uuid.clone().unwrap_or_default();
|
|
match &msg.message.content {
|
|
Content::Blocks(blocks) => blocks
|
|
.iter()
|
|
.filter_map(|block| {
|
|
if let ContentBlock::ToolResult { tool_use_id, content, is_error } = block {
|
|
// Extract text content from the tool result
|
|
let text_content = content.as_ref().and_then(|c| match c {
|
|
Content::Text(s) => Some(s.clone()),
|
|
Content::Blocks(bs) => {
|
|
// Concatenate text blocks
|
|
let texts: Vec<&str> = bs
|
|
.iter()
|
|
.filter_map(|b| {
|
|
if let ContentBlock::Text { text } = b {
|
|
Some(text.as_str())
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
if texts.is_empty() { None } else { Some(texts.join("\n")) }
|
|
}
|
|
});
|
|
Some(ToolResultData {
|
|
tool_use_id: tool_use_id.clone(),
|
|
content: text_content,
|
|
is_error: *is_error,
|
|
source_message_uuid: source_uuid.clone(),
|
|
})
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect(),
|
|
Content::Text(_) => Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Correlate tool calls with their results across a message sequence.
|
|
///
|
|
/// Iterates messages in order, collecting ToolUse blocks from assistant
|
|
/// messages and matching them by ID to ToolResult blocks in subsequent user
|
|
/// messages. Any tool calls that never received a result are emitted with
|
|
/// `result: None`.
|
|
pub fn correlate_tools(messages: &[RawMessage]) -> Vec<ToolExchange> {
|
|
let mut pending: HashMap<String, ToolCall> = HashMap::new();
|
|
let mut exchanges: Vec<ToolExchange> = Vec::new();
|
|
|
|
for msg in messages {
|
|
match msg {
|
|
RawMessage::Assistant(asst) => {
|
|
for call in extract_tool_calls(asst) {
|
|
pending.insert(call.id.clone(), call);
|
|
}
|
|
}
|
|
RawMessage::User(user) => {
|
|
for result in extract_tool_results(user) {
|
|
if let Some(call) = pending.remove(&result.tool_use_id) {
|
|
exchanges.push(ToolExchange { call, result: Some(result) });
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
// Emit unmatched calls (no result found)
|
|
for (_id, call) in pending {
|
|
exchanges.push(ToolExchange { call, result: None });
|
|
}
|
|
|
|
exchanges
|
|
}
|