Richard FeldmanOpus 5 BYOK Support (#61596)640286a3 months agoHistory
1572 lines · 58.1 KB · Rust
RawDownload
use anyhow::{Result, anyhow};use collections::HashMap;use futures::{Stream, StreamExt};use language_model_core::{    CompactedContext, CompactionUpdate, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelProviderId, LanguageModelProviderName, LanguageModelRequest,    LanguageModelRequestToolInput, LanguageModelToolChoice, LanguageModelToolResultContent,    LanguageModelToolUse, LanguageModelToolUseInput, MessageContent, ProviderCompactionState, Role,    SharedString, StopReason, TokenUsage,    util::{fix_streamed_json, parse_tool_arguments},};use std::pin::Pin;use std::str::FromStr;use std::sync::Arc; use crate::{    AdaptiveThinkingDisplay, AnthropicError, AnthropicModelMode, CacheControl, CacheControlType,    CacheTtl, CompactionTrigger, ContentDelta, ContextManagement, ContextManagementEdit, Event,    ImageSource, Message, RequestContent, ResponseContent, StringOrContents, Thinking, Tool,    ToolChoice, ToolResultContent, ToolResultPart, Usage, completion_error_from_anthropic,    completion_error_from_anthropic_api,}; pub const COMPACTION_STATE_FORMAT: &str = "anthropic.messages.encrypted-content.v1"; /// Packages a compaction block's opaque `encrypted_content` into provider/// state owned by `owner`.////// Anthropic requires the metadata to be round-tripped verbatim, and only the/// backend whose infrastructure produced it can make sense of it. The owner/// recorded here is what [`provider_compaction_encrypted_content`] later/// compares against, so it must identify that backend, not merely the wire/// protocol.pub fn provider_compaction_state_from_encrypted_content(    owner: LanguageModelProviderId,    encrypted_content: impl Into<Arc<str>>,) -> ProviderCompactionState {    ProviderCompactionState::new(        owner,        SharedString::new_static(COMPACTION_STATE_FORMAT),        encrypted_content,    )} /// Recovers the `encrypted_content` to round-trip from `state` if it is owned/// by `owner`, or `None` when the state belongs to a different backend and the/// summary should be replayed without it.pub fn provider_compaction_encrypted_content(    state: &ProviderCompactionState,    owner: &LanguageModelProviderId,) -> Result<Option<Arc<str>>> {    if state.provider_id() != owner {        return Ok(None);    }    if state.format() != COMPACTION_STATE_FORMAT {        return Err(anyhow!(            "unsupported Anthropic compaction state format: {}",            state.format()        ));    }    Ok(Some(state.payload().into()))} #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]pub enum AnthropicPromptCacheMode {    Disabled,    Legacy,    #[default]    Automatic,} fn set_cache_control(content: &mut RequestContent, cache_control: Option<CacheControl>) -> bool {    match content {        RequestContent::RedactedThinking { .. } => false,        RequestContent::Text {            cache_control: target,            ..        }        | RequestContent::Thinking {            cache_control: target,            ..        }        | RequestContent::Image {            cache_control: target,            ..        }        | RequestContent::ToolUse {            cache_control: target,            ..        }        | RequestContent::ToolResult {            cache_control: target,            ..        }        | RequestContent::Compaction {            cache_control: target,            ..        } => {            *target = cache_control;            true        }    }} fn mark_last_cacheable_content(content: &mut [RequestContent], cache_control: CacheControl) {    for content in content.iter_mut().rev() {        if set_cache_control(content, Some(cache_control)) {            break;        }    }} fn to_anthropic_content(    content: MessageContent,    compaction_state_owner: &LanguageModelProviderId,) -> Result<Option<RequestContent>> {    match content {        MessageContent::Text(text) => {            let text = if text.chars().last().is_some_and(|c| c.is_whitespace()) {                text.trim_end().to_string()            } else {                text            };            if !text.is_empty() {                Ok(Some(RequestContent::Text {                    text,                    cache_control: None,                }))            } else {                Ok(None)            }        }        MessageContent::Thinking {            text: thinking,            signature,        } => {            if let Some(signature) = signature                && !thinking.is_empty()            {                Ok(Some(RequestContent::Thinking {                    thinking,                    signature,                    cache_control: None,                }))            } else {                Ok(None)            }        }        MessageContent::RedactedThinking(data) => {            if !data.is_empty() {                Ok(Some(RequestContent::RedactedThinking { data }))            } else {                Ok(None)            }        }        MessageContent::Image(image) => Ok(Some(RequestContent::Image {            source: ImageSource {                source_type: "base64".to_string(),                media_type: "image/png".to_string(),                data: image.source.to_string(),            },            cache_control: None,        })),        MessageContent::ToolUse(tool_use) => match tool_use.input {            LanguageModelToolUseInput::Json(input) => Ok(Some(RequestContent::ToolUse {                id: tool_use.id.to_string(),                name: tool_use.name.to_string(),                input,                cache_control: None,            })),            LanguageModelToolUseInput::Text(_) => Err(anyhow::anyhow!(                "Anthropic does not support custom tool calls"            )),        },        MessageContent::ToolResult(tool_result) => {            let content = match tool_result.content.as_slice() {                [LanguageModelToolResultContent::Text(text)] => {                    ToolResultContent::Plain(text.to_string())                }                _ => {                    let parts = tool_result                        .content                        .into_iter()                        .map(|part| match part {                            LanguageModelToolResultContent::Text(text) => ToolResultPart::Text {                                text: text.to_string(),                            },                            LanguageModelToolResultContent::Image(image) => ToolResultPart::Image {                                source: ImageSource {                                    source_type: "base64".to_string(),                                    media_type: "image/png".to_string(),                                    data: image.source.to_string(),                                },                            },                        })                        .collect();                    ToolResultContent::Multipart(parts)                }            };            Ok(Some(RequestContent::ToolResult {                tool_use_id: tool_result.tool_use_id.to_string(),                is_error: tool_result.is_error,                content,                cache_control: None,            }))        }        MessageContent::Compaction(CompactedContext::Summary {            content,            provider_state,        }) => {            let encrypted_content = match &provider_state {                Some(state) => {                    provider_compaction_encrypted_content(state, compaction_state_owner)?                }                None => None,            };            Ok(Some(RequestContent::Compaction {                content: Some(content),                encrypted_content,                cache_control: None,            }))        }        MessageContent::Compaction(CompactedContext::ProviderState(_)) => Ok(None),    }} pub fn into_anthropic(    request: LanguageModelRequest,    model: String,    default_temperature: f32,    max_output_tokens: u64,    mode: AnthropicModelMode,    cache_mode: AnthropicPromptCacheMode,    compaction_state_owner: &LanguageModelProviderId,) -> Result<crate::Request> {    let mut new_messages: Vec<Message> = Vec::new();    let mut system_message = String::new();    let mut any_message_wants_cache = false;     for message in request.messages {        if message.contents_empty() {            continue;        }         any_message_wants_cache |= message.cache;         match message.role {            Role::User | Role::Assistant => {                let mut anthropic_message_content = Vec::new();                for content in message.content {                    if let Some(content) = to_anthropic_content(content, compaction_state_owner)? {                        anthropic_message_content.push(content);                    }                }                let anthropic_role = match message.role {                    Role::User => crate::Role::User,                    Role::Assistant => crate::Role::Assistant,                    Role::System => unreachable!("System role should never occur here"),                };                if anthropic_message_content.is_empty() {                    continue;                }                 if cache_mode == AnthropicPromptCacheMode::Legacy && message.cache {                    mark_last_cacheable_content(                        &mut anthropic_message_content,                        CacheControl {                            cache_type: CacheControlType::Ephemeral,                            ttl: None,                        },                    );                }                 if let Some(last_message) = new_messages.last_mut()                    && last_message.role == anthropic_role                {                    last_message.content.extend(anthropic_message_content);                    continue;                }                 new_messages.push(Message {                    role: anthropic_role,                    content: anthropic_message_content,                });            }            Role::System => {                if !system_message.is_empty() {                    system_message.push_str("\n\n");                }                system_message.push_str(&message.string_contents());            }        }    }     // When caching is enabled, mark the static prefix (tools + system) with an    // explicit long-TTL breakpoint, and let Anthropic's automatic top-level    // cache_control handle the short-TTL conversation breakpoint. Anthropic    // requires that longer TTLs appear earlier in the prefix, and the prefix    // order is tools → system → messages, so long-TTL tools/system before a    // short-TTL conversation breakpoint is a valid mix.    let long_lived_cache = (cache_mode == AnthropicPromptCacheMode::Automatic        && any_message_wants_cache)        .then_some(CacheControl {            cache_type: CacheControlType::Ephemeral,            ttl: Some(CacheTtl::OneHour),        });     let system = if system_message.is_empty() {        None    } else if let Some(cache_control) = long_lived_cache {        Some(StringOrContents::Content(vec![RequestContent::Text {            text: system_message,            cache_control: Some(cache_control),        }]))    } else {        Some(StringOrContents::String(system_message))    };     let mut tools: Vec<Tool> = request        .tools        .into_iter()        .map(|tool| match tool.input {            LanguageModelRequestToolInput::Function {                input_schema,                use_input_streaming,            } => Ok(Tool {                name: tool.name,                description: tool.description,                input_schema,                eager_input_streaming: use_input_streaming,                cache_control: None,            }),            LanguageModelRequestToolInput::Custom { .. } => {                Err(anyhow::anyhow!("Anthropic does not support custom tools"))            }        })        .collect::<Result<_>>()?;    if let Some(cache_control) = long_lived_cache        && let Some(last_tool) = tools.last_mut()    {        last_tool.cache_control = Some(cache_control);    }     let thinking = if request.thinking_allowed {        match mode {            AnthropicModelMode::Thinking { budget_tokens } => {                Some(Thinking::Enabled { budget_tokens })            }            AnthropicModelMode::AdaptiveThinking => Some(Thinking::Adaptive {                display: Some(AdaptiveThinkingDisplay::Summarized),            }),            AnthropicModelMode::Default => None,        }    } else if crate::requires_explicit_thinking_opt_out(&model) {        // On Claude Opus 5, omitting the `thinking` field no longer means        // "off": the model runs adaptive thinking by default, so features        // that suppress thinking (e.g. inline assist) must opt out        // explicitly. `disabled` is only accepted at effort `high` or below;        // that holds here because `output_config` is never sent when thinking        // is disallowed, and the server-side default effort is `high`.        // <https://platform.claude.com/docs/en/about-claude/models/migration-guide#migrating-to-claude-opus-5>        Some(Thinking::Disabled)    } else {        None    };     Ok(crate::Request {        model,        messages: new_messages,        max_tokens: max_output_tokens,        system,        // Opt into Anthropic's automatic prompt caching for the conversation        // tail. Omitting `ttl` uses the default (short) TTL, which refreshes        // for free on every cache hit — ideal for the rapidly-changing        // conversation suffix.        cache_control: (cache_mode == AnthropicPromptCacheMode::Automatic            && any_message_wants_cache)            .then_some(CacheControl {                cache_type: CacheControlType::Ephemeral,                ttl: None,            }),        thinking,        tools,        tool_choice: request.tool_choice.map(|choice| match choice {            LanguageModelToolChoice::Auto => ToolChoice::Auto,            LanguageModelToolChoice::Any => ToolChoice::Any,            LanguageModelToolChoice::None => ToolChoice::None,        }),        metadata: None,        output_config: if request.thinking_allowed            && matches!(mode, AnthropicModelMode::AdaptiveThinking)        {            request.thinking_effort.as_deref().and_then(|effort| {                let effort = match effort {                    "low" => Some(crate::Effort::Low),                    "medium" => Some(crate::Effort::Medium),                    "high" => Some(crate::Effort::High),                    "xhigh" => Some(crate::Effort::XHigh),                    "max" => Some(crate::Effort::Max),                    _ => None,                };                effort.map(|effort| crate::OutputConfig {                    effort: Some(effort),                })            })        } else {            None        },        stop_sequences: Vec::new(),        speed: request.speed.map(Into::into),        temperature: request.temperature.or(Some(default_temperature)),        top_k: None,        top_p: None,        context_management: request.compact_at_tokens.map(|value| ContextManagement {            edits: vec![ContextManagementEdit::Compact {                trigger: Some(CompactionTrigger::InputTokens { value }),            }],        }),    })} pub struct AnthropicEventMapper {    tool_uses_by_index: HashMap<usize, RawToolUse>,    compactions_by_index: HashMap<usize, RawCompaction>,    usage: Usage,    stop_reason: StopReason,    provider_name: LanguageModelProviderName,    compaction_state_owner: LanguageModelProviderId,} impl AnthropicEventMapper {    /// `compaction_state_owner` identifies the backend whose infrastructure    /// produced this stream, so that any `encrypted_content` it emits is only    /// ever round-tripped back to that same backend.    pub fn new(        provider_name: LanguageModelProviderName,        compaction_state_owner: LanguageModelProviderId,    ) -> Self {        Self {            tool_uses_by_index: HashMap::default(),            compactions_by_index: HashMap::default(),            usage: Usage::default(),            stop_reason: StopReason::EndTurn,            provider_name,            compaction_state_owner,        }    }     pub fn map_stream(        mut self,        events: Pin<Box<dyn Send + Stream<Item = Result<Event, AnthropicError>>>>,    ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>    {        events.flat_map(move |event| {            futures::stream::iter(match event {                Ok(event) => self.map_event(event),                Err(error) => vec![Err(completion_error_from_anthropic(                    error,                    self.provider_name.clone(),                ))],            })        })    }     pub fn map_event(        &mut self,        event: Event,    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {        match event {            Event::ContentBlockStart {                index,                content_block,            } => match content_block {                ResponseContent::Text { text } => {                    vec![Ok(LanguageModelCompletionEvent::Text(text))]                }                ResponseContent::Thinking { thinking } => {                    vec![Ok(LanguageModelCompletionEvent::Thinking {                        text: thinking,                        signature: None,                    })]                }                ResponseContent::RedactedThinking { data } => {                    vec![Ok(LanguageModelCompletionEvent::RedactedThinking { data })]                }                ResponseContent::ToolUse { id, name, .. } => {                    self.tool_uses_by_index.insert(                        index,                        RawToolUse {                            id,                            name,                            input_json: String::new(),                        },                    );                    Vec::new()                }                ResponseContent::Compaction {                    content,                    encrypted_content,                } => {

Lines 1–500 of 1572.