From 8836c6fb427a3a6b2ceca18931e1c5c77f00281d Mon Sep 17 00:00:00 2001 From: Nathan Sobo Date: Wed, 23 Apr 2025 20:30:48 -0600 Subject: [PATCH] Introduce LanguageModelToolUse::raw_input (#29322) This is to enable alternative streaming solutions at the application layer. I'm not sure we really should have performed parsing of the input at this layer. Either way I want to experiment with streaming approaches in a separate crate on a branch, and this will help. /cc @maxdeviant @bennetbo @rtfeldman Closes #ISSUE Release Notes: - N/A --- crates/agent/src/tool_use.rs | 1 + crates/google_ai/src/google_ai.rs | 1 + crates/language_model/src/language_model.rs | 1 + crates/language_models/src/provider/anthropic.rs | 2 ++ crates/language_models/src/provider/bedrock.rs | 1 + crates/language_models/src/provider/copilot_chat.rs | 1 + crates/language_models/src/provider/google.rs | 5 +++++ crates/language_models/src/provider/open_ai.rs | 1 + 8 files changed, 13 insertions(+) diff --git a/crates/agent/src/tool_use.rs b/crates/agent/src/tool_use.rs index 33472773ee..a6441bb5ae 100644 --- a/crates/agent/src/tool_use.rs +++ b/crates/agent/src/tool_use.rs @@ -72,6 +72,7 @@ impl ToolUseState { .map(|tool_use| LanguageModelToolUse { id: tool_use.id.clone(), name: tool_use.name.clone().into(), + raw_input: tool_use.input.to_string(), input: tool_use.input.clone(), is_input_complete: true, }) diff --git a/crates/google_ai/src/google_ai.rs b/crates/google_ai/src/google_ai.rs index 9f6045a288..7d39fab67c 100644 --- a/crates/google_ai/src/google_ai.rs +++ b/crates/google_ai/src/google_ai.rs @@ -338,6 +338,7 @@ pub struct CountTokensResponse { #[derive(Debug, Serialize, Deserialize)] pub struct FunctionCall { pub name: String, + pub raw_args: String, pub args: serde_json::Value, } diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 8dc641f88c..5cf0895f14 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -186,6 +186,7 @@ where pub struct LanguageModelToolUse { pub id: LanguageModelToolUseId, pub name: Arc, + pub raw_input: String, pub input: serde_json::Value, pub is_input_complete: bool, } diff --git a/crates/language_models/src/provider/anthropic.rs b/crates/language_models/src/provider/anthropic.rs index f255740c03..c029cd1402 100644 --- a/crates/language_models/src/provider/anthropic.rs +++ b/crates/language_models/src/provider/anthropic.rs @@ -727,6 +727,7 @@ pub fn map_to_language_model_completion_events( id: tool_use.id.clone().into(), name: tool_use.name.clone().into(), is_input_complete: false, + raw_input: tool_use.input_json.clone(), input, }, ))], @@ -757,6 +758,7 @@ pub fn map_to_language_model_completion_events( ) .map_err(|err| anyhow!("Error parsing tool call input JSON: {err:?} - JSON string was: {input_json:?}"))? }, + raw_input: tool_use.input_json.clone(), }, )) })], diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs index bb5c10cf93..edcb604919 100644 --- a/crates/language_models/src/provider/bedrock.rs +++ b/crates/language_models/src/provider/bedrock.rs @@ -894,6 +894,7 @@ pub fn map_to_language_model_completion_events( id: tool_use.id.into(), name: tool_use.name.into(), is_input_complete: true, + raw_input: tool_use.input_json.clone(), input: if tool_use.input_json.is_empty() { Value::Null } else { diff --git a/crates/language_models/src/provider/copilot_chat.rs b/crates/language_models/src/provider/copilot_chat.rs index eac138cd61..6fe86aec39 100644 --- a/crates/language_models/src/provider/copilot_chat.rs +++ b/crates/language_models/src/provider/copilot_chat.rs @@ -368,6 +368,7 @@ pub fn map_to_language_model_completion_events( id: tool_call.id.into(), name: tool_call.name.as_str().into(), is_input_complete: true, + raw_input: tool_call.arguments.clone(), input: serde_json::Value::from_str( &tool_call.arguments, )?, diff --git a/crates/language_models/src/provider/google.rs b/crates/language_models/src/provider/google.rs index 152eef1304..c3aeb0a36c 100644 --- a/crates/language_models/src/provider/google.rs +++ b/crates/language_models/src/provider/google.rs @@ -396,6 +396,7 @@ pub fn into_google( Some(Part::FunctionCallPart(google_ai::FunctionCallPart { function_call: google_ai::FunctionCall { name: tool_use.name.to_string(), + raw_args: tool_use.raw_input, args: tool_use.input, }, })) @@ -537,6 +538,10 @@ pub fn map_to_language_model_completion_events( id, name, is_input_complete: true, + raw_input: function_call_part + .function_call + .raw_args + .clone(), input: function_call_part.function_call.args, }, ))); diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs index 11ac5394e6..773c372d16 100644 --- a/crates/language_models/src/provider/open_ai.rs +++ b/crates/language_models/src/provider/open_ai.rs @@ -491,6 +491,7 @@ pub fn map_to_language_model_completion_events( id: tool_call.id.into(), name: tool_call.name.as_str().into(), is_input_complete: true, + raw_input: tool_call.arguments.clone(), input: serde_json::Value::from_str( &tool_call.arguments, )?,