Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

assistant2: Persist scripting tool uses in saved threads (#26404)

Browse source 
This PR makes it so the scripting tool uses are persisted to and
restored from saved threads.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-03-10 17:42:23 -04:00 • committed by GitHub
parent 6cfc4dc857
commit 082cc6184c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 52 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

44
crates/assistant2/src/tool_use.rs
View file

@ -46,25 +46,39 @@ impl ToolUseState {
}
}
pub fn from_saved_messages(messages: &[SavedMessage]) -> Self {
/// Constructs a [`ToolUseState`] from the given list of [`SavedMessage`]s.
///
/// Accepts a function to filter the tools that should be used to populate the state.
pub fn from_saved_messages(
messages: &[SavedMessage],
mut filter_by_tool_name: impl FnMut(&str) -> bool,
) -> Self {
let mut this = Self::new();
let mut tool_names_by_id = HashMap::default();
for message in messages {
match message.role {
Role::Assistant => {
if !message.tool_uses.is_empty() {
this.tool_uses_by_assistant_message.insert(
message.id,
message
.tool_uses
let tool_uses = message
.tool_uses
.iter()
.filter(|tool_use| (filter_by_tool_name)(tool_use.name.as_ref()))
.map(|tool_use| LanguageModelToolUse {
id: tool_use.id.clone(),
name: tool_use.name.clone().into(),
input: tool_use.input.clone(),
})
.collect::<Vec<_>>();
tool_names_by_id.extend(
tool_uses
.iter()
.map(|tool_use| LanguageModelToolUse {
id: tool_use.id.clone(),
name: tool_use.name.clone().into(),
input: tool_use.input.clone(),
})
.collect(),
.map(|tool_use| (tool_use.id.clone(), tool_use.name.clone())),
);
this.tool_uses_by_assistant_message
.insert(message.id, tool_uses);
}
}
Role::User => {
@ -76,6 +90,14 @@ impl ToolUseState {
for tool_result in &message.tool_results {
let tool_use_id = tool_result.tool_use_id.clone();
let Some(tool_use) = tool_names_by_id.get(&tool_use_id) else {
log::warn!("no tool name found for tool use: {tool_use_id:?}");
continue;
};
if !(filter_by_tool_name)(tool_use.as_ref()) {
continue;
}
tool_uses_by_user_message.push(tool_use_id.clone());
this.tool_results.insert(