M1: Anthropic codec and provider
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.
This commit is contained in:
Generated
+1013
-12
File diff suppressed because it is too large
Load Diff
+1
-2
@@ -59,8 +59,7 @@ dirs = "5"
|
||||
shell-words = "1"
|
||||
libc = "0.2"
|
||||
tempfile = "3"
|
||||
insta = "1"
|
||||
wiremock = "0.6"
|
||||
bytes = "1"
|
||||
|
||||
[workspace.lints.rust]
|
||||
unsafe_code = "deny"
|
||||
|
||||
@@ -574,18 +574,18 @@ pub async fn process_step(
|
||||
let mut result = StepResult::Stop;
|
||||
|
||||
loop {
|
||||
let cancelled = ctx.cancel.is_cancelled();
|
||||
if cancelled {
|
||||
run.mark_errored(&ProviderError::Cancelled).await;
|
||||
return Ok(StepOutcome {
|
||||
result: StepResult::Stop,
|
||||
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
||||
usage,
|
||||
aborted: true,
|
||||
});
|
||||
}
|
||||
|
||||
let item = stream.next().await;
|
||||
let item = tokio::select! {
|
||||
item = stream.next() => item,
|
||||
_ = ctx.cancel.cancelled() => {
|
||||
run.mark_errored(&ProviderError::Cancelled).await;
|
||||
return Ok(StepOutcome {
|
||||
result: StepResult::Stop,
|
||||
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
||||
usage,
|
||||
aborted: true,
|
||||
});
|
||||
}
|
||||
};
|
||||
let Some(item) = item else { break };
|
||||
|
||||
let event = match item {
|
||||
|
||||
@@ -6,6 +6,21 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
async-stream = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
eventsource-stream = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util", "macros"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::llm::{LlmEventStream, LlmRequest, Provider, ProviderError};
|
||||
use harness_core::types::ModelInfo;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::codec::anthropic::{build_request, decode, initiator_header};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
const ANTHROPIC_BETA: &str =
|
||||
"interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14";
|
||||
|
||||
pub struct AnthropicProvider {
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
impl AnthropicProvider {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
api_key: api_key.into(),
|
||||
base_url: DEFAULT_BASE_URL.to_string(),
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn with_base_url(api_key: impl Into<String>, base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
api_key: api_key.into(),
|
||||
base_url: base_url.into(),
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_error(
|
||||
status: reqwest::StatusCode,
|
||||
body: String,
|
||||
retry_after: Option<Duration>,
|
||||
) -> ProviderError {
|
||||
match status.as_u16() {
|
||||
401 | 403 => ProviderError::Auth(body),
|
||||
429 => ProviderError::RateLimited { retry_after },
|
||||
529 => ProviderError::Overloaded,
|
||||
s if (500..600).contains(&s) => ProviderError::Overloaded,
|
||||
s => ProviderError::Http { status: s, body },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for AnthropicProvider {
|
||||
fn id(&self) -> &str {
|
||||
"anthropic"
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
// Full models.dev metadata integration lands in M3; callers fall back to
|
||||
// config-specified model strings until then.
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
req: LlmRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError> {
|
||||
let body = build_request(&req);
|
||||
let url = format!("{}/v1/messages", self.base_url);
|
||||
let initiator = initiator_header(req.initiator);
|
||||
|
||||
let send = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("anthropic-beta", ANTHROPIC_BETA)
|
||||
.header("x-initiator", initiator)
|
||||
.json(&body)
|
||||
.send();
|
||||
|
||||
let response = tokio::select! {
|
||||
result = send => result.map_err(|e| ProviderError::Network(e.to_string()))?,
|
||||
_ = cancel.cancelled() => return Err(ProviderError::Cancelled),
|
||||
};
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let retry_after = response
|
||||
.headers()
|
||||
.get("retry-after")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.map(Duration::from_secs);
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(Self::classify_error(status, body, retry_after));
|
||||
}
|
||||
|
||||
Ok(decode(response.bytes_stream()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn classifies_status_codes() {
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::TOO_MANY_REQUESTS,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::RateLimited { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::UNAUTHORIZED,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Auth(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::SERVICE_UNAVAILABLE,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Overloaded
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::BAD_REQUEST,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Http { status: 400, .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn id_is_anthropic() {
|
||||
let provider = AnthropicProvider::new("test-key");
|
||||
assert_eq!(provider.id(), "anthropic");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn constructs_with_custom_base_url() {
|
||||
let provider = AnthropicProvider::with_base_url("test-key", "http://localhost:1234");
|
||||
assert_eq!(provider.base_url, "http://localhost:1234");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,499 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod anthropic;
|
||||
@@ -1 +1,6 @@
|
||||
// Provider implementations (Copilot, OpenAI, Anthropic) land here in M1/M3.
|
||||
pub mod anthropic;
|
||||
pub mod codec;
|
||||
pub mod registry;
|
||||
|
||||
pub use anthropic::AnthropicProvider;
|
||||
pub use registry::ProviderRegistry;
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use harness_core::llm::Provider;
|
||||
|
||||
#[derive(Default, Clone)]
|
||||
pub struct ProviderRegistry {
|
||||
providers: HashMap<String, Arc<dyn Provider>>,
|
||||
}
|
||||
|
||||
impl ProviderRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn register(&mut self, provider: Arc<dyn Provider>) {
|
||||
self.providers.insert(provider.id().to_string(), provider);
|
||||
}
|
||||
|
||||
pub fn get(&self, id: &str) -> Option<Arc<dyn Provider>> {
|
||||
self.providers.get(id).cloned()
|
||||
}
|
||||
|
||||
/// Resolves a `"provider/model"` string to its provider and bare model id.
|
||||
pub fn resolve(&self, model_ref: &str) -> Option<(Arc<dyn Provider>, String)> {
|
||||
let (provider_id, model_id) = model_ref.split_once('/')?;
|
||||
self.get(provider_id).map(|p| (p, model_id.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::anthropic::AnthropicProvider;
|
||||
|
||||
#[test]
|
||||
fn resolves_provider_slash_model() {
|
||||
let mut registry = ProviderRegistry::new();
|
||||
registry.register(Arc::new(AnthropicProvider::new("key")));
|
||||
let (provider, model) = registry.resolve("anthropic/claude-sonnet-4-5").unwrap();
|
||||
assert_eq!(provider.id(), "anthropic");
|
||||
assert_eq!(model, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_provider_resolves_to_none() {
|
||||
let registry = ProviderRegistry::new();
|
||||
assert!(registry.resolve("openai/gpt-4o").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_model_ref_resolves_to_none() {
|
||||
let mut registry = ProviderRegistry::new();
|
||||
registry.register(Arc::new(AnthropicProvider::new("key")));
|
||||
assert!(registry.resolve("no-slash-here").is_none());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user