gproxy_protocol/adapt/generate/stream/fanout/
images.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
12struct Images<'a, 'r, R: ResourceAccess> {
13    resources: &'a GenerationResources<'r, R>,
14    progress: &'a mut [ImageStreamProgress<R::PublishedHandle>],
15}
16
17impl<R: ResourceAccess + ResourceSync, S: StateStore + ResourceSync>
18    super::driver::ChildNext<ResponsesToGeminiStream, S> for Images<'_, '_, R>
19where
20    R::Scope: ResourceSync,
21    R::PublishedHandle: ResourceSend,
22    S::Scope: ResourceSync,
23{
24    async fn next(
25        &mut self,
26        child: &mut StreamInvocation<ResponsesToGeminiStream>,
27        index: usize,
28        state: &GenerationStateAccess<'_, S>,
29    ) -> Result<Option<StreamChunk<g::GenerateContentResponseBody>>, TransformError> {
30        let mut used_bytes = 0u64;
31        let mut used_count = 0usize;
32        for (other, progress) in self.progress.iter().enumerate() {
33            if other == index {
34                continue;
35            }
36            let (count, bytes) = progress.publication_budget()?;
37            used_bytes = used_bytes
38                .checked_add(bytes)
39                .ok_or_else(|| limit("fanout publication byte overflow"))?;
40            used_count = used_count
41                .checked_add(count)
42                .ok_or_else(|| limit("fanout publication count overflow"))?;
43        }
44        let mut limits = self.resources.limits;
45        limits.max_body_bytes = limits
46            .max_body_bytes
47            .checked_sub(used_bytes)
48            .ok_or_else(|| limit("fanout publication byte budget exceeded"))?;
49        let max_references = self
50            .resources
51            .max_references
52            .checked_sub(used_count)
53            .ok_or_else(|| limit("fanout publication count budget exceeded"))?;
54        let resources = GenerationResources {
55            access: self.resources.access,
56            scope: self.resources.scope,
57            limits,
58            max_references,
59            now: self.resources.now,
60        };
61        child
62            .next_with_image_resources(state, &resources, &mut self.progress[index])
63            .await
64    }
65}
66
67impl FanoutStream<ResponsesToGeminiStream> {
68    /// Keep one progress entry per child. Publication operation IDs and receipts
69    /// survive cancellation, and resource budgets apply to the whole group.
70    pub async fn next_with_image_resources<U: crate::capability::Upstream, S, R>(
71        &mut self,
72        upstream: &U,
73        target: &U::Target,
74        state: &GenerationStateAccess<'_, S>,
75        resources: &GenerationResources<'_, R>,
76        progress: &mut [ImageStreamProgress<R::PublishedHandle>],
77    ) -> Result<Option<StreamChunk<g::GenerateContentResponseBody>>, TransformError>
78    where
79        R: ResourceAccess + ResourceSync,
80        R::Scope: ResourceSync,
81        R::PublishedHandle: ResourceSend,
82        S: StateStore + ResourceSync,
83        S::Scope: ResourceSync,
84    {
85        if progress.len() != self.children.len() {
86            return Err(TransformError::shape(
87                "fanout.image_progress",
88                "one retained image progress entry per child required",
89            ));
90        }
91        for (child, image_progress) in self.children.iter().zip(progress.iter_mut()) {
92            if image_progress.revision != child.resource_revision {
93                return Err(conflict(
94                    "fanout image progress changed for an active or completed child",
95                ));
96            }
97            image_progress.bind(child.selected.identities.response.namespace())?;
98        }
99        self.next_using(
100            upstream,
101            target,
102            state,
103            &mut Images {
104                resources,
105                progress,
106            },
107        )
108        .await
109    }
110    pub async fn collect_with_image_resources<U: crate::capability::Upstream, S, R>(
111        &mut self,
112        upstream: &U,
113        target: &U::Target,
114        state: &GenerationStateAccess<'_, S>,
115        resources: &GenerationResources<'_, R>,
116        progress: &mut [ImageStreamProgress<R::PublishedHandle>],
117    ) -> Result<crate::transform::Converted<g::GenerateContentResponseBody>, TransformError>
118    where
119        R: ResourceAccess + ResourceSync,
120        R::Scope: ResourceSync,
121        R::PublishedHandle: ResourceSend,
122        S: StateStore + ResourceSync,
123        S::Scope: ResourceSync,
124    {
125        while self
126            .next_with_image_resources(upstream, target, state, resources, progress)
127            .await?
128            .is_some()
129        {}
130        Ok(crate::transform::Converted {
131            value: self
132                .client_result()
133                .ok_or_else(|| invalid("aggregate not acknowledged"))?
134                .clone(),
135            report: self.report.clone(),
136        })
137    }
138}
gproxy_protocol/adapt/generate/stream/fanout/images.rs

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