From eb9b72365f552fc030a4e928695e40ded0f6b369 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 30 May 2025 00:58:54 -0700 Subject: [PATCH] fix: chat completions API now also passes tools along --- codex-rs/core/src/chat_completions.rs | 231 ++++++++++++++++++++++---- codex-rs/core/src/client.rs | 116 +------------ codex-rs/core/src/codex.rs | 100 +++++++++-- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/openai_tools.rs | 121 ++++++++++++++ 5 files changed, 411 insertions(+), 158 deletions(-) create mode 100644 codex-rs/core/src/openai_tools.rs diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/chat_completions.rs index 7760c48fbf..42a612db10 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/chat_completions.rs @@ -25,10 +25,10 @@ use crate::flags::OPENAI_REQUEST_MAX_RETRIES; use crate::flags::OPENAI_STREAM_IDLE_TIMEOUT_MS; use crate::models::ContentItem; use crate::models::ResponseItem; +use crate::openai_tools::create_tools_json; use crate::util::backoff; -/// Implementation for the classic Chat Completions API. This is intentionally -/// minimal: we only stream back plain assistant text. +/// Implementation for the classic Chat Completions API. pub(crate) async fn stream_chat_completions( prompt: &Prompt, model: &str, @@ -42,31 +42,111 @@ pub(crate) async fn stream_chat_completions( messages.push(json!({"role": "system", "content": full_instructions})); for item in &prompt.input { - if let ResponseItem::Message { role, content } = item { - let mut text = String::new(); - for c in content { - match c { - ContentItem::InputText { text: t } | ContentItem::OutputText { text: t } => { - text.push_str(t); + match item { + ResponseItem::Message { role, content } => { + let mut text = String::new(); + for c in content { + match c { + ContentItem::InputText { text: t } + | ContentItem::OutputText { text: t } => { + text.push_str(t); + } + _ => {} } - _ => {} } + messages.push(json!({"role": role, "content": text})); + } + ResponseItem::FunctionCall { + name, + arguments, + call_id, + } => { + messages.push(json!({ + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": arguments, + } + }] + })); + } + ResponseItem::LocalShellCall { + id, + call_id: _, + status, + action, + } => { + // Confirm with API team. + messages.push(json!({ + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": id.clone().unwrap_or_else(|| "".to_string()), + "type": "local_shell_call", + "status": status, + "action": action, + }] + })); + } + ResponseItem::FunctionCallOutput { call_id, output } => { + messages.push(json!({ + "role": "tool", + "tool_call_id": call_id, + "content": output.content, + })); + } + ResponseItem::Reasoning { .. } | ResponseItem::Other => { + // Omit these items from the conversation history. + continue; } - messages.push(json!({"role": role, "content": text})); } } + let tools_json = create_tools_json(prompt, model)?; + // create_tools_json() returns JSON values that are compatible with + // Function Calling in the Responses API: + // https://platform.openai.com/docs/guides/function-calling?api-mode=responses + // So we must rewrite "tools" to match the chat completions tool call format: + // https://platform.openai.com/docs/guides/function-calling?api-mode=chat + let tools_json = tools_json + .into_iter() + .filter_map(|mut tool| { + if tool.get("type") != Some(&serde_json::Value::String("function".to_string())) { + return None; + } + + if let Some(map) = tool.as_object_mut() { + // Remove "type" field as it is not needed in chat completions. + map.remove("type"); + Some(json!({ + "type": "function", + "function": map, + })) + } else { + None + } + }) + .collect::>(); + let payload = json!({ "model": model, "messages": messages, - "stream": true + "stream": true, + "tools": tools_json, }); let base_url = provider.base_url.trim_end_matches('/'); let url = format!("{}/chat/completions", base_url); debug!(url, "POST (chat)"); - trace!("request payload: {}", payload); + trace!( + "request payload: {}", + serde_json::to_string_pretty(&payload).unwrap_or_default() + ); let api_key = provider.api_key()?; let mut attempt = 0; @@ -134,6 +214,21 @@ where let idle_timeout = *OPENAI_STREAM_IDLE_TIMEOUT_MS; + // State to accumulate a function call across streaming chunks. + // OpenAI may split the `arguments` string over multiple `delta` events + // until the chunk whose `finish_reason` is `tool_calls` is emitted. We + // keep collecting the pieces here and forward a single + // `ResponseItem::FunctionCall` once the call is complete. + #[derive(Default)] + struct FunctionCallState { + name: Option, + arguments: String, + call_id: Option, + active: bool, + } + + let mut fn_call_state = FunctionCallState::default(); + loop { let sse = match timeout(idle_timeout, stream.next()).await { Ok(Some(Ok(ev))) => ev, @@ -173,23 +268,89 @@ where Ok(v) => v, Err(_) => continue, }; + trace!("chat_completions received SSE chunk: {chunk:?}"); - let content_opt = chunk - .get("choices") - .and_then(|c| c.get(0)) - .and_then(|c| c.get("delta")) - .and_then(|d| d.get("content")) - .and_then(|c| c.as_str()); + let choice_opt = chunk.get("choices").and_then(|c| c.get(0)); - if let Some(content) = content_opt { - let item = ResponseItem::Message { - role: "assistant".to_string(), - content: vec![ContentItem::OutputText { - text: content.to_string(), - }], - }; + if let Some(choice) = choice_opt { + // Handle assistant content tokens. + if let Some(content) = choice + .get("delta") + .and_then(|d| d.get("content")) + .and_then(|c| c.as_str()) + { + let item = ResponseItem::Message { + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: content.to_string(), + }], + }; - let _ = tx_event.send(Ok(ResponseEvent::OutputItemDone(item))).await; + let _ = tx_event.send(Ok(ResponseEvent::OutputItemDone(item))).await; + } + + // Handle streaming function / tool calls. + if let Some(tool_calls) = choice + .get("delta") + .and_then(|d| d.get("tool_calls")) + .and_then(|tc| tc.as_array()) + { + if let Some(tool_call) = tool_calls.first() { + // Mark that we have an active function call in progress. + fn_call_state.active = true; + + // Extract call_id if present. + if let Some(id) = tool_call.get("id").and_then(|v| v.as_str()) { + fn_call_state.call_id.get_or_insert_with(|| id.to_string()); + } + + // Extract function details if present. + if let Some(function) = tool_call.get("function") { + if let Some(name) = function.get("name").and_then(|n| n.as_str()) { + fn_call_state.name.get_or_insert_with(|| name.to_string()); + } + + if let Some(args_fragment) = + function.get("arguments").and_then(|a| a.as_str()) + { + fn_call_state.arguments.push_str(args_fragment); + } + } + } + } + + // Emit end-of-turn when finish_reason signals completion. + if let Some(finish_reason) = choice.get("finish_reason").and_then(|v| v.as_str()) { + match finish_reason { + "tool_calls" if fn_call_state.active => { + // Build the FunctionCall response item. + let item = ResponseItem::FunctionCall { + name: fn_call_state.name.clone().unwrap_or_else(|| "".to_string()), + arguments: fn_call_state.arguments.clone(), + call_id: fn_call_state.call_id.clone().unwrap_or_else(String::new), + }; + + // Emit it downstream. + let _ = tx_event.send(Ok(ResponseEvent::OutputItemDone(item))).await; + } + "stop" => { + // Regular turn without tool-call. + } + _ => {} + } + + // Emit Completed regardless of reason so the agent can advance. + let _ = tx_event + .send(Ok(ResponseEvent::Completed { + response_id: String::new(), + })) + .await; + + // Prepare for potential next turn (should not happen in same stream). + // fn_call_state = FunctionCallState::default(); + + return; // End processing for this SSE stream. + } } } } @@ -236,9 +397,14 @@ where Poll::Ready(None) => return Poll::Ready(None), Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))) => { - // Accumulate *assistant* text but do not emit yet. - if let crate::models::ResponseItem::Message { role, content } = &item { - if role == "assistant" { + // If this is an incremental assistant message chunk, accumulate but + // do NOT emit yet. Forward any other item (e.g. FunctionCall) right + // away so downstream consumers see it. + + let is_assistant_delta = matches!(&item, crate::models::ResponseItem::Message { role, .. } if role == "assistant"); + + if is_assistant_delta { + if let crate::models::ResponseItem::Message { content, .. } = &item { if let Some(text) = content.iter().find_map(|c| match c { crate::models::ContentItem::OutputText { text } => Some(text), _ => None, @@ -246,10 +412,13 @@ where this.cumulative.push_str(text); } } + + // Swallow partial assistant chunk; keep polling. + continue; } - // Swallow partial event; keep polling. - continue; + // Not an assistant message – forward immediately. + return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))); } Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))) => { if !this.cumulative.is_empty() { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 72ce845fc8..f983662c3e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1,7 +1,5 @@ -use std::collections::BTreeMap; use std::io::BufRead; use std::path::Path; -use std::sync::LazyLock; use std::time::Duration; use bytes::Bytes; @@ -11,7 +9,6 @@ use reqwest::StatusCode; use serde::Deserialize; use serde::Serialize; use serde_json::Value; -use serde_json::json; use tokio::sync::mpsc; use tokio::time::timeout; use tokio_util::io::ReaderStream; @@ -36,71 +33,9 @@ use crate::flags::OPENAI_STREAM_IDLE_TIMEOUT_MS; use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::WireApi; use crate::models::ResponseItem; +use crate::openai_tools::create_tools_json; use crate::util::backoff; -/// When serialized as JSON, this produces a valid "Tool" in the OpenAI -/// Responses API. -#[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, - description: &'static str, - strict: bool, - parameters: JsonSchema, -} - -/// Generic JSON‑Schema subset needed for our tool definitions -#[derive(Debug, Clone, Serialize)] -#[serde(tag = "type", rename_all = "lowercase")] -enum JsonSchema { - String, - Number, - Array { - items: Box, - }, - Object { - properties: BTreeMap, - required: &'static [&'static str], - #[serde(rename = "additionalProperties")] - additional_properties: bool, - }, -} - -/// Tool usage specification -static DEFAULT_TOOLS: LazyLock> = LazyLock::new(|| { - let mut properties = BTreeMap::new(); - properties.insert( - "command".to_string(), - JsonSchema::Array { - items: Box::new(JsonSchema::String), - }, - ); - properties.insert("workdir".to_string(), JsonSchema::String); - properties.insert("timeout".to_string(), JsonSchema::Number); - - vec![OpenAiTool::Function(ResponsesApiTool { - name: "shell", - description: "Runs a shell command, and returns its output.", - strict: false, - parameters: JsonSchema::Object { - properties, - required: &["command"], - additional_properties: false, - }, - })] -}); - -static DEFAULT_CODEX_MODEL_TOOLS: LazyLock> = - LazyLock::new(|| vec![OpenAiTool::LocalShell {}]); - #[derive(Clone)] pub struct ModelClient { model: String, @@ -161,27 +96,8 @@ impl ModelClient { return stream_from_fixture(path).await; } - // Assemble tool list: built-in tools + any extra tools from the prompt. - 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( - prompt - .extra_tools - .clone() - .into_iter() - .map(|(name, tool)| mcp_tool_to_openai_tool(name, tool)), - ); - - debug!("tools_json: {}", serde_json::to_string_pretty(&tools_json)?); - let full_instructions = prompt.get_full_instructions(); + let tools_json = create_tools_json(prompt, &self.model)?; let payload = Payload { model: &self.model, instructions: &full_instructions, @@ -276,34 +192,6 @@ impl ModelClient { } } -fn mcp_tool_to_openai_tool( - fully_qualified_name: String, - tool: mcp_types::Tool, -) -> serde_json::Value { - let mcp_types::Tool { - description, - mut input_schema, - .. - } = tool; - - // OpenAI models mandate the "properties" field in the schema. The Agents - // SDK fixed this by inserting an empty object for "properties" if it is not - // already present https://github.com/openai/openai-agents-python/issues/449 - // so here we do the same. - if input_schema.properties.is_none() { - input_schema.properties = Some(serde_json::Value::Object(serde_json::Map::new())); - } - - // TODO(mbolin): Change the contract of this function to return - // ResponsesApiTool. - json!({ - "name": fully_qualified_name, - "description": description, - "parameters": input_schema, - "type": "function", - }) -} - #[derive(Debug, Deserialize, Serialize)] struct SseEvent { #[serde(rename = "type")] diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2699a9ce78..c58362fd9f 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -20,6 +20,7 @@ use codex_apply_patch::MaybeApplyPatchVerified; use codex_apply_patch::maybe_parse_apply_patch_verified; use codex_apply_patch::print_summary; use futures::prelude::*; +use mcp_types::CallToolResult; use serde::Serialize; use serde_json; use tokio::sync::Notify; @@ -388,7 +389,7 @@ impl Session { tool: &str, arguments: Option, timeout: Option, - ) -> anyhow::Result { + ) -> anyhow::Result { self.mcp_connection_manager .call_tool(server, tool, arguments, timeout) .await @@ -775,6 +776,7 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { let mut pending_response_input: Vec = vec![ResponseInputItem::from(input)]; let last_agent_message: Option; loop { + debug!("pending_response_input: {pending_response_input:?}"); let mut net_new_turn_input = pending_response_input .drain(..) .map(ResponseItem::from) @@ -828,31 +830,102 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { }) }) .collect(); + debug!("Turn input: {turn_input:?}"); match run_turn(&sess, sub_id.clone(), turn_input).await { Ok(turn_output) => { - let (items, responses): (Vec<_>, Vec<_>) = turn_output - .into_iter() - .map(|p| (p.item, p.response)) - .unzip(); - let responses = responses - .into_iter() - .flatten() - .collect::>(); + let mut items_to_record_to_conversation_history = Vec::::new(); + let mut responses = Vec::::new(); + for processed_response_item in turn_output { + let ProcessedResponseItem { item, response } = processed_response_item; + match (&item, &response) { + (ResponseItem::Message { role, content, .. }, None) + if role == "assistant" => + { + // If the model returned a message, we need to record it. + items_to_record_to_conversation_history.push(ResponseItem::Message { + content: content.clone(), + role: "assistant".to_string(), + }); + } + ( + ResponseItem::LocalShellCall { .. }, + Some(ResponseInputItem::FunctionCallOutput { call_id, output }), + ) => { + items_to_record_to_conversation_history.push(item); + items_to_record_to_conversation_history.push( + ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: output.clone(), + }, + ); + } + ( + ResponseItem::FunctionCall { .. }, + Some(ResponseInputItem::FunctionCallOutput { call_id, output }), + ) => { + items_to_record_to_conversation_history.push(item); + items_to_record_to_conversation_history.push( + ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: output.clone(), + }, + ); + } + ( + ResponseItem::FunctionCall { .. }, + Some(ResponseInputItem::McpToolCallOutput { call_id, result }), + ) => { + items_to_record_to_conversation_history.push(item); + // let (content, success): (String, Option) = match result { + // Ok(CallToolResult { content, is_error }) => { + // (content, is_error.or_else(false)) + // } + // Err(e) => (e.clone(), Some(true)), + // }; + // items_to_record_to_conversation_history.push( + // ResponseItem::FunctionCallOutput { + // call_id: call_id.clone(), + // output: FunctionCallOutputPayload { content, success }, + // }, + // ); + warn!( + "Skipping MCP tool call output: {call_id:?} with response: {result:?}" + ); + } + _ => { + warn!("Unexpected response item: {item:?} with response: {response:?}"); + } + }; + if let Some(response) = response { + responses.push(response); + } + } // Only attempt to take the lock if there is something to record. - if !items.is_empty() { + if !items_to_record_to_conversation_history.is_empty() { // First persist model-generated output to the rollout file – this only borrows. - sess.record_rollout_items(&items).await; + sess.record_rollout_items(&items_to_record_to_conversation_history) + .await; + + debug!( + "has transcript? {}", + sess.state.lock().unwrap().zdr_transcript.is_some() + ); // For ZDR we also need to keep a transcript clone. if let Some(transcript) = sess.state.lock().unwrap().zdr_transcript.as_mut() { - transcript.record_items(&items); + debug!( + "Recording items to transcript: {items_to_record_to_conversation_history:?}" + ); + transcript.record_items(&items_to_record_to_conversation_history); } } if responses.is_empty() { debug!("Turn completed"); - last_agent_message = get_last_assistant_message_from_turn(&items); + last_agent_message = get_last_assistant_message_from_turn( + &items_to_record_to_conversation_history, + ); sess.maybe_notify(UserNotification::AgentTurnComplete { turn_id: sub_id.clone(), input_messages: turn_input_messages, @@ -959,6 +1032,7 @@ async fn run_turn( /// events map to a `ResponseItem`. A `ResponseItem` may need to be /// "handled" such that it produces a `ResponseInputItem` that needs to be /// sent back to the model on the next turn. +#[derive(Debug)] struct ProcessedResponseItem { item: ResponseItem, response: Option, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 8398ff7650..77941a9a51 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -27,6 +27,7 @@ mod model_provider_info; pub use model_provider_info::ModelProviderInfo; pub use model_provider_info::WireApi; mod models; +mod openai_tools; mod project_doc; pub mod protocol; mod rollout; diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs new file mode 100644 index 0000000000..3bb4cd1a13 --- /dev/null +++ b/codex-rs/core/src/openai_tools.rs @@ -0,0 +1,121 @@ +use serde::Serialize; +use serde_json::json; +use std::collections::BTreeMap; +use std::sync::LazyLock; + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct ResponsesApiTool { + name: &'static str, + description: &'static str, + strict: bool, + parameters: JsonSchema, +} + +/// When serialized as JSON, this produces a valid "Tool" in the OpenAI +/// Responses API. +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type")] +pub(crate) enum OpenAiTool { + #[serde(rename = "function")] + Function(ResponsesApiTool), + #[serde(rename = "local_shell")] + LocalShell {}, +} + +/// Generic JSON‑Schema subset needed for our tool definitions +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub(crate) enum JsonSchema { + String, + Number, + Array { + items: Box, + }, + Object { + properties: BTreeMap, + required: &'static [&'static str], + #[serde(rename = "additionalProperties")] + additional_properties: bool, + }, +} + +/// Tool usage specification +static DEFAULT_TOOLS: LazyLock> = LazyLock::new(|| { + let mut properties = BTreeMap::new(); + properties.insert( + "command".to_string(), + JsonSchema::Array { + items: Box::new(JsonSchema::String), + }, + ); + properties.insert("workdir".to_string(), JsonSchema::String); + properties.insert("timeout".to_string(), JsonSchema::Number); + + vec![OpenAiTool::Function(ResponsesApiTool { + name: "shell", + description: "Runs a shell command, and returns its output.", + strict: false, + parameters: JsonSchema::Object { + properties, + required: &["command"], + additional_properties: false, + }, + })] +}); + +static DEFAULT_CODEX_MODEL_TOOLS: LazyLock> = + LazyLock::new(|| vec![OpenAiTool::LocalShell {}]); + +pub(crate) fn create_tools_json( + prompt: &crate::client_common::Prompt, + model: &str, +) -> crate::error::Result> { + // Assemble tool list: built-in tools + any extra tools from the prompt. + let default_tools = if 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( + prompt + .extra_tools + .clone() + .into_iter() + .map(|(name, tool)| mcp_tool_to_openai_tool(name, tool)), + ); + + tracing::debug!("tools_json: {}", serde_json::to_string_pretty(&tools_json)?); + Ok(tools_json) +} + +fn mcp_tool_to_openai_tool( + fully_qualified_name: String, + tool: mcp_types::Tool, +) -> serde_json::Value { + let mcp_types::Tool { + description, + mut input_schema, + .. + } = tool; + + // OpenAI models mandate the "properties" field in the schema. The Agents + // SDK fixed this by inserting an empty object for "properties" if it is not + // already present https://github.com/openai/openai-agents-python/issues/449 + // so here we do the same. + if input_schema.properties.is_none() { + input_schema.properties = Some(serde_json::Value::Object(serde_json::Map::new())); + } + + // TODO(mbolin): Change the contract of this function to return + // ResponsesApiTool. + json!({ + "name": fully_qualified_name, + "description": description, + "parameters": input_schema, + "type": "function", + }) +}