use futures::Stream;
use reqwest::Response;
use serde_json::Value;
use std::pin::Pin;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use crate::config::StreamingConfig;
use crate::error::{ClientError, Result};
#[derive(Debug, Clone)]
pub enum StreamEvent {
Content(String),
Role(String),
Start,
Finish(Option<String>),
Usage {
input_tokens: Option<u32>,
output_tokens: Option<u32>,
total_tokens: Option<u32>,
},
Error(String),
Raw(String),
}
fn extract_u32(value: &Value, path: &[&str]) -> Option<u32> {
let mut current = value;
for key in path {
current = current.get(*key)?;
}
current.as_u64().map(|t| t as u32)
}
fn extract_str<'a>(value: &'a Value, path: &[&str]) -> Option<&'a str> {
let mut current = value;
for key in path {
current = current.get(*key)?;
}
current.as_str()
}
fn parse_sse_data(data: &str) -> Option<Value> {
serde_json::from_str(data).ok()
}
fn extract_usage(value: &Value, input_path: &[&str], output_path: &[&str], total_path: Option<&[&str]>) -> StreamEvent {
StreamEvent::Usage {
input_tokens: extract_u32(value, input_path),
output_tokens: extract_u32(value, output_path),
total_tokens: total_path.and_then(|p| extract_u32(value, p)),
}
}
fn handle_anthropic_event(event_type: &str, data: &Value) -> Option<StreamEvent> {
match event_type {
"message_start" => Some(StreamEvent::Start),
"content_block_delta" => {
extract_str(data, &["delta", "text"])
.map(|text| StreamEvent::Content(text.to_string()))
}
"message_delta" => {
if data.get("delta").and_then(|d| d.get("usage")).is_some() {
Some(extract_usage(
data,
&["delta", "usage", "input_tokens"],
&["delta", "usage", "output_tokens"],
None,
))
} else {
None
}
}
"message_stop" => Some(StreamEvent::Finish(None)),
"error" => {
let error_msg = extract_str(data, &["error", "message"])
.unwrap_or("Unknown error");
Some(StreamEvent::Error(error_msg.to_string()))
}
_ => Some(StreamEvent::Raw(serde_json::to_string(data).unwrap_or_default())),
}
}
fn handle_openai_event(data: &Value) -> Option<StreamEvent> {
if let Some(content) = extract_str(data, &["choices", "0", "delta", "content"]) {
return Some(StreamEvent::Content(content.to_string()));
}
if let Some(content) = data.get("choices")
.and_then(|c| c.get(0))
.and_then(|choice| choice.get("delta"))
.and_then(|delta| delta.get("content"))
.and_then(|c| c.as_str())
{
return Some(StreamEvent::Content(content.to_string()));
}
if let Some(role) = data.get("choices")
.and_then(|c| c.get(0))
.and_then(|choice| choice.get("delta"))
.and_then(|delta| delta.get("role"))
.and_then(|r| r.as_str())
{
return Some(StreamEvent::Role(role.to_string()));
}
if let Some(finish_reason) = data.get("choices")
.and_then(|c| c.get(0))
.and_then(|choice| choice.get("finish_reason"))
.and_then(|r| r.as_str())
{
return Some(StreamEvent::Finish(Some(finish_reason.to_string())));
}
if data.get("usage").is_some() {
return Some(extract_usage(
data,
&["usage", "prompt_tokens"],
&["usage", "completion_tokens"],
Some(&["usage", "total_tokens"]),
));
}
None
}
fn parse_anthropic_sse(line: &str) -> Option<StreamEvent> {
if line.starts_with("event: ") {
return None;
}
let data = line.strip_prefix("data: ")?;
let json = parse_sse_data(data)?;
let event_type = extract_str(&json, &["type"])?;
handle_anthropic_event(event_type, &json)
}
fn parse_openai_sse(config: &StreamingConfig, line: &str) -> Option<StreamEvent> {
if let Some(prefix) = &config.line_prefix {
if !line.starts_with(prefix) {
return None;
}
}
let data = config.line_prefix.as_ref()
.map(|p| &line[p.len()..])
.unwrap_or(line);
if let Some(done_marker) = &config.done_marker {
if data.trim() == done_marker {
return Some(StreamEvent::Finish(None));
}
}
let json = parse_sse_data(data)?;
handle_openai_event(&json)
}
fn parse_generic_sse(line: &str) -> Option<StreamEvent> {
if !line.is_empty() {
Some(StreamEvent::Raw(line.to_string()))
} else {
None
}
}
pub struct StreamProcessor {
config: StreamingConfig,
}
impl StreamProcessor {
pub fn new(config: &StreamingConfig) -> Result<Self> {
Ok(Self {
config: config.clone(),
})
}
pub async fn process_response(&self, response: Response) -> Result<super::Stream> {
use crate::config::StreamingFormat;
match self.config.format.as_ref() {
Some(StreamingFormat::TextEventStream) => self.process_sse_response(response).await,
Some(StreamingFormat::Ndjson) => Err(ClientError::Stream(
"NDJSON streaming not yet implemented".to_string()
)),
Some(StreamingFormat::Custom(format)) => Err(ClientError::Stream(format!(
"Unsupported streaming format: {}",
format
))),
None => Err(ClientError::Stream(
"Streaming format not configured".to_string()
)),
}
}
async fn process_sse_response(&self, response: Response) -> Result<super::Stream> {
let (tx, rx) = mpsc::channel(100);
let config = self.config.clone();
tokio::spawn(async move {
use futures::StreamExt as FuturesStreamExt;
let mut stream = response.bytes_stream();
let mut buffer = String::new();
while let Some(chunk_result) = FuturesStreamExt::next(&mut stream).await {
match chunk_result {
Ok(chunk) => {
buffer.push_str(&String::from_utf8_lossy(&chunk));
while let Some(newline_pos) = buffer.find('\n') {
let line = buffer[..newline_pos].trim().to_string();
buffer = buffer[newline_pos + 1..].to_string();
if let Some(event) = Self::parse_sse_line(&config, &line) {
if tx.send(Ok(event)).await.is_err() {
return; }
}
}
}
Err(e) => {
let _ = tx.send(Err(ClientError::Http(e))).await;
return;
}
}
}
if !buffer.trim().is_empty() {
if let Some(event) = Self::parse_sse_line(&config, buffer.trim()) {
let _ = tx.send(Ok(event)).await;
}
}
});
Ok(Box::pin(ReceiverStream::new(rx)))
}
fn parse_sse_line(config: &StreamingConfig, line: &str) -> Option<StreamEvent> {
use crate::config::SseParser;
match config.parser.as_ref() {
Some(SseParser::AnthropicSse) => parse_anthropic_sse(line),
Some(SseParser::OpenAiSse) => parse_openai_sse(config, line),
Some(SseParser::GoogleSse) => parse_generic_sse(line),
Some(SseParser::Custom(_)) => parse_generic_sse(line),
None => None, }
}
}
pub trait StreamEventExt {
fn collect_content(self) -> impl std::future::Future<Output = Result<String>> + Send;
fn content_only(self) -> Pin<Box<dyn Stream<Item = Result<String>> + Send>>;
fn parse_json<T>(self) -> Pin<Box<dyn Stream<Item = super::StreamItem<T>> + Send>>
where
T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static;
}
impl<S> StreamEventExt for S
where
S: Stream<Item = Result<StreamEvent>> + Send + 'static,
{
async fn collect_content(self) -> Result<String> {
use futures::StreamExt as FuturesStreamExt;
let mut content = String::new();
let mut stream = Box::pin(self);
while let Some(event_result) = FuturesStreamExt::next(&mut stream).await {
match event_result? {
StreamEvent::Content(text) => content.push_str(&text),
StreamEvent::Error(error) => return Err(ClientError::Stream(error)),
_ => {}
}
}
Ok(content)
}
fn content_only(self) -> Pin<Box<dyn Stream<Item = Result<String>> + Send>> {
use futures::StreamExt as FuturesStreamExt;
let stream = FuturesStreamExt::filter_map(self, |event_result| async move {
match event_result {
Ok(StreamEvent::Content(text)) => Some(Ok(text)),
Ok(_) => None,
Err(e) => Some(Err(e)),
}
});
Box::pin(stream)
}
fn parse_json<T>(self) -> Pin<Box<dyn Stream<Item = super::StreamItem<T>> + Send>>
where
T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
{
use super::parsers::JsonStreamProcessor;
Box::pin(async_stream::stream! {
let mut processor = JsonStreamProcessor::<T>::new();
let mut stream = Box::pin(self);
while let Some(event_result) = futures::StreamExt::next(&mut stream).await {
match event_result {
Ok(StreamEvent::Content(chunk)) => {
for item in processor.process_chunk(&chunk) {
yield item;
}
}
Ok(StreamEvent::Error(e)) => {
yield super::StreamItem::Text(super::TextContent {
text: format!("[Error: {}]", e),
});
}
Ok(_) => {}
Err(_) => break,
}
}
for item in processor.finalize() {
yield item;
}
})
}
}