882 lines
31 KiB
Rust
882 lines
31 KiB
Rust
use async_trait::async_trait;
|
|
use futures_util::stream;
|
|
use reqwest::Client;
|
|
use serde::Serialize;
|
|
use serde_json::Value;
|
|
use std::collections::{BTreeMap, HashMap, VecDeque};
|
|
use std::time::Duration;
|
|
use thiserror::Error;
|
|
|
|
use super::stream::SseFramer;
|
|
use super::traits::Usage;
|
|
use super::{
|
|
ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk,
|
|
ProviderStream, Tool,
|
|
};
|
|
use crate::bus::{ProviderReasoningState, message::ContentBlock};
|
|
use crate::storage::Storage;
|
|
use std::sync::Arc;
|
|
|
|
const LLM_REQUEST_TIMEOUT_SECS: u64 = 300;
|
|
|
|
#[derive(Serialize)]
|
|
struct CacheControl {
|
|
#[serde(rename = "type")]
|
|
cache_type: String,
|
|
}
|
|
|
|
impl CacheControl {
|
|
fn ephemeral() -> Self {
|
|
Self {
|
|
cache_type: "ephemeral".to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn convert_content_blocks(blocks: &[ContentBlock], cacheable: bool) -> Vec<serde_json::Value> {
|
|
blocks
|
|
.iter()
|
|
.map(|b| match b {
|
|
ContentBlock::Text { text } => {
|
|
if cacheable {
|
|
serde_json::json!({
|
|
"type": "text",
|
|
"text": text,
|
|
"cache_control": CacheControl::ephemeral(),
|
|
})
|
|
} else {
|
|
serde_json::json!({ "type": "text", "text": text })
|
|
}
|
|
}
|
|
ContentBlock::ImageUrl { image_url } => convert_image_url_to_anthropic(&image_url.url),
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn convert_image_url_to_anthropic(url: &str) -> serde_json::Value {
|
|
// data:image/png;base64,... -> Anthropic image block
|
|
if let Some(caps) = regex::Regex::new(r"data:(image/\w+);base64,(.+)")
|
|
.ok()
|
|
.and_then(|re| re.captures(url))
|
|
{
|
|
let media_type = caps.get(1).map(|m| m.as_str()).unwrap_or("image/png");
|
|
let data = caps.get(2).map(|d| d.as_str()).unwrap_or("");
|
|
return serde_json::json!({
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": media_type,
|
|
"data": data
|
|
}
|
|
});
|
|
}
|
|
// Regular URL -> Anthropic image block with url source
|
|
serde_json::json!({
|
|
"type": "image",
|
|
"source": {
|
|
"type": "url",
|
|
"url": url
|
|
}
|
|
})
|
|
}
|
|
|
|
pub struct AnthropicProvider {
|
|
client: Client,
|
|
name: String,
|
|
api_key: String,
|
|
base_url: String,
|
|
extra_headers: HashMap<String, String>,
|
|
model_id: String,
|
|
temperature: Option<f32>,
|
|
max_tokens: Option<u32>,
|
|
model_extra: HashMap<String, serde_json::Value>,
|
|
storage: Option<Arc<Storage>>,
|
|
}
|
|
|
|
impl AnthropicProvider {
|
|
// Keep this constructor aligned with OpenAIProvider and LLMProviderConfig.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn new(
|
|
name: String,
|
|
api_key: String,
|
|
base_url: String,
|
|
extra_headers: HashMap<String, String>,
|
|
model_id: String,
|
|
temperature: Option<f32>,
|
|
max_tokens: Option<u32>,
|
|
model_extra: HashMap<String, serde_json::Value>,
|
|
) -> Self {
|
|
Self {
|
|
client: Client::builder()
|
|
.timeout(Duration::from_secs(LLM_REQUEST_TIMEOUT_SECS))
|
|
.build()
|
|
.unwrap_or_else(|_| Client::new()),
|
|
name,
|
|
api_key,
|
|
base_url,
|
|
extra_headers,
|
|
model_id,
|
|
temperature,
|
|
max_tokens,
|
|
model_extra,
|
|
storage: None,
|
|
}
|
|
}
|
|
|
|
pub fn set_storage(&mut self, storage: Arc<Storage>) {
|
|
self.storage = Some(storage);
|
|
}
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct AnthropicRequest {
|
|
model: String,
|
|
messages: Vec<AnthropicMessage>,
|
|
max_tokens: u32,
|
|
temperature: Option<f32>,
|
|
stream: bool,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
tools: Option<Vec<AnthropicTool>>,
|
|
#[serde(flatten)]
|
|
extra: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct AnthropicMessage {
|
|
role: String,
|
|
content: Vec<serde_json::Value>,
|
|
}
|
|
|
|
fn convert_messages(messages: &[Message]) -> Vec<AnthropicMessage> {
|
|
let mut converted = Vec::with_capacity(messages.len());
|
|
let mut index = 0;
|
|
|
|
while index < messages.len() {
|
|
let message = &messages[index];
|
|
|
|
// Anthropic requires all tool results for one assistant tool-use turn
|
|
// to be carried in a single `role: user` content array. Steering is
|
|
// represented as a normal user message in PicoBot history, so merge
|
|
// any immediately-following user messages into that same array at
|
|
// the provider boundary. Durable messages remain independent.
|
|
if message.role == "tool" && message.tool_call_id.is_some() {
|
|
let mut content = Vec::new();
|
|
while index < messages.len()
|
|
&& messages[index].role == "tool"
|
|
&& messages[index].tool_call_id.is_some()
|
|
{
|
|
let tool = &messages[index];
|
|
let tool_call_id = tool
|
|
.tool_call_id
|
|
.as_deref()
|
|
.expect("tool_call_id checked above");
|
|
content.push(serde_json::json!({
|
|
"type": "tool_result",
|
|
"tool_use_id": tool_call_id,
|
|
"content": convert_content_blocks(&tool.content, false),
|
|
}));
|
|
index += 1;
|
|
}
|
|
|
|
// One turn may receive more than one steering message before the
|
|
// next model request. Keep their order while emitting one native
|
|
// Anthropic user message alongside the tool_result blocks.
|
|
while index < messages.len() && messages[index].role == "user" {
|
|
let steering = &messages[index];
|
|
if let Some(native) = native_anthropic_content(steering) {
|
|
content.extend(native);
|
|
} else {
|
|
content.extend(convert_content_blocks(&steering.content, false));
|
|
}
|
|
index += 1;
|
|
}
|
|
|
|
converted.push(AnthropicMessage {
|
|
role: "user".to_string(),
|
|
content,
|
|
});
|
|
continue;
|
|
}
|
|
|
|
let role = message.role.clone();
|
|
let content = if let Some(native) = native_anthropic_content(message) {
|
|
native
|
|
} else {
|
|
let mut blocks = convert_content_blocks(&message.content, message.role == "system");
|
|
if let Some(tool_calls) = message
|
|
.tool_calls
|
|
.as_ref()
|
|
.filter(|calls| !calls.is_empty())
|
|
{
|
|
for tool_call in tool_calls {
|
|
blocks.push(serde_json::json!({
|
|
"type": "tool_use",
|
|
"id": tool_call.id,
|
|
"name": tool_call.name,
|
|
"input": tool_call.arguments,
|
|
}));
|
|
}
|
|
}
|
|
blocks
|
|
};
|
|
converted.push(AnthropicMessage { role, content });
|
|
index += 1;
|
|
}
|
|
|
|
converted
|
|
}
|
|
|
|
fn native_anthropic_content(message: &Message) -> Option<Vec<Value>> {
|
|
if message.role != "assistant" {
|
|
return None;
|
|
}
|
|
let state = message.provider_state.as_ref()?;
|
|
if state.provider != "anthropic" {
|
|
return None;
|
|
}
|
|
let blocks = state.payload.get("content")?.as_array()?;
|
|
blocks
|
|
.iter()
|
|
.all(|block| {
|
|
block
|
|
.as_object()
|
|
.and_then(|value| value.get("type"))
|
|
.is_some()
|
|
})
|
|
.then(|| blocks.clone())
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct AnthropicTool {
|
|
name: String,
|
|
description: String,
|
|
input_schema: serde_json::Value,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
cache_control: Option<CacheControl>,
|
|
}
|
|
|
|
#[derive(Debug, Error)]
|
|
enum AnthropicStreamError {
|
|
#[error("invalid UTF-8 in Anthropic SSE event: {0}")]
|
|
Utf8(#[from] std::string::FromUtf8Error),
|
|
#[error("invalid Anthropic SSE payload: {0}")]
|
|
Json(#[from] serde_json::Error),
|
|
#[error("Anthropic stream error: {0}")]
|
|
Api(String),
|
|
#[error("Anthropic stream ended without message_stop")]
|
|
MissingFinish,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct AnthropicSseDecoder {
|
|
framer: SseFramer,
|
|
blocks: BTreeMap<usize, Value>,
|
|
tool_json: HashMap<usize, String>,
|
|
usage: Usage,
|
|
finish_reason: Option<FinishReason>,
|
|
done_emitted: bool,
|
|
}
|
|
|
|
impl AnthropicSseDecoder {
|
|
fn push(&mut self, bytes: &[u8]) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
|
|
let frames = self.framer.push(bytes)?;
|
|
self.decode_frames(frames)
|
|
}
|
|
|
|
fn finish(&mut self) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
|
|
let frames = self.framer.finish()?;
|
|
let chunks = self.decode_frames(frames)?;
|
|
if !self.done_emitted {
|
|
return Err(AnthropicStreamError::MissingFinish);
|
|
}
|
|
Ok(chunks)
|
|
}
|
|
|
|
fn decode_frames(
|
|
&mut self,
|
|
frames: Vec<String>,
|
|
) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
|
|
let mut chunks = Vec::new();
|
|
for data in frames {
|
|
let payload: Value = serde_json::from_str(&data)?;
|
|
match payload.get("type").and_then(Value::as_str) {
|
|
Some("message_start") => self.message_start(&payload, &mut chunks),
|
|
Some("content_block_start") => self.block_start(&payload, &mut chunks),
|
|
Some("content_block_delta") => self.block_delta(&payload, &mut chunks),
|
|
Some("content_block_stop") => self.block_stop(&payload),
|
|
Some("message_delta") => self.message_delta(&payload, &mut chunks),
|
|
Some("message_stop") => self.message_stop(&mut chunks),
|
|
Some("error") => {
|
|
let message = payload
|
|
.pointer("/error/message")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("unknown streaming error");
|
|
return Err(AnthropicStreamError::Api(message.to_string()));
|
|
}
|
|
Some("ping") | None | Some(_) => {}
|
|
}
|
|
}
|
|
Ok(chunks)
|
|
}
|
|
|
|
fn message_start(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
|
|
let message = payload.get("message").unwrap_or(&Value::Null);
|
|
let id = message
|
|
.get("id")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let model = message
|
|
.get("model")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
chunks.push(ProviderChunk::Metadata {
|
|
id: id.to_string(),
|
|
model: model.to_string(),
|
|
});
|
|
if let Some(usage) = message.get("usage") {
|
|
update_anthropic_usage(&mut self.usage, usage);
|
|
chunks.push(ProviderChunk::Usage(self.usage.clone()));
|
|
}
|
|
}
|
|
|
|
fn block_start(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
|
|
let Some(index) = event_index(payload) else {
|
|
return;
|
|
};
|
|
let Some(block) = payload.get("content_block") else {
|
|
return;
|
|
};
|
|
self.blocks.insert(index, block.clone());
|
|
match block.get("type").and_then(Value::as_str) {
|
|
Some("text") => {
|
|
if let Some(text) = block.get("text").and_then(Value::as_str)
|
|
&& !text.is_empty()
|
|
{
|
|
chunks.push(ProviderChunk::Text(text.to_string()));
|
|
}
|
|
}
|
|
Some("thinking") => {
|
|
if let Some(thinking) = block.get("thinking").and_then(Value::as_str)
|
|
&& !thinking.is_empty()
|
|
{
|
|
chunks.push(ProviderChunk::Reasoning(thinking.to_string()));
|
|
}
|
|
}
|
|
Some("tool_use") => {
|
|
chunks.push(ProviderChunk::ToolCallStart {
|
|
index,
|
|
id: block.get("id").and_then(Value::as_str).map(str::to_string),
|
|
name: block
|
|
.get("name")
|
|
.and_then(Value::as_str)
|
|
.map(str::to_string),
|
|
});
|
|
let input = block.get("input").cloned().unwrap_or(Value::Null);
|
|
if !input.is_null() && input != serde_json::json!({}) {
|
|
chunks.push(ProviderChunk::ToolCallArguments {
|
|
index,
|
|
delta: input.to_string(),
|
|
});
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
fn block_delta(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
|
|
let Some(index) = event_index(payload) else {
|
|
return;
|
|
};
|
|
let Some(delta) = payload.get("delta") else {
|
|
return;
|
|
};
|
|
match delta.get("type").and_then(Value::as_str) {
|
|
Some("text_delta") => {
|
|
if let Some(text) = delta.get("text").and_then(Value::as_str) {
|
|
append_block_string(&mut self.blocks, index, "text", text);
|
|
if !text.is_empty() {
|
|
chunks.push(ProviderChunk::Text(text.to_string()));
|
|
}
|
|
}
|
|
}
|
|
Some("thinking_delta") => {
|
|
if let Some(thinking) = delta.get("thinking").and_then(Value::as_str) {
|
|
append_block_string(&mut self.blocks, index, "thinking", thinking);
|
|
if !thinking.is_empty() {
|
|
chunks.push(ProviderChunk::Reasoning(thinking.to_string()));
|
|
}
|
|
}
|
|
}
|
|
Some("signature_delta") => {
|
|
if let Some(signature) = delta.get("signature").and_then(Value::as_str) {
|
|
append_block_string(&mut self.blocks, index, "signature", signature);
|
|
}
|
|
}
|
|
Some("input_json_delta") => {
|
|
if let Some(partial) = delta.get("partial_json").and_then(Value::as_str) {
|
|
self.tool_json.entry(index).or_default().push_str(partial);
|
|
if !partial.is_empty() {
|
|
chunks.push(ProviderChunk::ToolCallArguments {
|
|
index,
|
|
delta: partial.to_string(),
|
|
});
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
fn block_stop(&mut self, payload: &Value) {
|
|
let Some(index) = event_index(payload) else {
|
|
return;
|
|
};
|
|
let Some(json) = self.tool_json.remove(&index) else {
|
|
return;
|
|
};
|
|
let input = serde_json::from_str(&json).unwrap_or(Value::Null);
|
|
if let Some(block) = self.blocks.get_mut(&index)
|
|
&& let Some(object) = block.as_object_mut()
|
|
{
|
|
object.insert("input".to_string(), input);
|
|
}
|
|
}
|
|
|
|
fn message_delta(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
|
|
if let Some(reason) = payload
|
|
.pointer("/delta/stop_reason")
|
|
.and_then(Value::as_str)
|
|
{
|
|
self.finish_reason = Some(FinishReason::from_provider(reason));
|
|
}
|
|
if let Some(usage) = payload.get("usage") {
|
|
update_anthropic_usage(&mut self.usage, usage);
|
|
chunks.push(ProviderChunk::Usage(self.usage.clone()));
|
|
}
|
|
}
|
|
|
|
fn message_stop(&mut self, chunks: &mut Vec<ProviderChunk>) {
|
|
if self.done_emitted {
|
|
return;
|
|
}
|
|
let content = self.blocks.values().cloned().collect::<Vec<_>>();
|
|
chunks.push(ProviderChunk::ProviderState(ProviderReasoningState {
|
|
provider: "anthropic".to_string(),
|
|
payload: serde_json::json!({ "version": 1, "content": content }),
|
|
}));
|
|
chunks.push(ProviderChunk::Done(
|
|
self.finish_reason.clone().unwrap_or(FinishReason::Stop),
|
|
));
|
|
self.done_emitted = true;
|
|
}
|
|
}
|
|
|
|
fn event_index(payload: &Value) -> Option<usize> {
|
|
payload
|
|
.get("index")
|
|
.and_then(Value::as_u64)
|
|
.and_then(|value| usize::try_from(value).ok())
|
|
}
|
|
|
|
fn append_block_string(blocks: &mut BTreeMap<usize, Value>, index: usize, key: &str, delta: &str) {
|
|
let Some(object) = blocks.get_mut(&index).and_then(Value::as_object_mut) else {
|
|
return;
|
|
};
|
|
let value = object
|
|
.entry(key.to_string())
|
|
.or_insert_with(|| Value::String(String::new()));
|
|
if let Some(current) = value.as_str() {
|
|
*value = Value::String(format!("{current}{delta}"));
|
|
}
|
|
}
|
|
|
|
fn update_anthropic_usage(usage: &mut Usage, value: &Value) {
|
|
if let Some(input) = json_u32(value, "input_tokens") {
|
|
usage.prompt_tokens = input;
|
|
}
|
|
if let Some(output) = json_u32(value, "output_tokens") {
|
|
usage.completion_tokens = output;
|
|
}
|
|
if let Some(cache_read) = json_u32(value, "cache_read_input_tokens") {
|
|
usage.cached_tokens = Some(cache_read);
|
|
usage.cache_read_input_tokens = Some(cache_read);
|
|
}
|
|
if let Some(cache_creation) = json_u32(value, "cache_creation_input_tokens") {
|
|
usage.cache_creation_input_tokens = Some(cache_creation);
|
|
}
|
|
usage.total_tokens = usage.prompt_tokens.saturating_add(usage.completion_tokens);
|
|
}
|
|
|
|
fn json_u32(value: &Value, key: &str) -> Option<u32> {
|
|
value
|
|
.get(key)
|
|
.and_then(Value::as_u64)
|
|
.and_then(|number| u32::try_from(number).ok())
|
|
}
|
|
|
|
struct AnthropicHttpStream {
|
|
response: reqwest::Response,
|
|
decoder: AnthropicSseDecoder,
|
|
pending: VecDeque<ProviderChunk>,
|
|
reached_eof: bool,
|
|
}
|
|
|
|
async fn next_anthropic_chunk(
|
|
mut state: AnthropicHttpStream,
|
|
) -> Result<Option<(ProviderChunk, AnthropicHttpStream)>, DynProviderError> {
|
|
loop {
|
|
if let Some(chunk) = state.pending.pop_front() {
|
|
return Ok(Some((chunk, state)));
|
|
}
|
|
if state.reached_eof {
|
|
return Ok(None);
|
|
}
|
|
match state.response.chunk().await? {
|
|
Some(bytes) => state.pending.extend(state.decoder.push(&bytes)?),
|
|
None => {
|
|
state.pending.extend(state.decoder.finish()?);
|
|
state.reached_eof = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl LLMProvider for AnthropicProvider {
|
|
async fn stream(
|
|
&self,
|
|
request: ChatCompletionRequest,
|
|
) -> Result<ProviderStream, DynProviderError> {
|
|
let start = std::time::Instant::now();
|
|
let url = format!("{}/v1/messages", self.base_url);
|
|
let max_tokens = request.max_tokens.or(self.max_tokens).unwrap_or(1024);
|
|
|
|
let tools = request.tools.map(|tools| {
|
|
tools
|
|
.iter()
|
|
.map(|t: &Tool| AnthropicTool {
|
|
name: t.function.name.clone(),
|
|
description: t.function.description.clone(),
|
|
input_schema: t.function.parameters.clone(),
|
|
cache_control: Some(CacheControl::ephemeral()),
|
|
})
|
|
.collect()
|
|
});
|
|
|
|
let body = AnthropicRequest {
|
|
model: self.model_id.clone(),
|
|
messages: convert_messages(&request.messages),
|
|
max_tokens,
|
|
temperature: request.temperature.or(self.temperature),
|
|
stream: true,
|
|
tools,
|
|
extra: self.model_extra.clone(),
|
|
};
|
|
|
|
let mut req_builder = self
|
|
.client
|
|
.post(&url)
|
|
.header("x-api-key", &self.api_key)
|
|
.header("anthropic-version", "2023-06-01")
|
|
.header("Content-Type", "application/json");
|
|
|
|
for (key, value) in &self.extra_headers {
|
|
req_builder = req_builder.header(key.as_str(), value.as_str());
|
|
}
|
|
|
|
let request_summary = super::stream::diagnostic_request_summary(
|
|
&self.model_id,
|
|
body.messages.len(),
|
|
body.tools.as_ref().map_or(0, Vec::len),
|
|
);
|
|
tracing::debug!(
|
|
message_count = body.messages.len(),
|
|
tool_count = body.tools.as_ref().map_or(0, Vec::len),
|
|
"Anthropic streaming request"
|
|
);
|
|
|
|
let resp = req_builder.json(&body).send().await.inspect_err(|e| {
|
|
let is_timeout = e.is_timeout();
|
|
tracing::error!(
|
|
provider = %self.name,
|
|
model = %self.model_id,
|
|
url = %url,
|
|
timeout = is_timeout,
|
|
error = %e,
|
|
elapsed_ms = %start.elapsed().as_millis(),
|
|
"LLM API request failed"
|
|
);
|
|
})?;
|
|
|
|
let status = resp.status();
|
|
if !status.is_success() {
|
|
let body_text = resp.text().await?;
|
|
let error_msg = serde_json::from_str::<serde_json::Value>(&body_text)
|
|
.ok()
|
|
.and_then(|v| {
|
|
v.get("error")
|
|
.and_then(|e| e.get("message"))
|
|
.and_then(|m| m.as_str())
|
|
.map(|s| s.to_string())
|
|
})
|
|
.unwrap_or_else(|| body_text.clone());
|
|
tracing::error!(
|
|
provider = %self.name,
|
|
model = %self.model_id,
|
|
http_status = %status,
|
|
error = %error_msg,
|
|
elapsed_ms = %start.elapsed().as_millis(),
|
|
"LLM API returned error"
|
|
);
|
|
if let Some(ref storage) = self.storage {
|
|
let _ = storage
|
|
.append_llm_call(
|
|
&self.name,
|
|
&self.model_id,
|
|
&request_summary,
|
|
Some(&body_text),
|
|
Some(&error_msg),
|
|
start.elapsed().as_millis() as u64,
|
|
)
|
|
.await;
|
|
}
|
|
return Err(format!("API error ({}): {}", status.as_u16(), error_msg).into());
|
|
}
|
|
tracing::debug!(status = %status, "Anthropic streaming response started");
|
|
Ok(Box::pin(stream::try_unfold(
|
|
AnthropicHttpStream {
|
|
response: resp,
|
|
decoder: AnthropicSseDecoder::default(),
|
|
pending: VecDeque::new(),
|
|
reached_eof: false,
|
|
},
|
|
next_anthropic_chunk,
|
|
)))
|
|
}
|
|
|
|
fn ptype(&self) -> &str {
|
|
"anthropic"
|
|
}
|
|
|
|
fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn model_id(&self) -> &str {
|
|
&self.model_id
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::providers::ProviderResponseAccumulator;
|
|
use serde_json::json;
|
|
|
|
#[test]
|
|
fn test_convert_content_blocks_adds_cache_control_for_system_text() {
|
|
let blocks = vec![ContentBlock::text("hello")];
|
|
let serialized = convert_content_blocks(&blocks, true);
|
|
|
|
assert_eq!(serialized[0]["type"], "text");
|
|
assert_eq!(serialized[0]["cache_control"]["type"], "ephemeral");
|
|
}
|
|
|
|
#[test]
|
|
fn test_convert_content_blocks_leaves_user_text_uncached() {
|
|
let blocks = vec![ContentBlock::text("hello")];
|
|
let serialized = convert_content_blocks(&blocks, false);
|
|
|
|
assert!(serialized[0].get("cache_control").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_anthropic_tool_serializes_cache_control() {
|
|
let tool = AnthropicTool {
|
|
name: "alpha".to_string(),
|
|
description: "desc".to_string(),
|
|
input_schema: json!({}),
|
|
cache_control: Some(CacheControl::ephemeral()),
|
|
};
|
|
|
|
let value = serde_json::to_value(tool).unwrap();
|
|
assert_eq!(value["cache_control"]["type"], "ephemeral");
|
|
}
|
|
|
|
#[test]
|
|
fn tool_result_preserves_native_image_blocks() {
|
|
let messages = vec![Message {
|
|
role: "tool".to_string(),
|
|
content: vec![
|
|
ContentBlock::text("image ready"),
|
|
ContentBlock::image_url("data:image/png;base64,AAAA"),
|
|
],
|
|
reasoning_content: None,
|
|
provider_state: None,
|
|
tool_call_id: Some("call_1".to_string()),
|
|
name: Some("file_read".to_string()),
|
|
tool_calls: None,
|
|
}];
|
|
|
|
let converted = convert_messages(&messages);
|
|
let result = &converted[0].content[0];
|
|
|
|
assert_eq!(converted[0].role, "user");
|
|
assert_eq!(result["type"], "tool_result");
|
|
assert_eq!(result["content"][0]["type"], "text");
|
|
assert_eq!(result["content"][1]["type"], "image");
|
|
assert_eq!(result["content"][1]["source"]["media_type"], "image/png");
|
|
assert_eq!(result["content"][1]["source"]["data"], "AAAA");
|
|
}
|
|
|
|
#[test]
|
|
fn tool_results_and_following_steering_share_one_user_content_array() {
|
|
let messages = vec![
|
|
Message::tool("call_1", "lookup", "first result"),
|
|
Message::tool("call_2", "lookup", "second result"),
|
|
Message::user("用户补充指令"),
|
|
Message::user("再补充一条"),
|
|
Message::assistant("最终回答"),
|
|
];
|
|
|
|
let converted = convert_messages(&messages);
|
|
|
|
assert_eq!(converted.len(), 2);
|
|
assert_eq!(converted[0].role, "user");
|
|
assert_eq!(converted[0].content.len(), 4);
|
|
assert_eq!(converted[0].content[0]["type"], "tool_result");
|
|
assert_eq!(converted[0].content[0]["tool_use_id"], "call_1");
|
|
assert_eq!(converted[0].content[1]["tool_use_id"], "call_2");
|
|
assert_eq!(
|
|
converted[0].content[2],
|
|
json!({"type": "text", "text": "用户补充指令"})
|
|
);
|
|
assert_eq!(
|
|
converted[0].content[3],
|
|
json!({"type": "text", "text": "再补充一条"})
|
|
);
|
|
assert_eq!(converted[1].role, "assistant");
|
|
}
|
|
|
|
#[test]
|
|
fn native_stream_decodes_thinking_signature_tools_usage_and_replay_state() {
|
|
let events = [
|
|
json!({"type":"message_start","message":{"id":"msg_1","model":"claude-test","usage":{"input_tokens":11,"output_tokens":0,"cache_read_input_tokens":3}}}),
|
|
json!({"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}),
|
|
json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"check "}}),
|
|
json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"facts"}}),
|
|
json!({"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig=="}}),
|
|
json!({"type":"content_block_stop","index":0}),
|
|
json!({"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_1","name":"lookup","input":{}}}),
|
|
json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}),
|
|
json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"rust\"}"}}),
|
|
json!({"type":"content_block_stop","index":1}),
|
|
json!({"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}),
|
|
json!({"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"answer"}}),
|
|
json!({"type":"content_block_stop","index":2}),
|
|
json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":7}}),
|
|
json!({"type":"message_stop"}),
|
|
];
|
|
let wire = events
|
|
.iter()
|
|
.map(|event| format!("event: ignored\ndata: {event}\n\n"))
|
|
.collect::<String>();
|
|
let mut decoder = AnthropicSseDecoder::default();
|
|
let mut accumulator = ProviderResponseAccumulator::default();
|
|
for bytes in wire.as_bytes().chunks(7) {
|
|
for chunk in decoder.push(bytes).unwrap() {
|
|
accumulator.push(chunk);
|
|
}
|
|
}
|
|
for chunk in decoder.finish().unwrap() {
|
|
accumulator.push(chunk);
|
|
}
|
|
let response = accumulator.finish();
|
|
|
|
assert_eq!(response.id, "msg_1");
|
|
assert_eq!(response.model, "claude-test");
|
|
assert_eq!(response.reasoning_content.as_deref(), Some("check facts"));
|
|
assert_eq!(response.content, "answer");
|
|
assert_eq!(response.tool_calls.len(), 1);
|
|
assert_eq!(response.tool_calls[0].id, "tool_1");
|
|
assert_eq!(response.tool_calls[0].arguments, json!({"q":"rust"}));
|
|
assert_eq!(response.usage.prompt_tokens, 11);
|
|
assert_eq!(response.usage.completion_tokens, 7);
|
|
assert_eq!(response.usage.total_tokens, 18);
|
|
assert_eq!(response.usage.cache_read_input_tokens, Some(3));
|
|
|
|
let state = response.provider_state.unwrap();
|
|
assert_eq!(state.provider, "anthropic");
|
|
assert_eq!(state.payload["content"][0]["thinking"], "check facts");
|
|
assert_eq!(state.payload["content"][0]["signature"], "sig==");
|
|
assert_eq!(state.payload["content"][1]["input"], json!({"q":"rust"}));
|
|
assert_eq!(state.payload["content"][2]["text"], "answer");
|
|
}
|
|
|
|
#[test]
|
|
fn matching_provider_state_replays_native_blocks_without_generic_duplicates() {
|
|
let native = json!([
|
|
{"type":"thinking","thinking":"signed thought","signature":"sig=="},
|
|
{"type":"tool_use","id":"tool_1","name":"lookup","input":{"q":"rust"}}
|
|
]);
|
|
let message = Message {
|
|
role: "assistant".into(),
|
|
content: vec![ContentBlock::text("generic text must not be appended")],
|
|
reasoning_content: Some("display copy".into()),
|
|
provider_state: Some(ProviderReasoningState {
|
|
provider: "anthropic".into(),
|
|
payload: json!({"version":1,"content":native}),
|
|
}),
|
|
tool_call_id: None,
|
|
name: None,
|
|
tool_calls: Some(vec![crate::providers::ToolCall {
|
|
id: "duplicate".into(),
|
|
name: "duplicate".into(),
|
|
arguments: json!({}),
|
|
}]),
|
|
};
|
|
|
|
let converted = convert_messages(&[message]);
|
|
|
|
assert_eq!(converted[0].content.len(), 2);
|
|
assert_eq!(converted[0].content[0]["signature"], "sig==");
|
|
assert_eq!(converted[0].content[1]["id"], "tool_1");
|
|
}
|
|
|
|
#[test]
|
|
fn foreign_provider_state_is_not_replayed_to_anthropic() {
|
|
let message = Message {
|
|
role: "assistant".into(),
|
|
content: vec![ContentBlock::text("answer")],
|
|
reasoning_content: Some("unsigned display reasoning".into()),
|
|
provider_state: Some(ProviderReasoningState {
|
|
provider: "openai".into(),
|
|
payload: json!({"private":"state"}),
|
|
}),
|
|
tool_call_id: None,
|
|
name: None,
|
|
tool_calls: None,
|
|
};
|
|
|
|
let converted = convert_messages(&[message]);
|
|
|
|
assert_eq!(
|
|
converted[0].content,
|
|
vec![json!({"type":"text","text":"answer"})]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn stream_requires_message_stop() {
|
|
let mut decoder = AnthropicSseDecoder::default();
|
|
decoder
|
|
.push(b"data: {\"type\":\"message_start\",\"message\":{}}\n\n")
|
|
.unwrap();
|
|
|
|
assert!(matches!(
|
|
decoder.finish(),
|
|
Err(AnthropicStreamError::MissingFinish)
|
|
));
|
|
}
|
|
}
|