diff --git a/crates/agent-core/src/dialect/anthropic.rs b/crates/agent-core/src/dialect/anthropic.rs new file mode 100644 index 0000000..b7a702b --- /dev/null +++ b/crates/agent-core/src/dialect/anthropic.rs @@ -0,0 +1,246 @@ +//! Anthropic Messages wire (`/v1/messages`). +//! +//! The harness's internal model was chosen to be Anthropic-shaped (content blocks, +//! `tool_use`/`tool_result`), so the request mapping is nearly an identity and the SSE decoder is a +//! direct translation of Anthropic's `content_block_*` / `message_*` events. + +use serde_json::{Map, Value}; + +use super::StreamDecoder; +use crate::message::{StopReason, StreamEvent}; +use crate::transport::ModelRequest; + +/// Build the streaming request body. `system` is hoisted to a top-level field (Anthropic keeps it +/// out of `messages`); `messages` and `tools` serialize straight from the internal model. +pub fn build_body(req: &ModelRequest) -> Value { + let mut map = Map::new(); + map.insert("model".into(), Value::String(req.model.clone())); + map.insert("max_tokens".into(), Value::from(req.max_tokens)); + map.insert("stream".into(), Value::Bool(true)); + map.insert( + "messages".into(), + serde_json::to_value(&req.messages).unwrap_or(Value::Null), + ); + if let Some(system) = &req.system { + map.insert("system".into(), Value::String(system.clone())); + } + if !req.tools.is_empty() { + map.insert( + "tools".into(), + serde_json::to_value(&req.tools).unwrap_or(Value::Null), + ); + } + Value::Object(map) +} + +fn map_stop_reason(s: Option<&str>) -> StopReason { + match s { + Some("end_turn") => StopReason::EndTurn, + Some("tool_use") => StopReason::ToolUse, + Some("max_tokens") => StopReason::MaxTokens, + Some("stop_sequence") => StopReason::StopSequence, + _ => StopReason::Other, + } +} + +/// Decodes Anthropic SSE. Tracks token counts (input from `message_start`, output from +/// `message_delta`) and the stop reason, emitting a single `Usage` + `MessageStop` at `message_stop`. +#[derive(Default)] +pub struct Decoder { + input_tokens: u32, + output_tokens: u32, + stop_reason: StopReason, +} + +impl StreamDecoder for Decoder { + fn push(&mut self, data: &Value) -> Vec { + let kind = data.get("type").and_then(Value::as_str).unwrap_or(""); + match kind { + "message_start" => { + let usage = data.get("message").and_then(|m| m.get("usage")); + self.input_tokens = u32_at(usage, "input_tokens"); + self.output_tokens = u32_at(usage, "output_tokens"); + vec![StreamEvent::MessageStart] + } + "content_block_start" => { + let block = data.get("content_block"); + if block.and_then(|b| b.get("type")).and_then(Value::as_str) == Some("tool_use") { + let id = str_at(block, "id").to_string(); + let name = str_at(block, "name").to_string(); + vec![StreamEvent::ToolUseStart { id, name }] + } else { + // A text block opening carries no event in our model — text accrues via deltas. + Vec::new() + } + } + "content_block_delta" => { + let delta = data.get("delta"); + match delta.and_then(|d| d.get("type")).and_then(Value::as_str) { + Some("text_delta") => { + vec![StreamEvent::TextDelta { + text: str_at(delta, "text").to_string(), + }] + } + Some("input_json_delta") => { + vec![StreamEvent::InputJsonDelta { + partial_json: str_at(delta, "partial_json").to_string(), + }] + } + _ => Vec::new(), + } + } + "content_block_stop" => vec![StreamEvent::ContentBlockStop], + "message_delta" => { + let delta = data.get("delta"); + self.stop_reason = map_stop_reason( + delta + .and_then(|d| d.get("stop_reason")) + .and_then(Value::as_str), + ); + let out = u32_at(data.get("usage"), "output_tokens"); + if out > 0 { + self.output_tokens = out; + } + Vec::new() + } + "message_stop" => vec![ + StreamEvent::Usage { + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + }, + StreamEvent::MessageStop { + stop_reason: self.stop_reason, + }, + ], + _ => Vec::new(), + } + } +} + +fn str_at<'a>(v: Option<&'a Value>, key: &str) -> &'a str { + v.and_then(|v| v.get(key)) + .and_then(Value::as_str) + .unwrap_or("") +} + +fn u32_at(v: Option<&Value>, key: &str) -> u32 { + v.and_then(|v| v.get(key)) + .and_then(Value::as_u64) + .unwrap_or(0) as u32 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dialect::decode_sse; + use crate::message::{ContentBlock, Message, ToolDef}; + use serde_json::json; + + #[test] + fn build_body_hoists_system_and_keeps_blocks() { + let req = ModelRequest::new("claude-opus-4-8", vec![Message::user("hi")], 256) + .with_system("be brief") + .with_tools(vec![ToolDef { + name: "read".into(), + description: "read a file".into(), + input_schema: json!({ "type": "object" }), + }]); + let body = build_body(&req); + assert_eq!(body["model"], "claude-opus-4-8"); + assert_eq!(body["stream"], true); + assert_eq!(body["system"], "be brief"); + assert_eq!(body["messages"][0]["role"], "user"); + assert_eq!(body["messages"][0]["content"][0]["type"], "text"); + assert_eq!(body["tools"][0]["name"], "read"); + assert_eq!(body["tools"][0]["input_schema"]["type"], "object"); + } + + #[test] + fn assistant_tool_use_round_trips_into_body() { + let req = ModelRequest::new( + "claude-opus-4-8", + vec![ + Message::user("weather?"), + Message::assistant(vec![ContentBlock::ToolUse { + id: "toolu_1".into(), + name: "get_weather".into(), + input: json!({ "city": "SF" }), + }]), + Message::tool_result("toolu_1", "72F", false), + ], + 256, + ); + let body = build_body(&req); + assert_eq!(body["messages"][1]["content"][0]["type"], "tool_use"); + assert_eq!(body["messages"][1]["content"][0]["id"], "toolu_1"); + assert_eq!(body["messages"][2]["content"][0]["type"], "tool_result"); + assert_eq!(body["messages"][2]["content"][0]["tool_use_id"], "toolu_1"); + } + + // A recorded text + tool_use streamed response. + const FIXTURE: &str = r#" +event: message_start +data: {"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":24,"output_tokens":1}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Let me check."}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: content_block_start +data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_42","name":"get_weather","input":{}}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"SF\"}"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":1} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":31}} + +event: message_stop +data: {"type":"message_stop"} +"#; + + #[test] + fn decodes_text_then_tool_use_stream() { + let mut dec = Decoder::default(); + let events = decode_sse(&mut dec, FIXTURE).unwrap(); + assert_eq!( + events, + vec![ + StreamEvent::MessageStart, + StreamEvent::TextDelta { + text: "Let me check.".into() + }, + StreamEvent::ContentBlockStop, + StreamEvent::ToolUseStart { + id: "toolu_42".into(), + name: "get_weather".into() + }, + StreamEvent::InputJsonDelta { + partial_json: "{\"city\":".into() + }, + StreamEvent::InputJsonDelta { + partial_json: "\"SF\"}".into() + }, + StreamEvent::ContentBlockStop, + StreamEvent::Usage { + input_tokens: 24, + output_tokens: 31 + }, + StreamEvent::MessageStop { + stop_reason: StopReason::ToolUse + }, + ] + ); + } +} diff --git a/crates/agent-core/src/dialect/mod.rs b/crates/agent-core/src/dialect/mod.rs new file mode 100644 index 0000000..3f308b5 --- /dev/null +++ b/crates/agent-core/src/dialect/mod.rs @@ -0,0 +1,108 @@ +//! Wire dialects. +//! +//! The gateway relays bytes verbatim, so the harness owns the provider wire. Each [`Dialect`] knows +//! how to (a) build a streaming request body from a [`ModelRequest`] and (b) decode that provider's +//! SSE stream into the dialect-agnostic [`StreamEvent`] sequence the loop consumes. +//! +//! Selection is by model id: Claude speaks Anthropic (`/v1/messages`); everything else speaks the +//! OpenAI wire (`/v1/chat/completions`), which is the lingua franca for the gateway's other providers. + +use serde_json::Value; + +use crate::error::{Error, Result}; +use crate::message::StreamEvent; +use crate::transport::ModelRequest; + +pub mod anthropic; +pub mod openai; + +/// Which provider wire a model speaks. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Dialect { + Anthropic, + OpenAi, +} + +impl Dialect { + /// Pick the dialect for a model id. Claude → Anthropic; everything else → OpenAI wire. + pub fn for_model(model: &str) -> Self { + if model.starts_with("claude") || model.contains("anthropic") { + Dialect::Anthropic + } else { + Dialect::OpenAi + } + } + + /// The gateway path this dialect POSTs to. Bare `/v1*` so the gateway's default-provider routing + /// (OpenAI on `/v1`, Anthropic on `/v1/messages`) applies; a `bai_v1` key picks the real provider. + pub fn endpoint_path(&self) -> &'static str { + match self { + Dialect::Anthropic => "/v1/messages", + Dialect::OpenAi => "/v1/chat/completions", + } + } + + /// Build the JSON body for a streaming completion request. + pub fn build_body(&self, req: &ModelRequest) -> Value { + match self { + Dialect::Anthropic => anthropic::build_body(req), + Dialect::OpenAi => openai::build_body(req), + } + } + + /// A fresh streaming decoder for this dialect. + pub fn decoder(&self) -> Box { + match self { + Dialect::Anthropic => Box::::default(), + Dialect::OpenAi => Box::::default(), + } + } +} + +/// A stateful decoder turning a provider's SSE `data:` payloads into [`StreamEvent`]s. One per +/// request (it carries cross-event state: token counts, the open content block, the stop reason). +pub trait StreamDecoder: Send { + /// Feed one parsed SSE `data:` JSON object; return any events it produced. + fn push(&mut self, data: &Value) -> Vec; + + /// Called once at end-of-stream. Flushes any held terminal event (OpenAI defers `MessageStop` + /// until the stream closes so it can land after the trailing usage chunk). + fn finish(&mut self) -> Vec { + Vec::new() + } +} + +/// Decode a complete SSE body into events. Splits on `data:` lines, skips comments/`event:` lines +/// and the OpenAI `[DONE]` sentinel, and flushes the decoder at the end. The streaming HTTP client +/// reuses the same decoder incrementally; this is the buffered form used by tests. +pub fn decode_sse(decoder: &mut dyn StreamDecoder, raw: &str) -> Result> { + let mut out = Vec::new(); + for line in raw.lines() { + let line = line.trim_end_matches('\r'); + let Some(payload) = line.strip_prefix("data:") else { + continue; + }; + let payload = payload.trim(); + if payload.is_empty() || payload == "[DONE]" { + continue; + } + let v: Value = serde_json::from_str(payload) + .map_err(|e| Error::Transport(format!("malformed SSE json: {e}")))?; + out.extend(decoder.push(&v)); + } + out.extend(decoder.finish()); + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dialect_selection_by_model() { + assert_eq!(Dialect::for_model("claude-opus-4-8"), Dialect::Anthropic); + assert_eq!(Dialect::for_model("gpt-4o"), Dialect::OpenAi); + assert_eq!(Dialect::for_model("llama-3.1-70b"), Dialect::OpenAi); + assert_eq!(Dialect::for_model("anthropic/claude"), Dialect::Anthropic); + } +} diff --git a/crates/agent-core/src/dialect/openai.rs b/crates/agent-core/src/dialect/openai.rs new file mode 100644 index 0000000..cf88003 --- /dev/null +++ b/crates/agent-core/src/dialect/openai.rs @@ -0,0 +1,366 @@ +//! OpenAI Chat Completions wire (`/v1/chat/completions`). +//! +//! OpenAI's shape is flatter than the internal model, so both directions need real translation: +//! - **Request:** the system prompt becomes a `system` message; an assistant turn's `ToolUse` blocks +//! become `tool_calls` (arguments as a JSON *string*); a user turn's `ToolResult` blocks fan out +//! into separate `role:"tool"` messages. +//! - **Stream:** `tool_calls` arrive as index-keyed deltas (id+name once, then `arguments` fragments) +//! with no explicit block-stop events, so the decoder synthesizes `ContentBlockStop` and defers +//! `MessageStop` to end-of-stream (after the trailing `usage` chunk). + +use serde_json::{Map, Value, json}; + +use super::StreamDecoder; +use crate::message::{ContentBlock, Role, StopReason, StreamEvent}; +use crate::transport::ModelRequest; + +fn text_of(blocks: &[ContentBlock]) -> String { + blocks + .iter() + .filter_map(|b| match b { + ContentBlock::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +/// Build the streaming request body, translating the internal messages into OpenAI's flat shape. +pub fn build_body(req: &ModelRequest) -> Value { + let mut messages: Vec = Vec::new(); + if let Some(system) = &req.system { + messages.push(json!({ "role": "system", "content": system })); + } + + for m in &req.messages { + match m.role { + Role::System => { + messages.push(json!({ "role": "system", "content": text_of(&m.content) })) + } + Role::User => { + let text = text_of(&m.content); + if !text.is_empty() { + messages.push(json!({ "role": "user", "content": text })); + } + // Tool results fan out into individual `role:"tool"` messages. + for b in &m.content { + if let ContentBlock::ToolResult { + tool_use_id, + content, + .. + } = b + { + messages.push(json!({ "role": "tool", "tool_call_id": tool_use_id, "content": content })); + } + } + } + Role::Assistant => { + let text = text_of(&m.content); + let tool_calls: Vec = m + .content + .iter() + .filter_map(|b| match b { + ContentBlock::ToolUse { id, name, input } => Some(json!({ + "id": id, + "type": "function", + "function": { + "name": name, + "arguments": serde_json::to_string(input).unwrap_or_else(|_| "{}".into()), + }, + })), + _ => None, + }) + .collect(); + let mut msg = Map::new(); + msg.insert("role".into(), json!("assistant")); + msg.insert( + "content".into(), + if text.is_empty() { + Value::Null + } else { + json!(text) + }, + ); + if !tool_calls.is_empty() { + msg.insert("tool_calls".into(), Value::Array(tool_calls)); + } + messages.push(Value::Object(msg)); + } + } + } + + let mut map = Map::new(); + map.insert("model".into(), json!(req.model)); + map.insert("max_tokens".into(), json!(req.max_tokens)); + map.insert("stream".into(), json!(true)); + // Ask for a trailing usage chunk so token accounting works on the streaming path. + map.insert("stream_options".into(), json!({ "include_usage": true })); + map.insert("messages".into(), Value::Array(messages)); + if !req.tools.is_empty() { + let tools: Vec = req + .tools + .iter() + .map(|t| json!({ + "type": "function", + "function": { "name": t.name, "description": t.description, "parameters": t.input_schema }, + })) + .collect(); + map.insert("tools".into(), Value::Array(tools)); + } + Value::Object(map) +} + +fn map_finish_reason(s: Option<&str>) -> StopReason { + match s { + Some("stop") => StopReason::EndTurn, + Some("tool_calls") | Some("function_call") => StopReason::ToolUse, + Some("length") => StopReason::MaxTokens, + _ => StopReason::Other, + } +} + +#[derive(Clone, Copy, PartialEq)] +enum Open { + None, + Text, + Tool, +} + +/// Decodes OpenAI SSE. Synthesizes the block-stop boundaries OpenAI omits and holds `MessageStop` +/// until `finish()` so it lands after the trailing usage chunk. +pub struct Decoder { + started: bool, + open: Open, + stop_reason: StopReason, + input_tokens: u32, + output_tokens: u32, +} + +impl Default for Decoder { + fn default() -> Self { + Self { + started: false, + open: Open::None, + stop_reason: StopReason::EndTurn, + input_tokens: 0, + output_tokens: 0, + } + } +} + +impl Decoder { + fn close_open(&mut self, out: &mut Vec) { + if self.open != Open::None { + out.push(StreamEvent::ContentBlockStop); + self.open = Open::None; + } + } +} + +impl StreamDecoder for Decoder { + fn push(&mut self, data: &Value) -> Vec { + let mut out = Vec::new(); + if !self.started { + self.started = true; + out.push(StreamEvent::MessageStart); + } + + // Trailing usage-only chunk (choices empty). + if let Some(usage) = data.get("usage").filter(|u| !u.is_null()) { + self.input_tokens = usage + .get("prompt_tokens") + .and_then(Value::as_u64) + .unwrap_or(0) as u32; + self.output_tokens = usage + .get("completion_tokens") + .and_then(Value::as_u64) + .unwrap_or(0) as u32; + out.push(StreamEvent::Usage { + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + }); + } + + let Some(choice) = data.get("choices").and_then(|c| c.get(0)) else { + return out; + }; + let delta = choice.get("delta"); + + // Plain text content. + if let Some(text) = delta.and_then(|d| d.get("content")).and_then(Value::as_str) { + if !text.is_empty() { + if self.open == Open::None { + self.open = Open::Text; + } + out.push(StreamEvent::TextDelta { + text: text.to_string(), + }); + } + } + + // Tool-call deltas: id+name on first sight of an index, then `arguments` fragments. + if let Some(calls) = delta + .and_then(|d| d.get("tool_calls")) + .and_then(Value::as_array) + { + for tc in calls { + let func = tc.get("function"); + if let Some(id) = tc.get("id").and_then(Value::as_str) { + self.close_open(&mut out); + self.open = Open::Tool; + let name = func + .and_then(|f| f.get("name")) + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + out.push(StreamEvent::ToolUseStart { + id: id.to_string(), + name, + }); + } + if let Some(args) = func + .and_then(|f| f.get("arguments")) + .and_then(Value::as_str) + { + if !args.is_empty() { + out.push(StreamEvent::InputJsonDelta { + partial_json: args.to_string(), + }); + } + } + } + } + + // Finish reason closes the open block and records why we stopped (MessageStop is held). + if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) { + self.stop_reason = map_finish_reason(Some(reason)); + self.close_open(&mut out); + } + + out + } + + fn finish(&mut self) -> Vec { + if self.started { + vec![StreamEvent::MessageStop { + stop_reason: self.stop_reason, + }] + } else { + Vec::new() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dialect::decode_sse; + use crate::message::{Message, ToolDef}; + + #[test] + fn build_body_maps_system_tool_calls_and_results() { + let req = ModelRequest::new( + "gpt-4o", + vec![ + Message::user("weather?"), + Message::assistant(vec![ + ContentBlock::Text { + text: "checking".into(), + }, + ContentBlock::ToolUse { + id: "call_1".into(), + name: "get_weather".into(), + input: json!({ "city": "SF" }), + }, + ]), + Message::tool_result("call_1", "72F", false), + ], + 256, + ) + .with_system("be brief") + .with_tools(vec![ToolDef { + name: "get_weather".into(), + description: "weather".into(), + input_schema: json!({ "type": "object" }), + }]); + let body = build_body(&req); + + assert_eq!( + body["messages"][0], + json!({ "role": "system", "content": "be brief" }) + ); + assert_eq!( + body["messages"][1], + json!({ "role": "user", "content": "weather?" }) + ); + // assistant: content + a tool_call whose arguments are a JSON *string*. + assert_eq!(body["messages"][2]["role"], "assistant"); + assert_eq!(body["messages"][2]["tool_calls"][0]["id"], "call_1"); + assert_eq!( + body["messages"][2]["tool_calls"][0]["function"]["arguments"], + "{\"city\":\"SF\"}" + ); + // tool result fanned out into a role:"tool" message. + assert_eq!( + body["messages"][3], + json!({ "role": "tool", "tool_call_id": "call_1", "content": "72F" }) + ); + // tool schema nested under function.parameters. + assert_eq!(body["tools"][0]["function"]["name"], "get_weather"); + assert_eq!(body["stream_options"]["include_usage"], true); + } + + // A recorded text + tool_call streamed response (trailing usage chunk + [DONE]). + const FIXTURE: &str = r#" +data: {"choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} + +data: {"choices":[{"index":0,"delta":{"content":"Let me check."},"finish_reason":null}]} + +data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_42","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]} + +data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":null}]} + +data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"SF\"}"}}]},"finish_reason":null}]} + +data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]} + +data: {"choices":[],"usage":{"prompt_tokens":24,"completion_tokens":31,"total_tokens":55}} + +data: [DONE] +"#; + + #[test] + fn decodes_text_then_tool_call_stream() { + let mut dec = Decoder::default(); + let events = decode_sse(&mut dec, FIXTURE).unwrap(); + assert_eq!( + events, + vec![ + StreamEvent::MessageStart, + StreamEvent::TextDelta { + text: "Let me check.".into() + }, + // Text block closes when the tool call begins — same shape the Anthropic decoder + // produces, so the loop assembles both dialects identically. + StreamEvent::ContentBlockStop, + StreamEvent::ToolUseStart { + id: "call_42".into(), + name: "get_weather".into() + }, + StreamEvent::InputJsonDelta { + partial_json: "{\"city\":".into() + }, + StreamEvent::InputJsonDelta { + partial_json: "\"SF\"}".into() + }, + StreamEvent::ContentBlockStop, + StreamEvent::Usage { + input_tokens: 24, + output_tokens: 31 + }, + StreamEvent::MessageStop { + stop_reason: StopReason::ToolUse + }, + ] + ); + } +} diff --git a/crates/agent-core/src/lib.rs b/crates/agent-core/src/lib.rs index 1d01a88..27694e6 100644 --- a/crates/agent-core/src/lib.rs +++ b/crates/agent-core/src/lib.rs @@ -18,6 +18,7 @@ // job is to assert a precondition holds, so `.unwrap()` *is* the assertion — allow it under `test`. #![cfg_attr(test, allow(clippy::unwrap_used, clippy::expect_used, clippy::panic))] +pub mod dialect; pub mod error; pub mod message; pub mod session; diff --git a/crates/agent-core/src/message.rs b/crates/agent-core/src/message.rs index 804df22..7c5b808 100644 --- a/crates/agent-core/src/message.rs +++ b/crates/agent-core/src/message.rs @@ -101,10 +101,11 @@ pub struct ToolDef { } /// Why the model ended a turn. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum StopReason { /// Natural end of the assistant's turn. + #[default] EndTurn, /// The model wants to call one or more tools; the loop should dispatch and continue. ToolUse,