Adds the Anthropic SSE codec (codec/anthropic.rs) translating Anthropic's event stream into LlmEvents, the AnthropicProvider, and registry wiring so the engine loop can run against the real API instead of just MockProvider.
500 lines
20 KiB
Rust
500 lines
20 KiB
Rust
//! Request builder + SSE decoder for Anthropic's `/v1/messages` streaming API.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use async_stream::try_stream;
|
|
use eventsource_stream::Eventsource;
|
|
use futures::Stream;
|
|
use harness_core::llm::{
|
|
FinishReason, Initiator, LlmEvent, LlmEventStream, LlmRequest, ProviderError, Role, WireContent,
|
|
};
|
|
use harness_core::types::TokenUsage;
|
|
use serde_json::{json, Value};
|
|
|
|
const DEFAULT_MAX_TOKENS: u32 = 8192;
|
|
/// opencode's `applyCaching`: cache breakpoints on the first 2 system blocks + last 2
|
|
/// non-system messages.
|
|
const CACHE_BREAKPOINTS: usize = 2;
|
|
|
|
fn cache_control() -> Value {
|
|
json!({"type": "ephemeral"})
|
|
}
|
|
|
|
fn build_system(system: &[String]) -> Vec<Value> {
|
|
system
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, block)| {
|
|
let mut b = json!({"type": "text", "text": block});
|
|
if i < CACHE_BREAKPOINTS {
|
|
b["cache_control"] = cache_control();
|
|
}
|
|
b
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn wire_role_to_anthropic(role: Role) -> &'static str {
|
|
match role {
|
|
Role::User | Role::Tool => "user",
|
|
Role::Assistant => "assistant",
|
|
Role::System => "user", // system blocks are carried separately in `system`, not here
|
|
}
|
|
}
|
|
|
|
fn build_messages(messages: &[harness_core::llm::WireMessage]) -> Vec<Value> {
|
|
let mut built: Vec<Value> = messages
|
|
.iter()
|
|
.map(|m| {
|
|
let content: Vec<Value> = m
|
|
.content
|
|
.iter()
|
|
.map(|c| match c {
|
|
WireContent::Text { text } => json!({"type": "text", "text": text}),
|
|
WireContent::ToolCall { call_id, name, input } => {
|
|
json!({"type": "tool_use", "id": call_id, "name": name, "input": input})
|
|
}
|
|
WireContent::ToolResult { call_id, output, is_error } => {
|
|
json!({"type": "tool_result", "tool_use_id": call_id, "content": output, "is_error": is_error})
|
|
}
|
|
WireContent::Image { mime_type, data } => {
|
|
json!({"type": "image", "source": {"type": "base64", "media_type": mime_type, "data": data}})
|
|
}
|
|
})
|
|
.collect();
|
|
json!({"role": wire_role_to_anthropic(m.role), "content": content})
|
|
})
|
|
.collect();
|
|
|
|
let n = built.len();
|
|
for msg in built.iter_mut().skip(n.saturating_sub(CACHE_BREAKPOINTS)) {
|
|
if let Some(content) = msg["content"].as_array_mut() {
|
|
if let Some(last) = content.last_mut() {
|
|
last["cache_control"] = cache_control();
|
|
}
|
|
}
|
|
}
|
|
built
|
|
}
|
|
|
|
pub fn build_request(req: &LlmRequest) -> Value {
|
|
let mut body = json!({
|
|
"model": req.model,
|
|
"system": build_system(&req.system),
|
|
"messages": build_messages(&req.messages),
|
|
"max_tokens": req.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS),
|
|
"stream": true,
|
|
});
|
|
|
|
if !req.tools.is_empty() {
|
|
body["tools"] = Value::Array(
|
|
req.tools
|
|
.iter()
|
|
.map(|t| json!({"name": t.name, "description": t.description, "input_schema": t.parameters}))
|
|
.collect(),
|
|
);
|
|
}
|
|
if let Some(temp) = req.temperature {
|
|
body["temperature"] = json!(temp);
|
|
}
|
|
if let Some(reasoning) = &req.reasoning {
|
|
if let Some(budget) = reasoning.budget_tokens {
|
|
body["thinking"] = json!({"type": "enabled", "budget_tokens": budget});
|
|
}
|
|
}
|
|
body
|
|
}
|
|
|
|
pub fn initiator_header(initiator: Initiator) -> &'static str {
|
|
match initiator {
|
|
Initiator::User => "user",
|
|
Initiator::Agent => "agent",
|
|
}
|
|
}
|
|
|
|
fn map_stop_reason(reason: Option<&str>) -> FinishReason {
|
|
match reason {
|
|
Some("end_turn") | Some("stop_sequence") => FinishReason::Stop,
|
|
Some("tool_use") => FinishReason::ToolCalls,
|
|
Some("max_tokens") => FinishReason::Length,
|
|
Some(other) => FinishReason::Unknown(other.to_string()),
|
|
None => FinishReason::Unknown("none".to_string()),
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct BlockState {
|
|
kind: BlockKind,
|
|
call_id: String,
|
|
name: String,
|
|
partial_json: String,
|
|
signature: Option<String>,
|
|
}
|
|
|
|
#[derive(Default, PartialEq, Eq)]
|
|
enum BlockKind {
|
|
#[default]
|
|
Text,
|
|
Thinking,
|
|
ToolUse,
|
|
}
|
|
|
|
/// Decodes a raw SSE byte stream into our normalized `LlmEvent` stream. Errors from the
|
|
/// underlying HTTP stream and any Anthropic `error` event both surface as `Err`.
|
|
pub fn decode<S, E>(byte_stream: S) -> LlmEventStream
|
|
where
|
|
S: Stream<Item = Result<bytes::Bytes, E>> + Send + 'static,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
let events = byte_stream.eventsource();
|
|
Box::pin(try_stream! {
|
|
futures::pin_mut!(events);
|
|
let mut blocks: HashMap<u64, BlockState> = HashMap::new();
|
|
let mut usage = TokenUsage::default();
|
|
|
|
while let Some(item) = futures::StreamExt::next(&mut events).await {
|
|
let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
|
if event.data.is_empty() {
|
|
continue;
|
|
}
|
|
let value: Value = serde_json::from_str(&event.data)
|
|
.map_err(|e| ProviderError::Decode(format!("{e}: {}", event.data)))?;
|
|
let kind = value["type"].as_str().unwrap_or_default();
|
|
|
|
match kind {
|
|
"message_start" => {
|
|
let u = &value["message"]["usage"];
|
|
usage.input = u["input_tokens"].as_u64().unwrap_or(0);
|
|
usage.cache_write = u["cache_creation_input_tokens"].as_u64().unwrap_or(0);
|
|
usage.cache_read = u["cache_read_input_tokens"].as_u64().unwrap_or(0);
|
|
}
|
|
"content_block_start" => {
|
|
let index = value["index"].as_u64().unwrap_or(0);
|
|
let block = &value["content_block"];
|
|
match block["type"].as_str().unwrap_or_default() {
|
|
"text" => {
|
|
blocks.insert(index, BlockState { kind: BlockKind::Text, ..Default::default() });
|
|
yield LlmEvent::TextStart { id: index.to_string() };
|
|
}
|
|
"thinking" => {
|
|
blocks.insert(index, BlockState { kind: BlockKind::Thinking, ..Default::default() });
|
|
yield LlmEvent::ReasoningStart { id: index.to_string() };
|
|
}
|
|
"tool_use" => {
|
|
let call_id = block["id"].as_str().unwrap_or_default().to_string();
|
|
let name = block["name"].as_str().unwrap_or_default().to_string();
|
|
blocks.insert(index, BlockState {
|
|
kind: BlockKind::ToolUse,
|
|
call_id: call_id.clone(),
|
|
name: name.clone(),
|
|
..Default::default()
|
|
});
|
|
yield LlmEvent::ToolInputStart { call_id, name };
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
"content_block_delta" => {
|
|
let index = value["index"].as_u64().unwrap_or(0);
|
|
let delta = &value["delta"];
|
|
match delta["type"].as_str().unwrap_or_default() {
|
|
"text_delta" => {
|
|
let text = delta["text"].as_str().unwrap_or_default().to_string();
|
|
yield LlmEvent::TextDelta { id: index.to_string(), text };
|
|
}
|
|
"thinking_delta" => {
|
|
let text = delta["thinking"].as_str().unwrap_or_default().to_string();
|
|
yield LlmEvent::ReasoningDelta { id: index.to_string(), text };
|
|
}
|
|
"signature_delta" => {
|
|
if let Some(state) = blocks.get_mut(&index) {
|
|
state.signature = Some(delta["signature"].as_str().unwrap_or_default().to_string());
|
|
}
|
|
}
|
|
"input_json_delta" => {
|
|
let partial = delta["partial_json"].as_str().unwrap_or_default();
|
|
if let Some(state) = blocks.get_mut(&index) {
|
|
state.partial_json.push_str(partial);
|
|
yield LlmEvent::ToolInputDelta { call_id: state.call_id.clone(), json: partial.to_string() };
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
"content_block_stop" => {
|
|
let index = value["index"].as_u64().unwrap_or(0);
|
|
if let Some(state) = blocks.remove(&index) {
|
|
match state.kind {
|
|
BlockKind::Text => yield LlmEvent::TextEnd { id: index.to_string() },
|
|
BlockKind::Thinking => {
|
|
yield LlmEvent::ReasoningEnd { id: index.to_string(), signature: state.signature };
|
|
}
|
|
BlockKind::ToolUse => {
|
|
let input: Value = if state.partial_json.trim().is_empty() {
|
|
json!({})
|
|
} else {
|
|
serde_json::from_str(&state.partial_json).unwrap_or(Value::Null)
|
|
};
|
|
yield LlmEvent::ToolCall { call_id: state.call_id, name: state.name, input };
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"message_delta" => {
|
|
if let Some(out) = value["usage"]["output_tokens"].as_u64() {
|
|
usage.output = out;
|
|
}
|
|
let stop_reason = value["delta"]["stop_reason"].as_str();
|
|
yield LlmEvent::Finish { reason: map_stop_reason(stop_reason), usage };
|
|
}
|
|
"error" => {
|
|
let message = value["error"]["message"].as_str().unwrap_or("unknown error").to_string();
|
|
let err_type = value["error"]["type"].as_str().unwrap_or("");
|
|
Err(match err_type {
|
|
"overloaded_error" => ProviderError::Overloaded,
|
|
"rate_limit_error" => ProviderError::RateLimited { retry_after: None },
|
|
"authentication_error" | "permission_error" => ProviderError::Auth(message),
|
|
_ => ProviderError::Http { status: 0, body: message },
|
|
})?;
|
|
}
|
|
_ => {} // ping, message_stop: nothing to emit
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use futures::StreamExt;
|
|
use harness_core::llm::{ReasoningOpts, ToolSchema, WireMessage};
|
|
|
|
fn sse_stream(raw: &'static str) -> LlmEventStream {
|
|
let chunks: Vec<Result<bytes::Bytes, std::io::Error>> =
|
|
vec![Ok(bytes::Bytes::from_static(raw.as_bytes()))];
|
|
decode(futures::stream::iter(chunks))
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn decodes_text_only_response() {
|
|
let raw = concat!(
|
|
"event: message_start\n",
|
|
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n",
|
|
"event: content_block_start\n",
|
|
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
|
|
"event: content_block_delta\n",
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n",
|
|
"event: content_block_stop\n",
|
|
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
|
"event: message_delta\n",
|
|
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n",
|
|
"event: message_stop\n",
|
|
"data: {\"type\":\"message_stop\"}\n\n",
|
|
);
|
|
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
|
assert_eq!(
|
|
events,
|
|
vec![
|
|
LlmEvent::TextStart { id: "0".into() },
|
|
LlmEvent::TextDelta {
|
|
id: "0".into(),
|
|
text: "Hello".into()
|
|
},
|
|
LlmEvent::TextEnd { id: "0".into() },
|
|
LlmEvent::Finish {
|
|
reason: FinishReason::Stop,
|
|
usage: TokenUsage {
|
|
input: 10,
|
|
output: 5,
|
|
..Default::default()
|
|
},
|
|
},
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn decodes_tool_call_with_streamed_json_input() {
|
|
let raw = concat!(
|
|
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
|
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"read\"}}\n\n",
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file\\\"\"}}\n\n",
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\":\\\"a.txt\\\"}\"}}\n\n",
|
|
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
|
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":8}}\n\n",
|
|
);
|
|
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
|
assert_eq!(
|
|
events,
|
|
vec![
|
|
LlmEvent::ToolInputStart {
|
|
call_id: "call_1".into(),
|
|
name: "read".into()
|
|
},
|
|
LlmEvent::ToolInputDelta {
|
|
call_id: "call_1".into(),
|
|
json: "{\"file\"".into()
|
|
},
|
|
LlmEvent::ToolInputDelta {
|
|
call_id: "call_1".into(),
|
|
json: ":\"a.txt\"}".into()
|
|
},
|
|
LlmEvent::ToolCall {
|
|
call_id: "call_1".into(),
|
|
name: "read".into(),
|
|
input: json!({"file": "a.txt"}),
|
|
},
|
|
LlmEvent::Finish {
|
|
reason: FinishReason::ToolCalls,
|
|
usage: TokenUsage {
|
|
input: 1,
|
|
output: 8,
|
|
..Default::default()
|
|
},
|
|
},
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn decodes_thinking_block_with_signature() {
|
|
let raw = concat!(
|
|
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
|
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\"}}\n\n",
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"pondering\"}}\n\n",
|
|
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig123\"}}\n\n",
|
|
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
|
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n",
|
|
);
|
|
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
|
assert_eq!(
|
|
events,
|
|
vec![
|
|
LlmEvent::ReasoningStart { id: "0".into() },
|
|
LlmEvent::ReasoningDelta {
|
|
id: "0".into(),
|
|
text: "pondering".into()
|
|
},
|
|
LlmEvent::ReasoningEnd {
|
|
id: "0".into(),
|
|
signature: Some("sig123".into())
|
|
},
|
|
LlmEvent::Finish {
|
|
reason: FinishReason::Stop,
|
|
usage: TokenUsage {
|
|
input: 1,
|
|
output: 2,
|
|
..Default::default()
|
|
},
|
|
},
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn mid_stream_error_event_surfaces_as_err() {
|
|
let raw = concat!(
|
|
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
|
"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"overloaded\"}}\n\n",
|
|
);
|
|
let events: Vec<Result<LlmEvent, ProviderError>> = sse_stream(raw).collect().await;
|
|
assert!(matches!(
|
|
events.last(),
|
|
Some(Err(ProviderError::Overloaded))
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn build_request_applies_cache_control_to_first_two_system_blocks() {
|
|
let req = LlmRequest {
|
|
model: "claude-sonnet".into(),
|
|
system: vec!["env".into(), "agent".into(), "instructions".into()],
|
|
messages: vec![],
|
|
tools: vec![],
|
|
temperature: None,
|
|
max_tokens: None,
|
|
reasoning: None,
|
|
initiator: Initiator::User,
|
|
};
|
|
let body = build_request(&req);
|
|
let system = body["system"].as_array().unwrap();
|
|
assert_eq!(system.len(), 3);
|
|
assert!(system[0]["cache_control"].is_object());
|
|
assert!(system[1]["cache_control"].is_object());
|
|
assert!(system[2].get("cache_control").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn build_request_applies_cache_control_to_last_two_messages() {
|
|
let messages = vec![
|
|
WireMessage {
|
|
role: Role::User,
|
|
content: vec![WireContent::Text { text: "1".into() }],
|
|
},
|
|
WireMessage {
|
|
role: Role::Assistant,
|
|
content: vec![WireContent::Text { text: "2".into() }],
|
|
},
|
|
WireMessage {
|
|
role: Role::User,
|
|
content: vec![WireContent::Text { text: "3".into() }],
|
|
},
|
|
];
|
|
let req = LlmRequest {
|
|
model: "claude-sonnet".into(),
|
|
system: vec![],
|
|
messages,
|
|
tools: vec![],
|
|
temperature: None,
|
|
max_tokens: None,
|
|
reasoning: None,
|
|
initiator: Initiator::User,
|
|
};
|
|
let body = build_request(&req);
|
|
let msgs = body["messages"].as_array().unwrap();
|
|
assert!(msgs[0]["content"][0].get("cache_control").is_none());
|
|
assert!(msgs[1]["content"][0]["cache_control"].is_object());
|
|
assert!(msgs[2]["content"][0]["cache_control"].is_object());
|
|
}
|
|
|
|
#[test]
|
|
fn build_request_maps_tool_schema_to_input_schema_key() {
|
|
let req = LlmRequest {
|
|
model: "claude-sonnet".into(),
|
|
system: vec![],
|
|
messages: vec![],
|
|
tools: vec![ToolSchema {
|
|
name: "read".into(),
|
|
description: "reads a file".into(),
|
|
parameters: json!({"type": "object"}),
|
|
}],
|
|
temperature: None,
|
|
max_tokens: None,
|
|
reasoning: None,
|
|
initiator: Initiator::User,
|
|
};
|
|
let body = build_request(&req);
|
|
assert_eq!(body["tools"][0]["name"], "read");
|
|
assert_eq!(body["tools"][0]["input_schema"], json!({"type": "object"}));
|
|
}
|
|
|
|
#[test]
|
|
fn build_request_includes_thinking_budget_when_reasoning_set() {
|
|
let req = LlmRequest {
|
|
model: "claude-sonnet".into(),
|
|
system: vec![],
|
|
messages: vec![],
|
|
tools: vec![],
|
|
temperature: None,
|
|
max_tokens: None,
|
|
reasoning: Some(ReasoningOpts {
|
|
effort: None,
|
|
budget_tokens: Some(2048),
|
|
}),
|
|
initiator: Initiator::User,
|
|
};
|
|
let body = build_request(&req);
|
|
assert_eq!(body["thinking"]["budget_tokens"], 2048);
|
|
}
|
|
}
|