gproxy_protocol/adapt/generate/stream/websocket/
image_resources.rs

1use super::*;
2use crate::{
3    adapt::generate::{
4        GenerationResources,
5        image_resources::{ImageStreamProgress, ResourceSend, ResourceSync},
6    },
7    capability::ResourceAccess,
8    transform::generate::gemini_responses::stream::ResponsesToGeminiStream,
9    wire::gemini as g,
10};
11
12impl GenerationWsTurn<'_, ResponsesToGeminiStream> {
13    /// Perform actual URI publications while retaining the native WS turn and
14    /// pending client events across cancellation of this individual future.
15    pub async fn next_with_image_resources<S, R>(
16        &mut self,
17        state: &GenerationStateAccess<'_, S>,
18        resources: &GenerationResources<'_, R>,
19        progress: &mut ImageStreamProgress<R::PublishedHandle>,
20    ) -> Result<Option<StreamChunk<g::GenerateContentResponseBody>>, TransformError>
21    where
22        R: ResourceAccess + ResourceSync,
23        R::Scope: ResourceSync,
24        R::PublishedHandle: ResourceSend,
25        S: StateStore + ResourceSync,
26        S::Scope: ResourceSync,
27    {
28        let Some(turn) = &mut self.turn else {
29            return Err(super::super::invalid("WebSocket generation turn failed"));
30        };
31        progress.bind(self.invocation.selected.identities.response.namespace())?;
32        self.invocation.image_resources_required = true;
33        let enabled = crate::adapt::generate::image_resources::wants_uri(&self.invocation.original);
34        let mut mapping = super::super::resource_map::PublishImages {
35            resources,
36            progress,
37            expires_at: state.expires_at,
38            enabled,
39        };
40        let result = self
41            .invocation
42            .next_mapped(state, Some(turn), Some(&mut mapping))
43            .await;
44        if result.is_err() && self.invocation.failed {
45            self.failure = turn.failure_event().cloned();
46            self.websocket_failure = turn.websocket_failure().cloned();
47            self.turn.take();
48        }
49        result
50    }
51}
gproxy_protocol/adapt/generate/stream/websocket/image_resources.rs

gproxy_protocol/adapt/generate/stream/websocket/
image_resources.rs
