byokey-proxy 2.0.0

//! Generic SSE stream tapping for token usage extraction.

use std::sync::Arc;

use byokey_types::ByokError;
use byokey_types::traits::ByteStream;
use futures_util::{StreamExt as _, stream::try_unfold};
use serde_json::Value;

use crate::UsageRecorder;

/// Implemented per-provider to extract (`input_tokens`, `output_tokens`) from SSE data lines.
pub(crate) trait UsageParser: Send + 'static {
    fn parse_line(&mut self, data: &Value);
    fn finish(self) -> (u64, u64);
}

/// Wraps a [`ByteStream`], scanning each SSE `data:` line through `parser`
/// and recording usage via [`UsageRecorder`] when the stream ends.
/// All bytes are forwarded unchanged.
pub(crate) fn tap_usage_stream<P: UsageParser>(
    inner: ByteStream,
    usage: Arc<UsageRecorder>,
    model: String,
    provider: String,
    account_id: String,
    parser: P,
) -> ByteStream {
    struct State<P> {
        inner: ByteStream,
        buf: Vec<u8>,
        usage: Arc<UsageRecorder>,
        model: String,
        provider: String,
        account_id: String,
        parser: P,
    }

    Box::pin(try_unfold(
        State {
            inner,
            buf: Vec::new(),
            usage,
            model,
            provider,
            account_id,
            parser,
        },
        |mut s| async move {
            match s.inner.next().await {
                Some(Ok(bytes)) => {
                    s.buf.extend_from_slice(&bytes);
                    while let Some(nl) = s.buf.iter().position(|&b| b == b'\n') {
                        let line: Vec<u8> = s.buf.drain(..=nl).collect();
                        parse_usage_sse_line(&mut s.parser, &line);
                    }
                    Ok(Some((bytes, s)))
                }
                Some(Err(e)) => {
                    tracing::error!(
                        model = %s.model,
                        provider = %s.provider,
                        account_id = %s.account_id,
                        error = %e,
                        "tap_usage_stream: upstream SSE stream yielded error"
                    );
                    s.usage
                        .record_failure_for(&s.model, &s.provider, &s.account_id);
                    Err(e)
                }
                None => {
                    if !s.buf.is_empty() {
                        let line = std::mem::take(&mut s.buf);
                        parse_usage_sse_line(&mut s.parser, &line);
                    }
                    let (input, output) = s.parser.finish();
                    s.usage
                        .record_success_for(&s.model, &s.provider, &s.account_id, input, output);
                    Ok(None)
                }
            }
        },
    ))
}

fn parse_usage_sse_line<P: UsageParser>(parser: &mut P, line: &[u8]) {
    let line = String::from_utf8_lossy(line);
    let line = line.trim();
    if let Some(data) = line.strip_prefix("data:").map(str::trim_start)
        && data != "[DONE]"
        && let Ok(ev) = serde_json::from_str::<Value>(data)
    {
        parser.parse_line(&ev);
    }
}

/// Converts a `wreq::Response` into a [`ByteStream`].
pub(crate) fn response_to_stream(resp: wreq::Response) -> ByteStream {
    Box::pin(resp.bytes_stream().map(|r| {
        r.map_err(|e| {
            tracing::error!(error = %e, "response_to_stream: wreq byte stream error");
            ByokError::from(e)
        })
    }))
}

// ── Parser implementations ──────────────────────────────────────────

pub(crate) struct OpenAIParser {
    input: u64,
    output: u64,
}

impl OpenAIParser {
    pub(crate) fn new() -> Self {
        Self {
            input: 0,
            output: 0,
        }
    }
}

impl UsageParser for OpenAIParser {
    fn parse_line(&mut self, ev: &Value) {
        if let Some(usage) = ev.get("usage") {
            if let Some(v) = usage.get("prompt_tokens").and_then(Value::as_u64) {
                self.input = v;
            }
            if let Some(v) = usage.get("completion_tokens").and_then(Value::as_u64) {
                self.output = v;
            }
        }
    }
    fn finish(self) -> (u64, u64) {
        (self.input, self.output)
    }
}

pub(crate) struct AnthropicParser {
    input: u64,
    output: u64,
}

impl AnthropicParser {
    pub(crate) fn new() -> Self {
        Self {
            input: 0,
            output: 0,
        }
    }
}

impl UsageParser for AnthropicParser {
    fn parse_line(&mut self, ev: &Value) {
        match ev.get("type").and_then(Value::as_str) {
            Some("message_start") => {
                if let Some(v) = ev
                    .pointer("/message/usage/input_tokens")
                    .and_then(Value::as_u64)
                {
                    self.input = v;
                }
            }
            Some("message_delta") => {
                if let Some(v) = ev.pointer("/usage/output_tokens").and_then(Value::as_u64) {
                    self.output = v;
                }
            }
            _ => {}
        }
    }
    fn finish(self) -> (u64, u64) {
        (self.input, self.output)
    }
}

pub(crate) struct CodexParser {
    input: u64,
    output: u64,
}

impl CodexParser {
    pub(crate) fn new() -> Self {
        Self {
            input: 0,
            output: 0,
        }
    }
}

impl UsageParser for CodexParser {
    fn parse_line(&mut self, ev: &Value) {
        if ev.get("type").and_then(Value::as_str) == Some("response.completed") {
            if let Some(v) = ev
                .pointer("/response/usage/input_tokens")
                .and_then(Value::as_u64)
            {
                self.input = v;
            }
            if let Some(v) = ev
                .pointer("/response/usage/output_tokens")
                .and_then(Value::as_u64)
            {
                self.output = v;
            }
        }
    }
    fn finish(self) -> (u64, u64) {
        (self.input, self.output)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use futures_util::stream;

    #[tokio::test]
    async fn tap_usage_stream_parses_final_line_without_newline() {
        let usage = Arc::new(UsageRecorder::new(None));
        let inner: ByteStream = Box::pin(stream::iter([Ok(Bytes::from_static(
            br#"data: {"usage":{"prompt_tokens":12,"completion_tokens":7}}"#,
        ))]));

        let chunks: Vec<_> = tap_usage_stream(
            inner,
            Arc::clone(&usage),
            "gpt-test".to_owned(),
            "openai".to_owned(),
            "default".to_owned(),
            OpenAIParser::new(),
        )
        .collect()
        .await;

        assert_eq!(chunks.len(), 1);
        assert!(chunks[0].is_ok());
        let snapshot = usage.snapshot();
        assert_eq!(snapshot.success_requests, 1);
        assert_eq!(snapshot.input_tokens, 12);
        assert_eq!(snapshot.output_tokens, 7);
    }
}