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"
|
shell-words = "1"
|
||||||
libc = "0.2"
|
libc = "0.2"
|
||||||
tempfile = "3"
|
tempfile = "3"
|
||||||
insta = "1"
|
bytes = "1"
|
||||||
wiremock = "0.6"
|
|
||||||
|
|
||||||
[workspace.lints.rust]
|
[workspace.lints.rust]
|
||||||
unsafe_code = "deny"
|
unsafe_code = "deny"
|
||||||
|
|||||||
@@ -574,8 +574,9 @@ pub async fn process_step(
|
|||||||
let mut result = StepResult::Stop;
|
let mut result = StepResult::Stop;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
let cancelled = ctx.cancel.is_cancelled();
|
let item = tokio::select! {
|
||||||
if cancelled {
|
item = stream.next() => item,
|
||||||
|
_ = ctx.cancel.cancelled() => {
|
||||||
run.mark_errored(&ProviderError::Cancelled).await;
|
run.mark_errored(&ProviderError::Cancelled).await;
|
||||||
return Ok(StepOutcome {
|
return Ok(StepOutcome {
|
||||||
result: StepResult::Stop,
|
result: StepResult::Stop,
|
||||||
@@ -584,8 +585,7 @@ pub async fn process_step(
|
|||||||
aborted: true,
|
aborted: true,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
};
|
||||||
let item = stream.next().await;
|
|
||||||
let Some(item) = item else { break };
|
let Some(item) = item else { break };
|
||||||
|
|
||||||
let event = match item {
|
let event = match item {
|
||||||
|
|||||||
@@ -6,6 +6,21 @@ license.workspace = true
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
harness-core = { workspace = true }
|
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]
|
[lints]
|
||||||
workspace = true
|
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