gproxy-transform 3.0.0-alpha.0

GPROXY v3 pairwise protocol transforms
Documentation
mod framing;

use bytes::Bytes;
use gproxy_protocol::{OperationKey, StreamFraming};

use self::framing::{FrameDecoder, FrameEncoder};
use super::SseFrame;
use crate::TransformError;
use crate::registry::{self, TransformPair};

pub struct ResponseStream {
    decoder: FrameDecoder,
    converter: Box<dyn Converter>,
    encoder: FrameEncoder,
}

pub(crate) trait Converter: Send {
    fn frame(&mut self, frame: SseFrame) -> Result<Vec<Bytes>, TransformError>;
    fn finish(&mut self) -> Result<Vec<Bytes>, TransformError>;
}

impl ResponseStream {
    pub fn new(source: OperationKey, target: OperationKey) -> Result<Self, TransformError> {
        Self::new_framed(source, target, StreamFraming::Sse, StreamFraming::Sse)
    }

    pub fn new_framed(
        source: OperationKey,
        target: OperationKey,
        source_framing: StreamFraming,
        target_framing: StreamFraming,
    ) -> Result<Self, TransformError> {
        let converter: Box<dyn Converter> = if source == target {
            Box::new(Passthrough)
        } else {
            let pair =
                registry::resolve(source, target).ok_or(TransformError::UnsupportedPair {
                    source_key: source,
                    target_key: target,
                })?;
            converter(pair, source, target)?
        };
        Ok(Self {
            decoder: FrameDecoder::new(target_framing)?,
            converter,
            encoder: FrameEncoder::new(source_framing)?,
        })
    }

    pub fn push(&mut self, chunk: Bytes) -> Result<Vec<Bytes>, TransformError> {
        let frames = self.decoder.push(&chunk)?;
        self.convert(frames)
    }

    pub fn finish(&mut self) -> Result<Vec<Bytes>, TransformError> {
        let frames = self.decoder.finish()?;
        let mut output = self.convert(frames)?;
        let terminal = self.converter.finish()?;
        output.extend(self.encoder.push(terminal)?);
        output.extend(self.encoder.finish()?);
        Ok(output)
    }

    fn convert(&mut self, frames: Vec<SseFrame>) -> Result<Vec<Bytes>, TransformError> {
        let mut output = Vec::new();
        for frame in frames {
            output.extend(self.encoder.push(self.converter.frame(frame)?)?);
        }
        Ok(output)
    }
}

struct Passthrough;

impl Converter for Passthrough {
    fn frame(&mut self, frame: SseFrame) -> Result<Vec<Bytes>, TransformError> {
        Ok(vec![SseFrame::encode(frame._event.as_deref(), &frame.data)])
    }

    fn finish(&mut self) -> Result<Vec<Bytes>, TransformError> {
        Ok(Vec::new())
    }
}

fn converter(
    pair: TransformPair,
    source: OperationKey,
    target: OperationKey,
) -> Result<Box<dyn Converter>, TransformError> {
    Ok(match pair {
        TransformPair::ChatToClaude => {
            crate::generate_content::openai_chat_to_claude_messages::stream::converter()
        }
        TransformPair::ResponsesToClaude => {
            crate::generate_content::openai_responses_to_claude_messages::stream::converter()
        }
        TransformPair::ClaudeToChat => {
            crate::generate_content::claude_messages_to_openai_chat::stream::converter()
        }
        TransformPair::ClaudeToResponses => {
            crate::generate_content::claude_messages_to_openai_responses::stream::converter()
        }
        TransformPair::ClaudeToGemini => {
            crate::generate_content::gemini_generate_content_to_claude_messages::stream::converter()
        }
        TransformPair::GeminiToClaude => {
            crate::generate_content::claude_messages_to_gemini_generate_content::stream::converter()
        }
        TransformPair::GeminiToChat => {
            crate::generate_content::gemini_generate_content_to_openai_chat::stream::converter()
        }
        TransformPair::ChatToGemini => {
            crate::generate_content::openai_chat_to_gemini_generate_content::stream::converter()
        }
        TransformPair::GeminiToResponses => {
            crate::generate_content::gemini_generate_content_to_openai_responses::stream::converter(
            )
        }
        TransformPair::ResponsesToGemini => {
            crate::generate_content::openai_responses_to_gemini_generate_content::stream::converter(
            )
        }
        TransformPair::OpenAiChatToResponses => {
            crate::generate_content::openai_chat_to_openai_responses::stream::converter()
        }
        TransformPair::OpenAiResponsesToChat => {
            crate::generate_content::openai_responses_to_openai_chat::stream::converter()
        }
        TransformPair::OpenAiCreateImageToGemini => crate::images::stream::from_gemini(false),
        TransformPair::OpenAiEditImageToGemini => crate::images::stream::from_gemini(true),
        TransformPair::OpenAiCreateImageToResponses => crate::images::stream::from_responses(false),
        TransformPair::OpenAiEditImageToResponses => crate::images::stream::from_responses(true),
        _ => {
            return Err(TransformError::UnsupportedPair {
                source_key: source,
                target_key: target,
            });
        }
    })
}