gproxy_protocol/adapt/generate/stream/
resource_map.rs

1//! Async resource work over the two concrete Gemini/Responses event types.
2
3use super::{StreamInvocation, bridge::StreamBridge};
4use crate::{
5    adapt::generate::{
6        GenerationResources, GenerationStateAccess,
7        image_resources::{ImageStreamProgress, ResourceSend, ResourceSync},
8    },
9    capability::{ResourceAccess, StateStore},
10    transform::{
11        TransformError,
12        generate::gemini_responses::stream::{GeminiToResponsesStream, ResponsesToGeminiStream},
13    },
14    wire::{gemini as g, openai::responses::stream::StreamEvent},
15};
16
17type Work<'a, T> = crate::capability::CapabilityFuture<'a, Result<T, TransformError>>;
18
19#[cfg(not(target_arch = "wasm32"))]
20pub(super) type ExternalSource<'a, E> =
21    dyn futures_core::Stream<Item = Result<E, TransformError>> + Unpin + Send + 'a;
22
23#[cfg(target_arch = "wasm32")]
24pub(super) type ExternalSource<'a, E> =
25    dyn futures_core::Stream<Item = Result<E, TransformError>> + Unpin + 'a;
26
27pub(super) trait ResourceMapping<B: StreamBridge, S: StateStore>: ResourceSend {
28    fn revision(&self) -> u64;
29    fn failed(&self) -> bool;
30    fn begin_step(&mut self) -> Result<u64, TransformError>;
31    fn source<'a>(&'a mut self, input: &'a B::NativeEvent) -> Work<'a, B::NativeEvent>;
32    fn client<'a>(&'a mut self, input: &'a [B::ClientEvent]) -> Work<'a, Vec<B::ClientEvent>>;
33}
34
35pub(super) struct ReadImages<'a, 'b, R: ResourceAccess> {
36    resources: &'a GenerationResources<'b, R>,
37    progress: &'a mut ImageStreamProgress<R::PublishedHandle>,
38}
39
40impl<R: ResourceAccess + ResourceSync, S: StateStore + ResourceSync>
41    ResourceMapping<GeminiToResponsesStream, S> for ReadImages<'_, '_, R>
42where
43    R::Scope: ResourceSync,
44    R::PublishedHandle: ResourceSend,
45    S::Scope: ResourceSync,
46{
47    fn revision(&self) -> u64 {
48        self.progress.revision
49    }
50    fn failed(&self) -> bool {
51        self.progress.failed()
52    }
53    fn begin_step(&mut self) -> Result<u64, TransformError> {
54        self.progress.committed(())?;
55        Ok(self.progress.revision)
56    }
57    fn source<'a>(
58        &'a mut self,
59        input: &'a g::GenerateContentResponseBody,
60    ) -> Work<'a, g::GenerateContentResponseBody> {
61        Box::pin(async move {
62            let value = self.progress.source(input, self.resources).await?;
63            self.progress.committed(value)
64        })
65    }
66    fn client<'a>(&'a mut self, input: &'a [StreamEvent]) -> Work<'a, Vec<StreamEvent>> {
67        Box::pin(async move { self.progress.committed(input.to_vec()) })
68    }
69}
70
71pub(super) struct PublishImages<'a, 'b, R: ResourceAccess> {
72    pub(super) resources: &'a GenerationResources<'b, R>,
73    pub(super) progress: &'a mut ImageStreamProgress<R::PublishedHandle>,
74    pub(super) expires_at: std::time::SystemTime,
75    pub(super) enabled: bool,
76}
77
78impl<R: ResourceAccess + ResourceSync, S: StateStore + ResourceSync>
79    ResourceMapping<ResponsesToGeminiStream, S> for PublishImages<'_, '_, R>
80where
81    R::Scope: ResourceSync,
82    R::PublishedHandle: ResourceSend,
83    S::Scope: ResourceSync,
84{
85    fn revision(&self) -> u64 {
86        self.progress.revision
87    }
88    fn failed(&self) -> bool {
89        self.progress.failed()
90    }
91    fn begin_step(&mut self) -> Result<u64, TransformError> {
92        self.progress.committed(())?;
93        Ok(self.progress.revision)
94    }
95    fn source<'a>(&'a mut self, input: &'a StreamEvent) -> Work<'a, StreamEvent> {
96        Box::pin(async move { self.progress.committed(input.clone()) })
97    }
98    fn client<'a>(
99        &'a mut self,
100        input: &'a [g::GenerateContentResponseBody],
101    ) -> Work<'a, Vec<g::GenerateContentResponseBody>> {
102        Box::pin(async move {
103            let value = if self.enabled {
104                self.progress
105                    .client(input, self.expires_at, self.resources)
106                    .await?
107            } else {
108                input.to_vec()
109            };
110            self.progress.committed(value)
111        })
112    }
113}
114
115impl StreamInvocation<GeminiToResponsesStream> {
116    pub async fn next_with_image_resources<S, R>(
117        &mut self,
118        state: &GenerationStateAccess<'_, S>,
119        resources: &GenerationResources<'_, R>,
120        progress: &mut ImageStreamProgress<R::PublishedHandle>,
121    ) -> Result<Option<super::StreamChunk<StreamEvent>>, TransformError>
122    where
123        R: ResourceAccess + ResourceSync,
124        R::Scope: ResourceSync,
125        R::PublishedHandle: ResourceSend,
126        S: StateStore + ResourceSync,
127        S::Scope: ResourceSync,
128    {
129        progress.bind(self.selected.identities.response.namespace())?;
130        self.image_resources_required = true;
131        self.next_mapped(
132            state,
133            None,
134            Some(&mut ReadImages {
135                resources,
136                progress,
137            }),
138        )
139        .await
140    }
141}
142
143impl StreamInvocation<ResponsesToGeminiStream> {
144    pub async fn next_with_image_resources<S, R>(
145        &mut self,
146        state: &GenerationStateAccess<'_, S>,
147        resources: &GenerationResources<'_, R>,
148        progress: &mut ImageStreamProgress<R::PublishedHandle>,
149    ) -> Result<Option<super::StreamChunk<g::GenerateContentResponseBody>>, TransformError>
150    where
151        R: ResourceAccess + ResourceSync,
152        R::Scope: ResourceSync,
153        R::PublishedHandle: ResourceSend,
154        S: StateStore + ResourceSync,
155        S::Scope: ResourceSync,
156    {
157        progress.bind(self.selected.identities.response.namespace())?;
158        self.image_resources_required = true;
159        let enabled = crate::adapt::generate::image_resources::wants_uri(&self.original);
160        let mut mapping = PublishImages {
161            resources,
162            progress,
163            expires_at: state.expires_at,
164            enabled,
165        };
166        if self.websocket_terminal {
167            self.next_mapped(
168                state,
169                Some(&mut futures_util::stream::poll_fn(|_| {
170                    std::task::Poll::Ready(None)
171                })),
172                Some(&mut mapping),
173            )
174            .await
175        } else {
176            self.next_mapped(state, None, Some(&mut mapping)).await
177        }
178    }
179}
gproxy_protocol/adapt/generate/stream/resource_map.rs

gproxy_protocol/adapt/generate/stream/
resource_map.rs
