feat(rust): scaffold anthropic stream transformation

litellm-rust/Cargo.lock

@@ -1953,6 +1953,8 @@ dependencies = [
 name = "litellm-core"
 version = "0.1.0"
 dependencies = [
+ "aws-smithy-eventstream",
+ "aws-smithy-types",
  "base64 0.22.1",
  "bytes",
  "data-url",
@@ -1961,6 +1963,7 @@ dependencies = [
  "litellm-auth-aws",
  "litellm-auth-azure",
  "litellm-auth-gcp",
+ "litellm-framing",
  "mime_guess",
  "moka",
  "rand 0.8.7",

litellm-rust/Cargo.toml

@@ -11,6 +11,7 @@ repository = "https://github.com/BerriAI/litellm"
 [workspace.dependencies]
 bytes = "1"
 litellm-core = { path = "crates/core" }
+litellm-framing = { path = "crates/framer" }
 litellm-auth = { path = "crates/auth" }
 litellm-auth-aws = { path = "crates/auth-aws" }
 litellm-auth-azure = { path = "crates/auth-azure" }

litellm-rust/crates/core/Cargo.toml

@@ -15,6 +15,7 @@ litellm-auth.workspace = true
 litellm-auth-aws.workspace = true
 litellm-auth-azure.workspace = true
 litellm-auth-gcp.workspace = true
+litellm-framing.workspace = true
 moka.workspace = true
 mime_guess = "2.0.5"
 rand.workspace = true
@@ -34,4 +35,6 @@ url.workspace = true
 veil.workspace = true
 
 [dev-dependencies]
+aws-smithy-eventstream = "=0.61.1"
+aws-smithy-types = "1.6.1"
 rstest.workspace = true

litellm-rust/crates/core/src/chat_completions/mod.rs

@@ -14,6 +14,7 @@ pub mod conversation;
 pub(crate) mod handler;
 mod prepare;
 pub mod response_utils;
+pub mod streaming;
 pub mod transformation;
 pub mod types;
 

litellm-rust/crates/core/src/chat_completions/streaming.rs

@@ -0,0 +1,9 @@
+pub trait StreamTransformer {
+    type Input;
+    type Output;
+    type Error;
+
+    fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error>;
+
+    fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error>;
+}

litellm-rust/crates/core/src/chat_completions/types.rs

@@ -120,3 +120,83 @@ pub struct ChatCompletionsResponse {
     pub choices: Vec<ChatCompletionsChoice>,
     pub usage: ChatCompletionsUsage,
 }
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ChatCompletionToolCallFunctionChunk {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub name: Option<String>,
+    pub arguments: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub provider_specific_fields: Option<Map<String, Value>>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ChatCompletionToolCallChunk {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub id: Option<String>,
+    #[serde(rename = "type")]
+    pub tool_type: String,
+    pub function: ChatCompletionToolCallFunctionChunk,
+    pub index: i64,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ChatCompletionThinkingBlock {
+    Thinking {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        thinking: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        signature: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        cache_control: Option<Value>,
+    },
+    RedactedThinking {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        data: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        cache_control: Option<Value>,
+    },
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub struct ChatCompletionDelta {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub content: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub role: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<ChatCompletionToolCallChunk>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reasoning_content: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub thinking_blocks: Option<Vec<ChatCompletionThinkingBlock>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub provider_specific_fields: Option<Map<String, Value>>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ChatCompletionStreamingChoice {
+    pub index: u64,
+    pub delta: ChatCompletionDelta,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub finish_reason: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub logprobs: Option<Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ChatCompletionChunk {
+    pub id: String,
+    pub created: u64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub model: Option<String>,
+    pub object: String,
+    pub choices: Vec<ChatCompletionStreamingChoice>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub usage: Option<ChatCompletionsUsage>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub provider_specific_fields: Option<Map<String, Value>>,
+}

litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs

@@ -1 +1,2 @@
+pub mod streaming;
 pub mod transformation;

litellm-rust/crates/core/src/providers/anthropic/chat_completions/streaming.rs

@@ -0,0 +1,164 @@
+use std::collections::HashMap;
+
+use serde_json::Value;
+
+use crate::chat_completions::Error;
+use crate::chat_completions::streaming::StreamTransformer;
+use crate::chat_completions::types::{
+    ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
+    ChatCompletionsUsage,
+};
+use crate::providers::anthropic::messages::streaming::{
+    AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
+    AnthropicStreamUsage,
+};
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum AnthropicJsonChunkType {
+    ValidJson,
+    AccumulatedJson,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum AnthropicContentBlockType {
+    Text,
+    ToolUse,
+    ServerToolUse,
+    Thinking,
+    RedactedThinking,
+    Compaction,
+    ToolResult(String),
+    Other(String),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct AnthropicContentBlockDeltaEvent {
+    pub index: u64,
+    pub delta: AnthropicContentBlockDelta,
+}
+
+pub struct AnthropicChatCompletionsStreamTransformer {
+    pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
+    pub tool_index: i64,
+    pub json_mode: bool,
+    pub speed: Option<String>,
+    pub tool_name_reverse_map: HashMap<String, String>,
+    pub response_id: String,
+    pub served_model: Option<String>,
+    pub is_response_format_tool: bool,
+    pub converted_response_format_tool: bool,
+    pub accumulated_json: String,
+    pub chunk_type: AnthropicJsonChunkType,
+    pub current_content_block_type: Option<AnthropicContentBlockType>,
+    pub web_search_results: Vec<Value>,
+    pub web_search_calls: HashMap<String, Value>,
+    pub compaction_blocks: Vec<Value>,
+    pub reasoning_content_chunks: Vec<String>,
+    pub server_tool_inputs: HashMap<String, Value>,
+    pub tool_results: Vec<Value>,
+    pub current_server_tool_id: Option<String>,
+    pub container_id: Option<String>,
+}
+
+impl AnthropicChatCompletionsStreamTransformer {
+    pub fn new(
+        _json_mode: bool,
+        _speed: Option<String>,
+        _tool_name_reverse_map: HashMap<String, String>,
+    ) -> Self {
+        todo!()
+    }
+
+    pub fn check_empty_tool_call_args(&self) -> bool {
+        todo!()
+    }
+
+    pub fn handle_usage(&mut self, _usage: AnthropicStreamUsage) -> ChatCompletionsUsage {
+        todo!()
+    }
+
+    pub fn handle_content_block_delta(
+        &mut self,
+        _index: u64,
+        _delta: AnthropicContentBlockDelta,
+    ) -> (
+        String,
+        Option<ChatCompletionToolCallChunk>,
+        Vec<ChatCompletionThinkingBlock>,
+        Option<Value>,
+        Option<String>,
+    ) {
+        todo!()
+    }
+
+    pub fn handle_content_block_start(
+        &mut self,
+        _index: u64,
+        _content_block: AnthropicContentBlock,
+    ) -> Result<ChatCompletionChunk, Error> {
+        todo!()
+    }
+
+    pub fn handle_json_mode_chunk(
+        &mut self,
+        _text: String,
+        _tool_use: Option<ChatCompletionToolCallChunk>,
+    ) -> (String, Option<ChatCompletionToolCallChunk>) {
+        todo!()
+    }
+
+    pub fn handle_accumulated_json_chunk(
+        &mut self,
+        _data: &str,
+        _is_final: bool,
+    ) -> Result<Option<ChatCompletionChunk>, Error> {
+        todo!()
+    }
+
+    pub fn handle_redacted_thinking_content(
+        &mut self,
+        _content_block: &AnthropicContentBlock,
+    ) -> Vec<ChatCompletionThinkingBlock> {
+        todo!()
+    }
+
+    pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
+        todo!()
+    }
+
+    pub fn complete_web_search_call(&mut self, _result: Value) {
+        todo!()
+    }
+
+    pub fn build_code_interpreter_results(&self) -> Vec<Value> {
+        todo!()
+    }
+
+    pub fn handle_message_delta(
+        &mut self,
+        _event: AnthropicMessagesStreamEvent,
+    ) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
+        todo!()
+    }
+
+    pub fn chunk_parser(
+        &mut self,
+        _event: AnthropicMessagesStreamEvent,
+    ) -> Result<ChatCompletionChunk, Error> {
+        todo!()
+    }
+}
+
+impl StreamTransformer for AnthropicChatCompletionsStreamTransformer {
+    type Input = AnthropicMessagesStreamEvent;
+    type Output = ChatCompletionChunk;
+    type Error = Error;
+
+    fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
+        todo!()
+    }
+
+    fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
+        todo!()
+    }
+}

litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs

@@ -1 +1,2 @@
+pub mod streaming;
 pub mod transformation;

litellm-rust/crates/core/src/providers/anthropic/messages/streaming.rs

@@ -0,0 +1,256 @@
+use base64::Engine;
+use bytes::Buf;
+use futures_util::{Stream, StreamExt};
+use litellm_framing::Framer;
+use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
+use litellm_framing::sse::{SseFrame, SseFramer};
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+#[derive(Debug, thiserror::Error)]
+pub enum AnthropicStreamDecodeError {
+    #[error("stream framing failed: {0}")]
+    Framing(#[from] litellm_framing::Error),
+    #[error("Anthropic SSE frame has no data")]
+    MissingSseData,
+    #[error("Anthropic stream event is invalid: {0}")]
+    InvalidEvent(#[from] serde_json::Error),
+    #[error("Bedrock event payload has invalid base64: {0}")]
+    InvalidBedrockPayload(#[from] base64::DecodeError),
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicStreamUsage {
+    #[serde(default)]
+    pub input_tokens: u64,
+    #[serde(default)]
+    pub output_tokens: u64,
+    #[serde(default)]
+    pub cache_creation_input_tokens: u64,
+    #[serde(default)]
+    pub cache_read_input_tokens: u64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub server_tool_use: Option<Value>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicStreamMessage {
+    pub id: String,
+    #[serde(rename = "type")]
+    pub message_type: String,
+    pub role: String,
+    pub model: String,
+    pub content: Vec<Value>,
+    pub stop_reason: Option<String>,
+    pub stop_sequence: Option<String>,
+    pub usage: AnthropicStreamUsage,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum AnthropicContentBlockDelta {
+    TextDelta { text: String },
+    InputJsonDelta { partial_json: String },
+    Citations { citation: Value },
+    ThinkingDelta { thinking: String },
+    SignatureDelta { signature: String },
+    CompactionDelta { content: String },
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicContentBlock {
+    #[serde(rename = "type")]
+    pub block_type: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub name: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub text: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub input: Option<Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub thinking: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub signature: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub data: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub content: Option<Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub caller: Option<Value>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicMessageDelta {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub stop_reason: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub stop_sequence: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub stop_details: Option<Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub container: Option<Value>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicStreamError {
+    #[serde(rename = "type")]
+    pub error_type: String,
+    pub message: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub details: Option<Value>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum AnthropicMessagesStreamEvent {
+    MessageStart {
+        message: AnthropicStreamMessage,
+    },
+    ContentBlockStart {
+        index: u64,
+        content_block: AnthropicContentBlock,
+    },
+    ContentBlockDelta {
+        index: u64,
+        delta: AnthropicContentBlockDelta,
+    },
+    ContentBlockStop {
+        index: u64,
+    },
+    MessageDelta {
+        delta: AnthropicMessageDelta,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        usage: Option<AnthropicStreamUsage>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        context_management: Option<Value>,
+    },
+    MessageStop,
+    Ping,
+    Error {
+        error: AnthropicStreamError,
+    },
+}
+
+#[derive(Deserialize)]
+struct BedrockChunkPayload {
+    bytes: String,
+}
+
+pub fn decode_anthropic_sse_frame(
+    frame: SseFrame,
+) -> Result<AnthropicMessagesStreamEvent, AnthropicStreamDecodeError> {
+    let data = frame
+        .data
+        .ok_or(AnthropicStreamDecodeError::MissingSseData)?;
+    Ok(serde_json::from_str(&data)?)
+}
+
+pub fn decode_bedrock_anthropic_frame(
+    frame: AwsEventStreamFrame,
+) -> Result<AnthropicMessagesStreamEvent, AnthropicStreamDecodeError> {
+    let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)?;
+    let event = base64::engine::general_purpose::STANDARD.decode(payload.bytes)?;
+    Ok(serde_json::from_slice(&event)?)
+}
+
+pub fn direct_anthropic_event_stream<S, B, E>(
+    input: S,
+) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, AnthropicStreamDecodeError>> + Send
+where
+    S: Stream<Item = Result<B, E>> + Send,
+    B: Buf + Send,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    SseFramer
+        .frame(input)
+        .map(|frame| decode_anthropic_sse_frame(frame?))
+}
+
+pub fn bedrock_anthropic_event_stream<S, B, E>(
+    input: S,
+) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, AnthropicStreamDecodeError>> + Send
+where
+    S: Stream<Item = Result<B, E>> + Send,
+    B: Buf + Send,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    AwsEventStreamFramer
+        .frame(input)
+        .map(|frame| decode_bedrock_anthropic_frame(frame?))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io;
+
+    use aws_smithy_eventstream::frame::write_message_to;
+    use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
+    use base64::engine::general_purpose::STANDARD;
+    use bytes::Bytes;
+    use futures_util::TryStreamExt;
+
+    use super::*;
+
+    const TEXT_DELTA: &str =
+        r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
+
+    #[tokio::test]
+    async fn direct_anthropic_sse_frames_into_typed_events() {
+        let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
+        let events = direct_anthropic_event_stream(futures_util::stream::iter(
+            wire.as_bytes().chunks(3).map(Ok::<_, io::Error>),
+        ))
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+
+        assert_eq!(
+            events,
+            vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
+                index: 0,
+                delta: AnthropicContentBlockDelta::TextDelta {
+                    text: "hello".into(),
+                },
+            }]
+        );
+    }
+
+    #[tokio::test]
+    async fn bedrock_aws_frames_into_the_same_typed_events() {
+        let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)});
+        let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(
+            Header::new(":event-type", HeaderValue::String("chunk".into())),
+        );
+        let mut wire = Vec::new();
+        write_message_to(&message, &mut wire).unwrap();
+
+        let events = bedrock_anthropic_event_stream(futures_util::stream::iter(
+            wire.chunks(3).map(Ok::<_, io::Error>),
+        ))
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+
+        assert_eq!(
+            events,
+            vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
+                index: 0,
+                delta: AnthropicContentBlockDelta::TextDelta {
+                    text: "hello".into(),
+                },
+            }]
+        );
+    }
+}