gproxy_transform/envelope/collector/
mod.rs

1use bytes::Bytes;
2use gproxy_protocol::{ContentGenerationKind, claude as claude_wire, openai};
3
4use self::chat::ChatCollector;
5use self::claude::ClaudeCollector;
6use self::gemini::GeminiCollector;
7use self::responses::ResponsesCollector;
8
9mod chat;
10mod claude;
11mod gemini;
12mod responses;
13
14use super::{SseDecoder, SseFrame};
15use crate::TransformError;
16
17pub enum BufferedResponse {
18    OpenAiChat(Box<openai::ChatCompletionResponse>),
19    OpenAiResponses(Box<openai::ResponseObject>),
20    Claude(Box<claude_wire::CreateMessageResponseBody>),
21    Gemini(Box<gproxy_protocol::gemini::GenerateContentResponse>),
22}
23
24impl BufferedResponse {
25    pub fn into_bytes(self) -> Result<Bytes, TransformError> {
26        Ok(Bytes::from(match self {
27            Self::OpenAiChat(response) => serde_json::to_vec(&response)?,
28            Self::OpenAiResponses(response) => serde_json::to_vec(&response)?,
29            Self::Claude(response) => serde_json::to_vec(&response)?,
30            Self::Gemini(response) => serde_json::to_vec(&response)?,
31        }))
32    }
33}
34
35pub struct ResponseCollector {
36    decoder: SseDecoder,
37    state: Collector,
38}
39
40enum Collector {
41    Chat(Box<ChatCollector>),
42    Responses(Box<ResponsesCollector>),
43    Claude(Box<ClaudeCollector>),
44    Gemini(Box<GeminiCollector>),
45}
46
47impl ResponseCollector {
48    pub fn new(kind: ContentGenerationKind) -> Result<Self, TransformError> {
49        let state = match kind {
50            ContentGenerationKind::OpenAiChat => Collector::Chat(Box::default()),
51            ContentGenerationKind::OpenAiResponses
52            | ContentGenerationKind::OpenAiResponsesWebSocket => {
53                Collector::Responses(Box::default())
54            }
55            ContentGenerationKind::ClaudeMessages => Collector::Claude(Box::default()),
56            ContentGenerationKind::GeminiGenerateContent => Collector::Gemini(Box::default()),
57            #[cfg(not(feature = "exhaustive"))]
58            _ => {
59                return Err(crate::TransformError::unsupported(
60                    "protocol enum",
61                    "unrecognized external variant",
62                ));
63            }
64        };
65        Ok(Self {
66            decoder: SseDecoder::default(),
67            state,
68        })
69    }
70
71    pub fn push(&mut self, chunk: Bytes) -> Result<(), TransformError> {
72        for frame in self.decoder.push(&chunk)? {
73            self.state.frame(frame)?;
74        }
75        Ok(())
76    }
77
78    pub fn is_complete(&self) -> bool {
79        self.state.is_complete()
80    }
81
82    pub fn finish(mut self) -> Result<BufferedResponse, TransformError> {
83        if let Some(frame) = self.decoder.finish()? {
84            self.state.frame(frame)?;
85        }
86        self.state.finish()
87    }
88}
89
90impl Collector {
91    fn frame(&mut self, frame: SseFrame) -> Result<(), TransformError> {
92        match self {
93            Self::Chat(state) => state.frame(frame),
94            Self::Responses(state) => state.frame(frame),
95            Self::Claude(state) => state.frame(frame),
96            Self::Gemini(state) => state.frame(frame),
97        }
98    }
99
100    fn is_complete(&self) -> bool {
101        match self {
102            Self::Chat(state) => state.is_complete(),
103            Self::Responses(state) => state.response.is_some(),
104            Self::Claude(state) => state.complete,
105            Self::Gemini(state) => state.is_complete(),
106        }
107    }
108
109    fn finish(self) -> Result<BufferedResponse, TransformError> {
110        match self {
111            Self::Chat(state) => state
112                .finish()
113                .map(Box::new)
114                .map(BufferedResponse::OpenAiChat),
115            Self::Responses(state) => state
116                .finish()
117                .map(Box::new)
118                .map(BufferedResponse::OpenAiResponses),
119            Self::Claude(state) => state.finish().map(Box::new).map(BufferedResponse::Claude),
120            Self::Gemini(state) => state.finish().map(Box::new).map(BufferedResponse::Gemini),
121        }
122    }
123}
gproxy_transform/envelope/collector/mod.rs

gproxy_transform/envelope/collector/
mod.rs
