Skip to main content

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

1use super::super::super::{
2    GenerationResources, GenerationStateAccess,
3    gemini_responses::{GeminiViaResponses, ResponsesViaGemini},
4};
5use super::super::{StreamInvocation, StreamSettings, StreamTarget};
6use super::RequestMode;
7use crate::{
8    capability::{ResourceAccess, StateStore},
9    transform::{
10        TransformError,
11        generate::{gemini_responses as pair, gemini_responses::stream as p},
12    },
13    wire::{DeclaredFields, gemini as g, openai::responses as r},
14};
15
16pub struct ResponsesViaGeminiStreamFacts {
17    pub request: pair::GeminiReplayContext,
18    pub response: p::GeminiToResponsesContext,
19}
20
21impl GeminiViaResponses {
22    pub async fn prepare_stream<S: StateStore>(
23        input: g::GenerateContentRequestBody,
24        mut target: StreamTarget,
25        mut context: p::ResponsesToGeminiContext,
26        settings: StreamSettings,
27        state: &GenerationStateAccess<'_, S>,
28    ) -> Result<StreamInvocation<p::ResponsesToGeminiStream>, TransformError> {
29        let original = input.into_declared();
30        context.response_modalities = original
31            .generation_config
32            .as_ref()
33            .and_then(|v| v.response_modalities.clone());
34        context.image_mime = original
35            .generation_config
36            .as_ref()
37            .and_then(|v| v.response_format.as_ref())
38            .and_then(|v| v.image.as_ref())
39            .and_then(|v| v.mime_type.clone());
40        let prepared = Self::prepare_with_state(
41            original.clone().buffered(),
42            target.endpoint.clone(),
43            target.identities.clone(),
44            state,
45        )
46        .await?;
47        target.identities = prepared.identities().clone();
48        let bridge = p::ResponsesToGeminiStream::new_with_policy(
49            context,
50            target.identities.response.clone(),
51            target.identities.response_policy.clone(),
52            settings.events.into(),
53        )?;
54        let needs_images = crate::adapt::generate::image_resources::wants_uri(&original);
55        let mut invocation = StreamInvocation::new(
56            original,
57            prepared.target_request().clone().streaming(),
58            target,
59            settings,
60            bridge,
61            prepared.report().clone(),
62            state,
63        )
64        .await?;
65        invocation.image_resources_required = needs_images;
66        Ok(invocation)
67    }
68    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
69        input: g::GenerateContentRequestBody,
70        mut target: StreamTarget,
71        mut context: p::ResponsesToGeminiContext,
72        settings: StreamSettings,
73        state: &GenerationStateAccess<'_, S>,
74        resources: &GenerationResources<'_, R>,
75    ) -> Result<StreamInvocation<p::ResponsesToGeminiStream>, TransformError> {
76        let original = input.into_declared();
77        context.response_modalities = original
78            .generation_config
79            .as_ref()
80            .and_then(|v| v.response_modalities.clone());
81        context.image_mime = original
82            .generation_config
83            .as_ref()
84            .and_then(|v| v.response_format.as_ref())
85            .and_then(|v| v.image.as_ref())
86            .and_then(|v| v.mime_type.clone());
87        let prepared = Self::prepare_with_capabilities(
88            original.clone().buffered(),
89            target.endpoint.clone(),
90            target.identities.clone(),
91            state,
92            resources,
93        )
94        .await?;
95        target.identities = prepared.identities().clone();
96        let bridge = p::ResponsesToGeminiStream::new_with_policy(
97            context,
98            target.identities.response.clone(),
99            target.identities.response_policy.clone(),
100            settings.events.into(),
101        )?;
102        let needs_images = crate::adapt::generate::image_resources::wants_uri(&original);
103        let mut invocation = StreamInvocation::new(
104            original,
105            prepared.target_request().clone().streaming(),
106            target,
107            settings,
108            bridge,
109            prepared.report().clone(),
110            state,
111        )
112        .await?;
113        invocation.image_resources_required = needs_images;
114        Ok(invocation)
115    }
116}
117
118impl ResponsesViaGemini {
119    pub async fn prepare_stream<S: StateStore>(
120        input: r::GenerateContentRequestBody,
121        target: StreamTarget,
122        context: ResponsesViaGeminiStreamFacts,
123        settings: StreamSettings,
124        state: &GenerationStateAccess<'_, S>,
125    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
126        Self::prepare_stream_inner(input, target, context, settings, state, None).await
127    }
128    pub async fn prepare_stream_with_history_cache<S: StateStore>(
129        input: r::GenerateContentRequestBody,
130        target: StreamTarget,
131        context: ResponsesViaGeminiStreamFacts,
132        settings: StreamSettings,
133        state: &GenerationStateAccess<'_, S>,
134        history_cache: &super::super::ResponsesHistoryCache,
135    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
136        Self::prepare_stream_inner(input, target, context, settings, state, Some(history_cache))
137            .await
138    }
139    async fn prepare_stream_inner<S: StateStore>(
140        input: r::GenerateContentRequestBody,
141        mut target: StreamTarget,
142        mut context: ResponsesViaGeminiStreamFacts,
143        settings: StreamSettings,
144        state: &GenerationStateAccess<'_, S>,
145        history_cache: Option<&super::super::ResponsesHistoryCache>,
146    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
147        let original = input.into_declared();
148        let (history, expanded) = super::super::history::History::prepare_with_cache(
149            &original,
150            state,
151            settings.codec,
152            history_cache,
153        )
154        .await?;
155        context.response.response.request = original.clone();
156        context.response.response.request.input = expanded.input.clone();
157        if context.response.actual_model.is_none() {
158            context.response.actual_model = Some(state.target.model.clone());
159        }
160        let prepared = Self::prepare_with_state(
161            expanded.buffered(),
162            target.endpoint.clone(),
163            target.identities.clone(),
164            state,
165            context.request,
166        )
167        .await?;
168        target.identities = prepared.identities().clone();
169        history.claim_response_id(&mut target.identities, crate::Dialect::Gemini)?;
170        let bridge = p::GeminiToResponsesStream::new_with_policy(
171            context.response,
172            target.identities.response.clone(),
173            target.identities.response_policy.clone(),
174            settings.events.into(),
175        )?;
176        let mut invocation = StreamInvocation::new(
177            original,
178            prepared.target_request().clone().streaming(),
179            target,
180            settings,
181            bridge,
182            prepared.report().clone(),
183            state,
184        )
185        .await?;
186        history.bind(&invocation.binding, state)?;
187        invocation.history = Some(history);
188        Ok(invocation)
189    }
190    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
191        input: r::GenerateContentRequestBody,
192        target: StreamTarget,
193        context: ResponsesViaGeminiStreamFacts,
194        settings: StreamSettings,
195        state: &GenerationStateAccess<'_, S>,
196        resources: &GenerationResources<'_, R>,
197    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
198        Self::prepare_stream_with_capabilities_inner(
199            input, target, context, settings, state, resources, None,
200        )
201        .await
202    }
203    pub async fn prepare_stream_with_capabilities_and_history_cache<
204        S: StateStore,
205        R: ResourceAccess,
206    >(
207        input: r::GenerateContentRequestBody,
208        target: StreamTarget,
209        context: ResponsesViaGeminiStreamFacts,
210        settings: StreamSettings,
211        state: &GenerationStateAccess<'_, S>,
212        resources: &GenerationResources<'_, R>,
213        history_cache: &super::super::ResponsesHistoryCache,
214    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
215        Self::prepare_stream_with_capabilities_inner(
216            input,
217            target,
218            context,
219            settings,
220            state,
221            resources,
222            Some(history_cache),
223        )
224        .await
225    }
226    async fn prepare_stream_with_capabilities_inner<S: StateStore, R: ResourceAccess>(
227        input: r::GenerateContentRequestBody,
228        mut target: StreamTarget,
229        mut context: ResponsesViaGeminiStreamFacts,
230        settings: StreamSettings,
231        state: &GenerationStateAccess<'_, S>,
232        resources: &GenerationResources<'_, R>,
233        history_cache: Option<&super::super::ResponsesHistoryCache>,
234    ) -> Result<StreamInvocation<p::GeminiToResponsesStream>, TransformError> {
235        let original = input.into_declared();
236        let (history, expanded) = super::super::history::History::prepare_with_cache(
237            &original,
238            state,
239            settings.codec,
240            history_cache,
241        )
242        .await?;
243        context.response.response.request = original.clone();
244        context.response.response.request.input = expanded.input.clone();
245        if context.response.actual_model.is_none() {
246            context.response.actual_model = Some(state.target.model.clone());
247        }
248        let prepared = Self::prepare_with_capabilities(
249            expanded.buffered(),
250            target.endpoint.clone(),
251            target.identities.clone(),
252            state,
253            resources,
254            context.request,
255        )
256        .await?;
257        target.identities = prepared.identities().clone();
258        history.claim_response_id(&mut target.identities, crate::Dialect::Gemini)?;
259        let bridge = p::GeminiToResponsesStream::new_with_policy(
260            context.response,
261            target.identities.response.clone(),
262            target.identities.response_policy.clone(),
263            settings.events.into(),
264        )?;
265        let mut invocation = StreamInvocation::new(
266            original,
267            prepared.target_request().clone().streaming(),
268            target,
269            settings,
270            bridge,
271            prepared.report().clone(),
272            state,
273        )
274        .await?;
275        history.bind(&invocation.binding, state)?;
276        invocation.history = Some(history);
277        Ok(invocation)
278    }
279}