Skip to main content

gproxy_protocol/adapt/generate/stream/fanout/
prepare.rs

1use super::super::{GeminiViaClaudeStreamFacts, StreamTarget};
2use super::*;
3use crate::{
4    adapt::generate::{
5        GenerationResources,
6        chat_claude::ChatViaClaude,
7        chat_responses::ChatViaResponses,
8        claude_gemini::GeminiViaClaude,
9        fanout::{
10            ChatViaClaudeFanout, ChatViaResponsesFanout, FanoutTarget, GeminiViaClaudeFanout,
11            GeminiViaResponsesFanout,
12        },
13        gemini_responses::GeminiViaResponses,
14    },
15    capability::ResourceAccess,
16    transform::generate::{
17        chat_responses::stream as hr, claude_chat::stream as ch, claude_gemini::stream as cg,
18        gemini_responses::stream as gr,
19    },
20    wire::{DeclaredFields, gemini as g, openai::chat as h},
21};
22
23impl ChatViaClaudeFanout {
24    pub async fn prepare_stream<S: StateStore>(
25        input: h::GenerateContentRequestBody,
26        target: FanoutTarget,
27        mut context: impl FnMut(usize) -> ch::ClaudeToChatContext,
28        settings: StreamSettings,
29        state: &GenerationStateAccess<'_, S>,
30    ) -> Result<FanoutStream<ch::ClaudeToChatStream>, TransformError> {
31        let input = input.into_declared();
32
33        let mut children = Vec::new();
34        for index in 0..input.n.flatten().unwrap_or(1) {
35            let identities = target.options.child_identity(index, crate::Dialect::Claude);
36            let context = context(index as usize);
37            let mut request = input.clone();
38            request.n = Some(Some(1));
39            children.push(
40                ChatViaClaude::prepare_stream(
41                    request,
42                    StreamTarget {
43                        endpoint: target.endpoint.clone(),
44                        identities,
45                    },
46                    context,
47                    settings,
48                    state,
49                )
50                .await?,
51            );
52        }
53        FanoutStream::new(children, target.options, state).await
54    }
55    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
56        input: h::GenerateContentRequestBody,
57        target: FanoutTarget,
58        mut context: impl FnMut(usize) -> ch::ClaudeToChatContext,
59        settings: StreamSettings,
60        state: &GenerationStateAccess<'_, S>,
61        resources: &GenerationResources<'_, R>,
62    ) -> Result<FanoutStream<ch::ClaudeToChatStream>, TransformError> {
63        let input = input.into_declared();
64
65        let mut children = Vec::new();
66        for index in 0..input.n.flatten().unwrap_or(1) {
67            let identities = target.options.child_identity(index, crate::Dialect::Claude);
68            let context = context(index as usize);
69            let mut request = input.clone();
70            request.n = Some(Some(1));
71            children.push(
72                ChatViaClaude::prepare_stream_with_capabilities(
73                    request,
74                    StreamTarget {
75                        endpoint: target.endpoint.clone(),
76                        identities,
77                    },
78                    context,
79                    settings,
80                    state,
81                    resources,
82                )
83                .await?,
84            );
85        }
86        FanoutStream::new(children, target.options, state).await
87    }
88}
89
90impl ChatViaResponsesFanout {
91    pub async fn prepare_stream<S: StateStore>(
92        input: h::GenerateContentRequestBody,
93        target: FanoutTarget,
94        settings: StreamSettings,
95        state: &GenerationStateAccess<'_, S>,
96    ) -> Result<FanoutStream<hr::ResponsesToChatStream>, TransformError> {
97        let input = input.into_declared();
98
99        let mut children = Vec::new();
100        for index in 0..input.n.flatten().unwrap_or(1) {
101            let identities = target.options.child_identity(index, crate::Dialect::OpenAi);
102            let mut request = input.clone();
103            request.n = Some(Some(1));
104            children.push(
105                ChatViaResponses::prepare_stream(
106                    request,
107                    StreamTarget {
108                        endpoint: target.endpoint.clone(),
109                        identities,
110                    },
111                    settings,
112                    state,
113                )
114                .await?,
115            );
116        }
117        FanoutStream::new(children, target.options, state).await
118    }
119    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
120        input: h::GenerateContentRequestBody,
121        target: FanoutTarget,
122        settings: StreamSettings,
123        state: &GenerationStateAccess<'_, S>,
124        resources: &GenerationResources<'_, R>,
125    ) -> Result<FanoutStream<hr::ResponsesToChatStream>, TransformError> {
126        let input = input.into_declared();
127
128        let mut children = Vec::new();
129        for index in 0..input.n.flatten().unwrap_or(1) {
130            let identities = target.options.child_identity(index, crate::Dialect::OpenAi);
131            let mut request = input.clone();
132            request.n = Some(Some(1));
133            children.push(
134                ChatViaResponses::prepare_stream_with_capabilities(
135                    request,
136                    StreamTarget {
137                        endpoint: target.endpoint.clone(),
138                        identities,
139                    },
140                    settings,
141                    state,
142                    resources,
143                )
144                .await?,
145            );
146        }
147        FanoutStream::new(children, target.options, state).await
148    }
149}
150
151impl GeminiViaClaudeFanout {
152    pub async fn prepare_stream<S: StateStore>(
153        input: g::GenerateContentRequestBody,
154        target: FanoutTarget,
155        mut context: impl FnMut(usize) -> GeminiViaClaudeStreamFacts,
156        settings: StreamSettings,
157        state: &GenerationStateAccess<'_, S>,
158    ) -> Result<FanoutStream<cg::ClaudeToGeminiStream>, TransformError> {
159        let input = input.into_declared();
160
161        let mut children = Vec::new();
162        for index in 0..input
163            .generation_config
164            .as_ref()
165            .and_then(|v| v.candidate_count)
166            .unwrap_or(1)
167        {
168            let identities = target.options.child_identity(index, crate::Dialect::Claude);
169            let context = context(index as usize);
170            let mut request = input.clone();
171            request
172                .generation_config
173                .get_or_insert_with(|| g::GenerationConfig::builder().build())
174                .candidate_count = Some(1);
175            children.push(
176                GeminiViaClaude::prepare_stream(
177                    request,
178                    StreamTarget {
179                        endpoint: target.endpoint.clone(),
180                        identities,
181                    },
182                    context,
183                    settings,
184                    state,
185                )
186                .await?,
187            );
188        }
189        FanoutStream::new(children, target.options, state).await
190    }
191    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
192        input: g::GenerateContentRequestBody,
193        target: FanoutTarget,
194        mut context: impl FnMut(usize) -> GeminiViaClaudeStreamFacts,
195        settings: StreamSettings,
196        state: &GenerationStateAccess<'_, S>,
197        resources: &GenerationResources<'_, R>,
198    ) -> Result<FanoutStream<cg::ClaudeToGeminiStream>, TransformError> {
199        let input = input.into_declared();
200
201        let mut children = Vec::new();
202        for index in 0..input
203            .generation_config
204            .as_ref()
205            .and_then(|v| v.candidate_count)
206            .unwrap_or(1)
207        {
208            let identities = target.options.child_identity(index, crate::Dialect::Claude);
209            let context = context(index as usize);
210            let mut request = input.clone();
211            request
212                .generation_config
213                .get_or_insert_with(|| g::GenerationConfig::builder().build())
214                .candidate_count = Some(1);
215            children.push(
216                GeminiViaClaude::prepare_stream_with_capabilities(
217                    request,
218                    StreamTarget {
219                        endpoint: target.endpoint.clone(),
220                        identities,
221                    },
222                    context,
223                    settings,
224                    state,
225                    resources,
226                )
227                .await?,
228            );
229        }
230        FanoutStream::new(children, target.options, state).await
231    }
232}
233
234impl GeminiViaResponsesFanout {
235    pub async fn prepare_stream<S: StateStore>(
236        input: g::GenerateContentRequestBody,
237        target: FanoutTarget,
238        mut context: impl FnMut(usize) -> gr::ResponsesToGeminiContext,
239        settings: StreamSettings,
240        state: &GenerationStateAccess<'_, S>,
241    ) -> Result<FanoutStream<gr::ResponsesToGeminiStream>, TransformError> {
242        let input = input.into_declared();
243
244        let mut children = Vec::new();
245        for index in 0..input
246            .generation_config
247            .as_ref()
248            .and_then(|v| v.candidate_count)
249            .unwrap_or(1)
250        {
251            let identities = target.options.child_identity(index, crate::Dialect::OpenAi);
252            let context = context(index as usize);
253            let mut request = input.clone();
254            request
255                .generation_config
256                .get_or_insert_with(|| g::GenerationConfig::builder().build())
257                .candidate_count = Some(1);
258            children.push(
259                GeminiViaResponses::prepare_stream(
260                    request,
261                    StreamTarget {
262                        endpoint: target.endpoint.clone(),
263                        identities,
264                    },
265                    context,
266                    settings,
267                    state,
268                )
269                .await?,
270            );
271        }
272        FanoutStream::new(children, target.options, state).await
273    }
274    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
275        input: g::GenerateContentRequestBody,
276        target: FanoutTarget,
277        mut context: impl FnMut(usize) -> gr::ResponsesToGeminiContext,
278        settings: StreamSettings,
279        state: &GenerationStateAccess<'_, S>,
280        resources: &GenerationResources<'_, R>,
281    ) -> Result<FanoutStream<gr::ResponsesToGeminiStream>, TransformError> {
282        let input = input.into_declared();
283
284        let mut children = Vec::new();
285        for index in 0..input
286            .generation_config
287            .as_ref()
288            .and_then(|v| v.candidate_count)
289            .unwrap_or(1)
290        {
291            let identities = target.options.child_identity(index, crate::Dialect::OpenAi);
292            let context = context(index as usize);
293            let mut request = input.clone();
294            request
295                .generation_config
296                .get_or_insert_with(|| g::GenerationConfig::builder().build())
297                .candidate_count = Some(1);
298            children.push(
299                GeminiViaResponses::prepare_stream_with_capabilities(
300                    request,
301                    StreamTarget {
302                        endpoint: target.endpoint.clone(),
303                        identities,
304                    },
305                    context,
306                    settings,
307                    state,
308                    resources,
309                )
310                .await?,
311            );
312        }
313        FanoutStream::new(children, target.options, state).await
314    }
315}