feat: add support for OpenAI tool type, local_shell

This commit is contained in:
Michael Bolin
2025-05-16 12:18:00 -07:00
parent 89ed7e54ac
commit 8470d27eeb
5 changed files with 112 additions and 33 deletions

View File

@@ -40,10 +40,18 @@ use crate::util::backoff;
/// When serialized as JSON, this produces a valid "Tool" in the OpenAI
/// Responses API.
#[derive(Debug, Serialize)]
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
enum OpenAiTool {
#[serde(rename = "function")]
Function(ResponsesApiTool),
#[serde(rename = "local_shell")]
LocalShell {},
}
#[derive(Debug, Clone, Serialize)]
struct ResponsesApiTool {
name: &'static str,
r#type: &'static str, // "function"
description: &'static str,
strict: bool,
parameters: JsonSchema,
@@ -67,7 +75,7 @@ enum JsonSchema {
}
/// Tool usage specification
static DEFAULT_TOOLS: LazyLock<Vec<ResponsesApiTool>> = LazyLock::new(|| {
static DEFAULT_TOOLS: LazyLock<Vec<OpenAiTool>> = LazyLock::new(|| {
let mut properties = BTreeMap::new();
properties.insert(
"command".to_string(),
@@ -78,9 +86,8 @@ static DEFAULT_TOOLS: LazyLock<Vec<ResponsesApiTool>> = LazyLock::new(|| {
properties.insert("workdir".to_string(), JsonSchema::String);
properties.insert("timeout".to_string(), JsonSchema::Number);
vec![ResponsesApiTool {
vec![OpenAiTool::Function(ResponsesApiTool {
name: "shell",
r#type: "function",
description: "Runs a shell command, and returns its output.",
strict: false,
parameters: JsonSchema::Object {
@@ -88,9 +95,12 @@ static DEFAULT_TOOLS: LazyLock<Vec<ResponsesApiTool>> = LazyLock::new(|| {
required: &["command"],
additional_properties: false,
},
}]
})]
});
static DEFAULT_CODEX_MODEL_TOOLS: LazyLock<Vec<OpenAiTool>> =
LazyLock::new(|| vec![OpenAiTool::LocalShell {}]);
#[derive(Clone)]
pub struct ModelClient {
model: String,
@@ -152,8 +162,13 @@ impl ModelClient {
}
// Assemble tool list: built-in tools + any extra tools from the prompt.
let mut tools_json = Vec::with_capacity(DEFAULT_TOOLS.len() + prompt.extra_tools.len());
for t in DEFAULT_TOOLS.iter() {
let default_tools = if self.model.starts_with("codex") {
&DEFAULT_CODEX_MODEL_TOOLS
} else {
&DEFAULT_TOOLS
};
let mut tools_json = Vec::with_capacity(default_tools.len() + prompt.extra_tools.len());
for t in default_tools.iter() {
tools_json.push(serde_json::to_value(t)?);
}
tools_json.extend(

View File

@@ -51,6 +51,7 @@ use crate::mcp_connection_manager::try_parse_fully_qualified_tool_name;
use crate::mcp_tool_call::handle_mcp_tool_call;
use crate::models::ContentItem;
use crate::models::FunctionCallOutputPayload;
use crate::models::LocalShellAction;
use crate::models::ReasoningItemReasoningSummary;
use crate::models::ResponseInputItem;
use crate::models::ResponseItem;
@@ -1022,10 +1023,44 @@ async fn handle_response_item(
arguments,
call_id,
} => {
tracing::info!("FunctionCall: {arguments}");
output = Some(
handle_function_call(sess, sub_id.to_string(), name, arguments, call_id).await,
);
}
ResponseItem::LocalShellCall {
id,
call_id,
status: _,
action,
} => {
let LocalShellAction::Exec(action) = action;
tracing::info!("LocalShellCall: {action:?}");
let params = ShellToolCallParams {
command: action.command,
workdir: action.working_directory,
timeout_ms: action.timeout_ms,
};
let effective_call_id = match (call_id, id) {
(Some(call_id), _) => call_id,
(None, Some(id)) => id,
(None, None) => {
error!("LocalShellCall without call_id or id");
todo!("Respond to model to tell it about this error");
}
};
let exec_params = to_exec_params(params, sess);
output = Some(
handle_container_exec_with_params(
exec_params,
sess,
sub_id.to_string(),
effective_call_id,
)
.await,
)
}
ResponseItem::FunctionCallOutput { .. } => {
debug!("unexpected FunctionCallOutput from stream");
}
@@ -1043,7 +1078,13 @@ async fn handle_function_call(
) -> ResponseInputItem {
match name.as_str() {
"container.exec" | "shell" => {
handle_container_exec_function_call(sess, sub_id, arguments, call_id).await
let params = match parse_container_exec_arguments(arguments, sess, &call_id) {
Ok(params) => params,
Err(output) => {
return output;
}
};
handle_container_exec_with_params(params, sess, sub_id, call_id).await
}
_ => {
match try_parse_fully_qualified_tool_name(&name) {
@@ -1070,6 +1111,14 @@ async fn handle_function_call(
}
}
fn to_exec_params(params: ShellToolCallParams, sess: &Session) -> ExecParams {
ExecParams {
command: params.command,
cwd: sess.resolve_path(params.workdir.clone()),
timeout_ms: params.timeout_ms,
}
}
fn parse_container_exec_arguments(
arguments: String,
sess: &Session,
@@ -1077,11 +1126,7 @@ fn parse_container_exec_arguments(
) -> Result<ExecParams, ResponseInputItem> {
// parse command
match serde_json::from_str::<ShellToolCallParams>(&arguments) {
Ok(shell_tool_call_params) => Ok(ExecParams {
command: shell_tool_call_params.command,
cwd: sess.resolve_path(shell_tool_call_params.workdir.clone()),
timeout_ms: shell_tool_call_params.timeout_ms,
}),
Ok(shell_tool_call_params) => Ok(to_exec_params(shell_tool_call_params, sess)),
Err(e) => {
// allow model to re-sample
let output = ResponseInputItem::FunctionCallOutput {
@@ -1096,22 +1141,6 @@ fn parse_container_exec_arguments(
}
}
async fn handle_container_exec_function_call(
sess: &Session,
sub_id: String,
arguments: String,
call_id: String,
) -> ResponseInputItem {
let params = match parse_container_exec_arguments(arguments, sess, &call_id) {
Ok(params) => params,
Err(output) => {
return output;
}
};
handle_container_exec_with_params(params, sess, sub_id, call_id).await
}
async fn handle_container_exec_with_params(
params: ExecParams,
sess: &Session,

View File

@@ -41,8 +41,9 @@ impl ConversationHistory {
fn is_api_message(message: &ResponseItem) -> bool {
match message {
ResponseItem::Message { role, .. } => role.as_str() != "system",
ResponseItem::FunctionCall { .. } => true,
ResponseItem::FunctionCallOutput { .. } => true,
_ => false,
ResponseItem::FunctionCallOutput { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::LocalShellCall { .. } => true,
ResponseItem::Reasoning { .. } | ResponseItem::Other => false,
}
}

View File

@@ -1,3 +1,5 @@
use std::collections::HashMap;
use base64::Engine;
use serde::Deserialize;
use serde::Serialize;
@@ -37,6 +39,14 @@ pub enum ResponseItem {
id: String,
summary: Vec<ReasoningItemReasoningSummary>,
},
LocalShellCall {
/// Set when using the chat completions API.
id: Option<String>,
/// Set when using the Responses API.
call_id: Option<String>,
status: LocalShellStatus,
action: LocalShellAction,
},
FunctionCall {
name: String,
// The Responses API returns the function call arguments as a *string* that contains
@@ -71,6 +81,29 @@ impl From<ResponseInputItem> for ResponseItem {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LocalShellStatus {
Completed,
InProgress,
Incomplete,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum LocalShellAction {
Exec(LocalShellExecAction),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalShellExecAction {
pub command: Vec<String>,
pub timeout_ms: Option<u64>,
pub working_directory: Option<String>,
pub env: Option<HashMap<String, String>>,
pub user: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ReasoningItemReasoningSummary {

View File

@@ -115,6 +115,7 @@ impl RolloutRecorder {
// "fully qualified MCP tool calls," so we could consider
// reformatting them in that case.
ResponseItem::Message { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::FunctionCallOutput { .. } => {}
ResponseItem::Reasoning { .. } | ResponseItem::Other => {