use futures_util::{StreamExt as _, TryStreamExt as _, stream};
use std::sync::{Arc, Mutex};
use anyhow::{Context as _, Result};
use async_anthropic::types::{
CreateMessagesRequestBuilder, Message, MessageBuilder, MessageContent, MessageContentList,
MessageRole, MessagesStreamEvent, ToolChoice, ToolResultBuilder, ToolUseBuilder,
};
use async_trait::async_trait;
use serde_json::{Value, json};
use swiftide_core::{
ChatCompletion, ChatCompletionStream,
chat_completion::{
ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ToolCall, ToolOutput, ToolSpec,
Usage, UsageBuilder, errors::LanguageModelError,
},
};
use super::Anthropic;
use super::tool_schema::AnthropicToolSchema;
#[cfg(feature = "metrics")]
use swiftide_core::metrics::emit_usage;
#[async_trait]
impl ChatCompletion for Anthropic {
#[tracing::instrument(skip_all, err)]
async fn complete(
&self,
request: &ChatCompletionRequest<'_>,
) -> Result<ChatCompletionResponse, LanguageModelError> {
let model = &self.default_options.prompt_model;
let request = self
.build_request(request)
.and_then(|b| b.build().map_err(LanguageModelError::permanent))?;
tracing::debug!(
model = &model,
messages = serde_json::to_string_pretty(&request).expect("Infallible"),
"[ChatCompletion] Request to anthropic"
);
let response = self
.client
.messages()
.create(request)
.await
.map_err(LanguageModelError::permanent)?;
tracing::debug!(
response = serde_json::to_string_pretty(&response).expect("Infallible"),
"[ChatCompletion] Response from anthropic"
);
let maybe_tool_calls = response
.messages()
.iter()
.flat_map(Message::tool_uses)
.map(|atool| {
ToolCall::builder()
.id(atool.id)
.name(atool.name)
.args(atool.input.to_string())
.build()
.expect("infallible")
})
.collect::<Vec<_>>();
let maybe_tool_calls = if maybe_tool_calls.is_empty() {
None
} else {
Some(maybe_tool_calls)
};
let mut builder = ChatCompletionResponse::builder()
.maybe_message(response.messages().iter().find_map(Message::text))
.maybe_tool_calls(maybe_tool_calls)
.to_owned();
if let Some(usage) = &response.usage {
let input_tokens = usage.input_tokens.unwrap_or_default();
let output_tokens = usage.output_tokens.unwrap_or_default();
let total_tokens = input_tokens + output_tokens;
#[cfg(feature = "metrics")]
emit_usage(
model,
input_tokens.into(),
output_tokens.into(),
total_tokens.into(),
self.metric_metadata.as_ref(),
);
let usage = Usage {
prompt_tokens: input_tokens,
completion_tokens: output_tokens,
total_tokens,
details: None,
};
if let Some(callback) = &self.on_usage {
callback(&usage).await?;
}
let usage = UsageBuilder::default()
.prompt_tokens(input_tokens)
.completion_tokens(output_tokens)
.total_tokens(total_tokens)
.build()
.map_err(LanguageModelError::permanent)?;
builder.usage(usage);
}
builder.build().map_err(LanguageModelError::from)
}
#[tracing::instrument(skip_all)]
async fn complete_stream(&self, request: &ChatCompletionRequest<'_>) -> ChatCompletionStream {
let model = &self.default_options.prompt_model;
let request = match self
.build_request(request)
.and_then(|b| b.build().map_err(LanguageModelError::permanent))
{
Ok(request) => request,
Err(e) => {
return e.into();
}
};
tracing::debug!(
model = &model,
messages = serde_json::to_string_pretty(&request).expect("Infallible"),
"[ChatCompletion] Request to anthropic"
);
let response = self.client.messages().create_stream(request).await;
let accumulating_response = Arc::new(Mutex::new(ChatCompletionResponse::default()));
let final_response = Arc::clone(&accumulating_response);
#[cfg(feature = "metrics")]
let model = model.clone();
#[cfg(feature = "metrics")]
let metric_metadata = self.metric_metadata.clone();
let maybe_usage_callback = self.on_usage.clone();
response
.map_ok(move |chunk| {
let accumulating_response = Arc::clone(&accumulating_response);
let mut lock = accumulating_response.lock().unwrap();
append_delta_from_chunk(&chunk, &mut lock);
lock.clone()
})
.map_err(LanguageModelError::permanent)
.chain(
stream::iter(vec![final_response]).map(move |final_response| {
if let Some(usage) = final_response.lock().unwrap().usage.as_ref() {
let usage = usage.clone();
if let Some(callback) = maybe_usage_callback.as_ref() {
let usage = usage.clone();
let callback = callback.clone();
tokio::spawn(async move {
if let Err(e) = callback(&usage).await {
tracing::error!("Error in on_usage callback: {}", e);
}
});
}
#[cfg(feature = "metrics")]
emit_usage(
&model,
usage.prompt_tokens.into(),
usage.completion_tokens.into(),
usage.total_tokens.into(),
metric_metadata.as_ref(),
);
}
Ok(final_response.lock().unwrap().clone())
}),
)
.boxed()
}
}
#[allow(clippy::collapsible_match)]
fn append_delta_from_chunk(chunk: &MessagesStreamEvent, lock: &mut ChatCompletionResponse) {
match chunk {
MessagesStreamEvent::ContentBlockStart {
index,
content_block,
} => match content_block {
MessageContent::ToolUse(tool_use) => {
lock.append_tool_call_delta(*index, Some(&tool_use.id), Some(&tool_use.name), None);
}
MessageContent::Text(text) => {
lock.append_message_delta(Some(&text.text));
}
MessageContent::ToolResult(_tool_result) => (),
},
MessagesStreamEvent::ContentBlockDelta { index, delta } => match delta {
async_anthropic::types::ContentBlockDelta::TextDelta { text } => {
lock.append_message_delta(Some(text));
}
async_anthropic::types::ContentBlockDelta::InputJsonDelta { partial_json } => {
lock.append_tool_call_delta(*index, None, None, Some(partial_json));
}
},
#[allow(clippy::cast_possible_truncation)]
MessagesStreamEvent::MessageDelta { usage, .. } => {
if let Some(usage) = usage {
let input_tokens = usage.input_tokens.unwrap_or_default();
let output_tokens = usage.output_tokens.unwrap_or_default();
let total_tokens = input_tokens + output_tokens;
lock.append_usage_delta(input_tokens, output_tokens, total_tokens);
}
}
MessagesStreamEvent::MessageStart { message, usage } => {
if let Some(usage) = usage {
let input_tokens = usage.input_tokens.unwrap_or_default();
let output_tokens = usage.output_tokens.unwrap_or_default();
let total_tokens = input_tokens + output_tokens;
lock.append_usage_delta(input_tokens, output_tokens, total_tokens);
}
if let Some(message_usage) = &message.usage {
let input_tokens = message_usage.input_tokens.unwrap_or_default();
let output_tokens = message_usage.output_tokens.unwrap_or_default();
let total_tokens = input_tokens + output_tokens;
lock.append_usage_delta(input_tokens, output_tokens, total_tokens);
}
}
_ => {}
}
}
impl Anthropic {
fn build_request(
&self,
request: &ChatCompletionRequest<'_>,
) -> Result<async_anthropic::types::CreateMessagesRequestBuilder, LanguageModelError> {
let model = &self.default_options.prompt_model;
let mut messages = request.messages().to_vec();
let maybe_system = messages
.iter()
.position(ChatMessage::is_system)
.map(|idx| messages.remove(idx));
let messages = messages_to_antropic(&messages)?;
let mut anthropic_request = CreateMessagesRequestBuilder::default()
.model(model)
.messages(messages)
.to_owned();
if let Some(ChatMessage::System(system)) = maybe_system {
anthropic_request.system(system);
}
if !request.tools_spec().is_empty() {
anthropic_request
.tools(
request
.tools_spec()
.iter()
.map(tools_to_anthropic)
.collect::<Result<Vec<_>>>()?,
)
.tool_choice(ToolChoice::Auto);
}
Ok(anthropic_request)
}
}
fn messages_to_antropic(messages: &[ChatMessage]) -> Result<Vec<Message>> {
let mut anthropic_messages = Vec::with_capacity(messages.len());
let mut messages = messages.iter().peekable();
while let Some(message) = messages.next() {
match message {
ChatMessage::ToolOutput(tool_call, tool_output) => {
let mut content = vec![tool_result_to_anthropic(tool_call, tool_output)?];
while let Some(ChatMessage::ToolOutput(tool_call, tool_output)) = messages.peek() {
content.push(tool_result_to_anthropic(tool_call, tool_output)?);
messages.next();
}
anthropic_messages.push(
MessageBuilder::default()
.role(MessageRole::User)
.content(MessageContentList(content))
.build()
.context("Failed to build message")?,
);
}
_ => {
if let Some(message) = message_to_antropic(message)? {
anthropic_messages.push(message);
}
}
}
}
Ok(anthropic_messages)
}
fn tool_result_to_anthropic(
tool_call: &ToolCall,
tool_output: &ToolOutput,
) -> Result<MessageContent> {
Ok(ToolResultBuilder::default()
.tool_use_id(tool_call.id())
.content(tool_output.content().unwrap_or("Success"))
.build()?
.into())
}
#[allow(clippy::items_after_statements)]
fn message_to_antropic(message: &ChatMessage) -> Result<Option<Message>> {
let mut builder = MessageBuilder::default().role(MessageRole::User).to_owned();
match message {
ChatMessage::ToolOutput(tool_call, tool_output) => builder.content(MessageContentList(
vec![tool_result_to_anthropic(tool_call, tool_output)?],
)),
ChatMessage::Summary(msg) | ChatMessage::System(msg) => builder.content(msg.as_str()),
ChatMessage::User(content) => builder.content(content.as_str()),
ChatMessage::UserWithParts(parts) => {
if parts.iter().any(|part| {
!matches!(
part,
swiftide_core::chat_completion::ChatMessageContentPart::Text { .. }
)
}) {
anyhow::bail!("Anthropic chat completions only support text message parts");
}
let text_parts = parts
.iter()
.filter_map(|part| match part {
swiftide_core::chat_completion::ChatMessageContentPart::Text { text } => {
Some(text.as_ref())
}
swiftide_core::chat_completion::ChatMessageContentPart::Image { .. }
| swiftide_core::chat_completion::ChatMessageContentPart::Document { .. }
| swiftide_core::chat_completion::ChatMessageContentPart::Audio { .. }
| swiftide_core::chat_completion::ChatMessageContentPart::Video { .. } => None,
})
.collect::<Vec<_>>();
builder.content(text_parts.join(" "))
}
ChatMessage::Assistant(content, tool_calls) => {
builder.role(MessageRole::Assistant);
let mut content_list: Vec<MessageContent> = Vec::new();
if let Some(content) = content.as_ref() {
content_list.push(content.clone().into());
}
if let Some(tool_calls) = tool_calls.as_ref() {
for tool_call in tool_calls {
let tool_call = ToolUseBuilder::default()
.id(tool_call.id())
.name(tool_call.name())
.input(tool_call.args().and_then(|v| v.parse::<Value>().ok()))
.build()?;
content_list.push(tool_call.into());
}
}
if content_list.is_empty() {
return Ok(None);
}
let content_list = MessageContentList(content_list);
builder.content(content_list)
}
ChatMessage::Reasoning(_) => return Ok(None),
};
builder.build().context("Failed to build message").map(Some)
}
fn tools_to_anthropic(
spec: &ToolSpec,
) -> Result<serde_json::value::Map<String, serde_json::Value>> {
let mut map = json!({
"name": &spec.name,
"description": &spec.description,
})
.as_object_mut()
.context("Failed to build tool")?
.to_owned();
let schema = AnthropicToolSchema::try_from(spec)
.context("tool schema must be Anthropic compatible")?
.into_value();
map.insert("input_schema".to_string(), schema);
Ok(map)
}
#[cfg(test)]
mod tests {
use super::*;
use schemars::{JsonSchema, schema_for};
use swiftide_core::{
AgentContext, Tool,
chat_completion::{ChatCompletionRequest, ChatMessage},
};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_partial_json, method, path},
};
#[derive(Clone)]
struct FakeTool();
#[derive(Clone)]
struct AlphaTool();
#[derive(JsonSchema, serde::Serialize, serde::Deserialize)]
struct LocationArgs {
#[schemars(description = "City or region to retrieve weather for")]
location: String,
}
#[derive(JsonSchema, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct NestedCommentArgs {
request: NestedCommentRequest,
}
#[derive(JsonSchema, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct NestedCommentRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
body: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
page_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
block_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
discussion_id: Option<String>,
}
#[async_trait]
impl Tool for FakeTool {
async fn invoke(
&self,
_agent_context: &dyn AgentContext,
_tool_call: &ToolCall,
) -> std::result::Result<
swiftide_core::chat_completion::ToolOutput,
swiftide_core::chat_completion::errors::ToolError,
> {
todo!()
}
fn name(&self) -> std::borrow::Cow<'_, str> {
"get_weather".into()
}
fn tool_spec(&self) -> ToolSpec {
ToolSpec::builder()
.description("Gets the weather")
.name("get_weather")
.parameters_schema(schema_for!(LocationArgs))
.build()
.unwrap()
}
}
#[async_trait]
impl Tool for AlphaTool {
async fn invoke(
&self,
_agent_context: &dyn AgentContext,
_tool_call: &ToolCall,
) -> std::result::Result<
swiftide_core::chat_completion::ToolOutput,
swiftide_core::chat_completion::errors::ToolError,
> {
todo!()
}
fn name(&self) -> std::borrow::Cow<'_, str> {
"alpha_tool".into()
}
fn tool_spec(&self) -> ToolSpec {
ToolSpec::builder()
.name("alpha_tool")
.description("Alpha tool")
.parameters_schema(schemars::schema_for!(LocationArgs))
.build()
.unwrap()
}
}
#[test_log::test(tokio::test)]
async fn test_complete_without_tools() {
let mock_server = MockServer::start().await;
let mock_response = ResponseTemplate::new(200).set_body_json(serde_json::json!({
"content": [{"type": "text", "text": "mocked response"}]
}));
Mock::given(method("POST"))
.and(path("/v1/messages")) .respond_with(mock_response)
.mount(&mock_server)
.await;
let client = async_anthropic::Client::builder()
.base_url(mock_server.uri())
.build()
.unwrap();
let mut client_builder = Anthropic::builder();
client_builder.client(client);
let client = client_builder.build().unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hello".into())])
.build()
.unwrap();
let result = client.complete(&request).await.unwrap();
assert_eq!(result.message, Some("mocked response".into()));
assert!(result.tool_calls.is_none());
}
#[test_log::test(tokio::test)]
async fn test_complete_with_tools() {
let mock_server = MockServer::start().await;
let mock_response = ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": "msg_016zKNb88WhhgBQXhSaQf1rs",
"content": [
{
"type": "text",
"text": "I'll check the current weather in San Francisco, CA for you."
},
{
"type": "tool_use",
"id": "toolu_01E1yxpxXU4hBgCMLzPL1FuR",
"input": {
"location": "San Francisco, CA"
},
"name": "get_weather"
}
],
"model": "claude-3-5-sonnet-20241022",
"stop_reason": "tool_use",
"stop_sequence": null,
"usage": {
"input_tokens": 403,
"output_tokens": 71
}
}));
Mock::given(method("POST"))
.and(path("/v1/messages")) .respond_with(mock_response)
.mount(&mock_server)
.await;
let client = async_anthropic::Client::builder()
.base_url(mock_server.uri())
.build()
.unwrap();
let mut client_builder = Anthropic::builder();
client_builder.client(client);
let client = client_builder.build().unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hello".into())])
.tool_specs([FakeTool().tool_spec()])
.build()
.unwrap();
let result = client.complete(&request).await.unwrap();
assert_eq!(
result.message,
Some("I'll check the current weather in San Francisco, CA for you.".into())
);
assert!(result.tool_calls.is_some());
let Some(tool_call) = result.tool_calls.and_then(|f| f.first().cloned()) else {
panic!("No tool call found")
};
assert_eq!(tool_call.name(), "get_weather");
assert_eq!(
tool_call.args(),
Some(
json!({"location": "San Francisco, CA"})
.to_string()
.as_str()
)
);
}
#[test]
fn test_build_request_orders_tools_deterministically() {
let client = Anthropic::builder().build().unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hello".into())])
.tool_specs([FakeTool().tool_spec(), AlphaTool().tool_spec()])
.build()
.unwrap();
let built = client.build_request(&request).unwrap().build().unwrap();
let tool_names = built
.tools
.expect("tools present")
.into_iter()
.map(|tool| {
tool.get("name")
.and_then(serde_json::Value::as_str)
.expect("tool name")
.to_owned()
})
.collect::<Vec<_>>();
assert_eq!(tool_names, vec!["alpha_tool", "get_weather"]);
}
#[test_log::test(tokio::test)]
async fn test_complete_with_system_prompt() {
let mock_server = MockServer::start().await;
let mock_response = ResponseTemplate::new(200).set_body_json(serde_json::json!({
"content": [{"type": "text", "text": "Response with system prompt"}],
"usage": {
"input_tokens": 19,
"output_tokens": 10,
}
}));
Mock::given(method("POST"))
.and(path("/v1/messages")) .and(body_partial_json(json!({
"system": "System message",
"messages":[{"role":"user","content":[{"type":"text","text":"Hello"}]}]
})))
.respond_with(mock_response)
.mount(&mock_server)
.await;
let client = async_anthropic::Client::builder()
.base_url(mock_server.uri())
.build()
.unwrap();
let mut client_builder = Anthropic::builder();
client_builder.client(client);
let client = client_builder.build().unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![
ChatMessage::System("System message".into()),
ChatMessage::User("Hello".into()),
])
.build()
.unwrap();
let response = client.complete(&request).await.unwrap();
assert_eq!(response.message, Some("Response with system prompt".into()));
let usage = response.usage.unwrap();
assert_eq!(usage.prompt_tokens, 19);
assert_eq!(usage.completion_tokens, 10);
assert_eq!(usage.total_tokens, 29);
}
#[test]
fn test_tools_to_anthropic() {
let tool_spec = ToolSpec::builder()
.description("Gets the weather")
.name("get_weather")
.parameters_schema(schema_for!(LocationArgs))
.build()
.unwrap();
let result = tools_to_anthropic(&tool_spec).unwrap();
let expected_schema = tool_spec.strict_parameters_schema().unwrap().into_json();
assert_eq!(
expected_schema["properties"]["location"]["description"],
json!("City or region to retrieve weather for")
);
let expected = json!({
"name": "get_weather",
"description": "Gets the weather",
"input_schema": expected_schema,
});
assert_eq!(serde_json::Value::Object(result), expected);
}
#[test]
fn test_tools_to_anthropic_preserves_optional_nested_fields() {
let tool_spec = ToolSpec::builder()
.description("Creates a comment")
.name("create_comment")
.parameters_schema(schema_for!(NestedCommentArgs))
.build()
.unwrap();
let result = tools_to_anthropic(&tool_spec).unwrap();
let input_schema = result
.get("input_schema")
.and_then(Value::as_object)
.expect("anthropic tool should contain input_schema");
assert_eq!(
input_schema.get("type"),
Some(&Value::String("object".into()))
);
assert_eq!(
input_schema.get("required"),
Some(&Value::Array(vec![Value::String("request".into())]))
);
let nested_ref = input_schema["properties"]["request"]["$ref"]
.as_str()
.expect("nested request should be referenced");
let nested_name = nested_ref
.rsplit('/')
.next()
.expect("nested request ref name");
assert!(input_schema["$defs"][nested_name].get("required").is_none());
}
#[test]
fn test_build_request_groups_adjacent_tool_outputs() {
let first_tool = ToolCall::builder()
.id("tool_1")
.name("shell_command")
.args("{\"cmd\":\"pwd\"}")
.build()
.unwrap();
let second_tool = ToolCall::builder()
.id("tool_2")
.name("git")
.args("{\"command\":\"status\"}")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![
ChatMessage::Assistant(None, Some(vec![first_tool.clone(), second_tool.clone()])),
ChatMessage::new_tool_output(
first_tool,
ToolOutput::Text("pwd output".to_string()),
),
ChatMessage::new_tool_output(
second_tool,
ToolOutput::Text("git output".to_string()),
),
])
.build()
.unwrap();
let client = Anthropic::builder().build().unwrap();
let built = client.build_request(&request).unwrap().build().unwrap();
assert_eq!(built.messages.len(), 2);
assert_eq!(built.messages[1].content.len(), 2);
}
}