assistant2: Avoid unnecessary `String` cloning in tool use (#25725)

Marshall Bowers created 10 months ago

This PR removes some unnecessary `String` cloning in the tool use paths.

We now store the data in `Arc<str>`s for cheap cloning.

Release Notes:

- N/A

Change summary

crates/assistant2/src/thread.rs                  | 11 ++++++-----
crates/language_model/src/language_model.rs      |  2 +-
crates/language_model/src/request.rs             | 14 ++++++++------
crates/language_models/src/provider/anthropic.rs |  6 +++---
crates/language_models/src/provider/bedrock.rs   |  2 +-
5 files changed, 19 insertions(+), 16 deletions(-)

Detailed changes

crates/assistant2/src/thread.rs

@@ -636,7 +636,7 @@ impl Thread {
                             Ok(output) => {
                                 tool_results.push(LanguageModelToolResult {
                                     tool_use_id: tool_use_id.clone(),
-                                    content: output,
+                                    content: output.into(),
                                     is_error: false,
                                 });
                                 thread.pending_tool_uses_by_id.remove(&tool_use_id);
@@ -646,14 +646,15 @@ impl Thread {
                             Err(err) => {
                                 tool_results.push(LanguageModelToolResult {
                                     tool_use_id: tool_use_id.clone(),
-                                    content: err.to_string(),
+                                    content: err.to_string().into(),
                                     is_error: true,
                                 });
 
                                 if let Some(tool_use) =
                                     thread.pending_tool_uses_by_id.get_mut(&tool_use_id)
                                 {
-                                    tool_use.status = PendingToolUseStatus::Error(err.to_string());
+                                    tool_use.status =
+                                        PendingToolUseStatus::Error(err.to_string().into());
                                 }
 
                                 cx.emit(ThreadEvent::ToolFinished { tool_use_id });
@@ -716,7 +717,7 @@ pub struct PendingToolUse {
     pub id: LanguageModelToolUseId,
     /// The ID of the Assistant message in which the tool use was requested.
     pub assistant_message_id: MessageId,
-    pub name: String,
+    pub name: Arc<str>,
     pub input: serde_json::Value,
     pub status: PendingToolUseStatus,
 }
@@ -725,7 +726,7 @@ pub struct PendingToolUse {
 pub enum PendingToolUseStatus {
     Idle,
     Running { _task: Shared<Task<()>> },
-    Error(#[allow(unused)] String),
+    Error(#[allow(unused)] Arc<str>),
 }
 
 impl PendingToolUseStatus {

crates/language_model/src/language_model.rs

@@ -90,7 +90,7 @@ where
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
 pub struct LanguageModelToolUse {
     pub id: LanguageModelToolUseId,
-    pub name: String,
+    pub name: Arc<str>,
     pub input: serde_json::Value,
 }
 

crates/language_model/src/request.rs

@@ -1,4 +1,5 @@
 use std::io::{Cursor, Write};
+use std::sync::Arc;
 
 use crate::role::Role;
 use crate::{LanguageModelToolUse, LanguageModelToolUseId};
@@ -167,7 +168,7 @@ impl LanguageModelImage {
 pub struct LanguageModelToolResult {
     pub tool_use_id: LanguageModelToolUseId,
     pub is_error: bool,
-    pub content: String,
+    pub content: Arc<str>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
@@ -199,15 +200,16 @@ pub struct LanguageModelRequestMessage {
 
 impl LanguageModelRequestMessage {
     pub fn string_contents(&self) -> String {
-        let mut string_buffer = String::new();
+        let mut buffer = String::new();
         for string in self.content.iter().filter_map(|content| match content {
-            MessageContent::Text(text) => Some(text),
-            MessageContent::ToolResult(tool_result) => Some(&tool_result.content),
+            MessageContent::Text(text) => Some(text.as_str()),
+            MessageContent::ToolResult(tool_result) => Some(tool_result.content.as_ref()),
             MessageContent::ToolUse(_) | MessageContent::Image(_) => None,
         }) {
-            string_buffer.push_str(string.as_str())
+            buffer.push_str(string);
         }
-        string_buffer
+
+        buffer
     }
 
     pub fn contents_empty(&self) -> bool {

crates/language_models/src/provider/anthropic.rs

@@ -506,7 +506,7 @@ pub fn into_anthropic(
                         MessageContent::ToolUse(tool_use) => {
                             Some(anthropic::RequestContent::ToolUse {
                                 id: tool_use.id.to_string(),
-                                name: tool_use.name,
+                                name: tool_use.name.to_string(),
                                 input: tool_use.input,
                                 cache_control,
                             })
@@ -515,7 +515,7 @@ pub fn into_anthropic(
                             Some(anthropic::RequestContent::ToolResult {
                                 tool_use_id: tool_result.tool_use_id.to_string(),
                                 is_error: tool_result.is_error,
-                                content: tool_result.content,
+                                content: tool_result.content.to_string(),
                                 cache_control,
                             })
                         }
@@ -636,7 +636,7 @@ pub fn map_to_language_model_completion_events(
                                         Ok(LanguageModelCompletionEvent::ToolUse(
                                             LanguageModelToolUse {
                                                 id: tool_use.id.into(),
-                                                name: tool_use.name,
+                                                name: tool_use.name.into(),
                                                 input: if tool_use.input_json.is_empty() {
                                                     serde_json::Value::Null
                                                 } else {

crates/language_models/src/provider/bedrock.rs

@@ -731,7 +731,7 @@ pub fn map_to_language_model_completion_events(
                                                     Ok(LanguageModelCompletionEvent::ToolUse(
                                                         LanguageModelToolUse {
                                                             id: tool_use.id.into(),
-                                                            name: tool_use.name,
+                                                            name: tool_use.name.into(),
                                                             input: if tool_use.input_json.is_empty()
                                                             {
                                                                 Value::Null