diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index dd177a7..300ecaf 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -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 = 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 { + 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 { + 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.