Aether/apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs

use serde_json::Value;

use crate::ai_serving::{
    maybe_build_ai_surface_stream_rewriter, AiSurfaceFinalizeError, AiSurfaceStreamRewriter,
    ResponseHistoryRecord,
};
use crate::GatewayError;

pub(crate) struct LocalStreamRewriter<'a> {
    inner: AiSurfaceStreamRewriter<'a>,
}

pub(crate) fn maybe_build_local_stream_rewriter<'a>(
    report_context: Option<&'a Value>,
) -> Option<LocalStreamRewriter<'a>> {
    maybe_build_ai_surface_stream_rewriter(report_context)
        .map(|inner| LocalStreamRewriter { inner })
}

impl LocalStreamRewriter<'_> {
    pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
        self.inner.push_chunk(chunk).map_err(map_surface_error)
    }

    pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
        self.inner.finish().map_err(map_surface_error)
    }

    pub(crate) fn take_response_history_record(&mut self) -> Option<ResponseHistoryRecord> {
        self.inner.take_response_history_record()
    }
}

fn map_surface_error(error: AiSurfaceFinalizeError) -> GatewayError {
    error.into()
}

#[cfg(test)]
#[path = "../tests_stream.rs"]
mod tests;