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