Flatten partial_findings_from_history's nested for/match/if-let pyramid

Split into two filter_map-based helpers (assistant_text, tool_result_text)
and a shared truncate() so the extraction reads as a flat iterator chain
instead of four levels of nesting.
This commit is contained in:
Austin Schaefer 2026-08-18 11:30:45 +02:00
parent 187d327144
commit 3ddc48fe18

View file

@ -112,36 +112,14 @@ struct GatheredFindings {
fn partial_findings_from_history(chat_history: &[Message]) -> String {
const MAX_TOOL_RESULT_CHARS: usize = 2000;
let mut sections = Vec::new();
for message in chat_history {
match message {
Message::Assistant { content, .. } => {
for item in content.iter() {
if let AssistantContent::Text(text) = item {
sections.push(text.text().to_string());
}
}
}
Message::User { content } => {
for item in content.iter() {
if let UserContent::ToolResult(result) = item {
for part in result.content.iter() {
if let ToolResultContent::Text(text) = part {
let mut snippet = text.text().to_string();
if snippet.len() > MAX_TOOL_RESULT_CHARS {
snippet.truncate(MAX_TOOL_RESULT_CHARS);
snippet.push_str(" [...truncated]");
}
sections.push(snippet);
}
}
}
}
}
Message::System { .. } => {}
}
}
let sections: Vec<String> = chat_history
.iter()
.flat_map(|message| {
assistant_text(message)
.into_iter()
.chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS))
})
.collect();
if sections.is_empty() {
"The researcher exhausted its turn budget before gathering any usable evidence."
@ -151,6 +129,49 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String {
}
}
/// Plain-text blocks from an assistant message, if any.
fn assistant_text(message: &Message) -> Vec<String> {
let Message::Assistant { content, .. } = message else {
return Vec::new();
};
content
.iter()
.filter_map(|item| match item {
AssistantContent::Text(text) => Some(text.text().to_string()),
_ => None,
})
.collect()
}
/// Plain-text tool results attached to a user message (that's where rig
/// places them), each truncated to `max_chars` so one huge fetched page
/// can't crowd out everything else that was found.
fn tool_result_text(message: &Message, max_chars: usize) -> Vec<String> {
let Message::User { content } = message else {
return Vec::new();
};
content
.iter()
.filter_map(|item| match item {
UserContent::ToolResult(result) => Some(result),
_ => None,
})
.flat_map(|result| result.content.iter())
.filter_map(|part| match part {
ToolResultContent::Text(text) => Some(truncate(text.text(), max_chars)),
_ => None,
})
.collect()
}
fn truncate(text: &str, max_chars: usize) -> String {
if text.len() <= max_chars {
text.to_string()
} else {
format!("{} [...truncated]", &text[..max_chars])
}
}
/// Wraps the tool-calling research loop in its own span so it's visible as a
/// single unit in traces, distinct from the writing and review phases and
/// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it.