systemprompt-api 0.47.0

//! Decodes Anthropic SSE frames into canonical events.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use futures_util::StreamExt;
// JSON: protocol boundary — event shapes are owned by the models::wire
// Anthropic codec.
use serde_json::Value;
use systemprompt_models::wire::anthropic::AnthropicStreamState;

use super::super::super::canonical_response::CanonicalEvent;

#[cfg_attr(
    not(feature = "test-api"),
    expect(
        unreachable_pub,
        reason = "items are re-exported via `test_api` only when the feature is on"
    )
)]
pub fn sse_to_canonical_events<S>(
    stream: S,
) -> futures_util::stream::BoxStream<'static, Result<CanonicalEvent, String>>
where
    S: futures_util::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send + 'static,
{
    use futures_util::stream;
    let s = stream
        .map(|chunk| chunk.map_err(|e| e.to_string()))
        .scan(
            (Vec::<u8>::new(), AnthropicStreamState::default()),
            |state, item| {
                let (buf, codec) = state;
                let res = match item {
                    Ok(bytes) => {
                        buf.extend_from_slice(&bytes);
                        Some(drain_frames(buf, codec))
                    },
                    Err(e) => Some(vec![Err(e)]),
                };
                futures_util::future::ready(res)
            },
        )
        .flat_map(stream::iter);
    s.boxed()
}

pub(in crate::services::gateway) fn raw_sse_stream<S>(
    stream: S,
) -> futures_util::stream::BoxStream<'static, Result<bytes::Bytes, String>>
where
    S: futures_util::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send + 'static,
{
    stream.map(|chunk| chunk.map_err(|e| e.to_string())).boxed()
}

#[derive(Debug, Default)]
pub(in crate::services::gateway) struct SseDecoder {
    buf: Vec<u8>,
    codec: AnthropicStreamState,
}

impl SseDecoder {
    pub(in crate::services::gateway) fn push(&mut self, chunk: &[u8]) -> Vec<CanonicalEvent> {
        self.buf.extend_from_slice(chunk);
        drain_frames(&mut self.buf, &mut self.codec)
            .into_iter()
            .flatten()
            .collect()
    }
}

fn drain_frames(
    buf: &mut Vec<u8>,
    codec: &mut AnthropicStreamState,
) -> Vec<Result<CanonicalEvent, String>> {
    let mut events: Vec<Result<CanonicalEvent, String>> = Vec::new();
    while let Some(end) = systemprompt_models::wire::sse::frame_end(buf) {
        let frame: Vec<u8> = buf.drain(..end).collect();
        let frame_str = String::from_utf8_lossy(&frame);
        for line in frame_str.lines() {
            if let Some(data) = line.strip_prefix("data: ") {
                if data.trim() == "[DONE]" {
                    continue;
                }
                if let Ok(value) = serde_json::from_str::<Value>(data) {
                    events.extend(codec.events_from_sse(&value).into_iter().map(Ok));
                }
            }
        }
    }
    events
}