Skip to main content

gproxy_protocol/adapt/generate/fanout/prepare/
chat_responses.rs

1use super::*;
2
3#[derive(Debug)]
4pub struct ChatViaResponsesFanout(Fanout<ChatViaResponses>);
5
6impl ChatViaResponsesFanout {
7    pub async fn prepare<S: StateStore>(
8        input: h::GenerateContentRequestBody,
9        target: FanoutTarget,
10        state: &GenerationStateAccess<'_, S>,
11    ) -> Result<Self, TransformError> {
12        let mut input = input.into_declared();
13        let id = group_id(&target.options)?;
14        let count = input.n.flatten().unwrap_or(1);
15        input.n = Some(Some(1));
16        let mut children = Vec::new();
17        for index in 0..count {
18            let identities = target.options.child_identity(index, crate::Dialect::OpenAi);
19            children.push(
20                ChatViaResponses::prepare_with_state(
21                    input.clone(),
22                    target.endpoint.clone(),
23                    identities,
24                    state,
25                )
26                .await?,
27            );
28        }
29        Ok(Self(Fanout {
30            children,
31            endpoint: target.endpoint,
32            group_id: id,
33            started: false,
34        }))
35    }
36    /// Resources are read once before cloning the single-candidate request.
37    pub async fn prepare_with_capabilities<S: StateStore, R: ResourceAccess>(
38        input: h::GenerateContentRequestBody,
39        target: FanoutTarget,
40        state: &GenerationStateAccess<'_, S>,
41        resources: &GenerationResources<'_, R>,
42    ) -> Result<Self, TransformError> {
43        let input = input.into_declared();
44        Self::prepare(resources.chat(input).await?, target, state).await
45    }
46    pub fn response_id(&self) -> &str {
47        &self.0.group_id
48    }
49    pub fn children(&self) -> &[ChatViaResponses] {
50        &self.0.children
51    }
52    /// Sends each child at most once and aggregates their results in order.
53    pub async fn invoke<U: Upstream, S: StateStore>(
54        &mut self,
55        upstream: &U,
56        target: &U::Target,
57        limits: CodecLimits,
58        state: &GenerationStateAccess<'_, S>,
59        progress: &mut FanoutProgress<r::GenerateContentResponseBody>,
60        facts: impl FnMut(usize, &r::GenerateContentResponseBody) -> Result<(), TransformError>,
61    ) -> Result<Converted<h::GenerateContentResponseBody>, TransformError> {
62        self.0
63            .run((upstream, target), limits, state, progress, facts)
64            .await
65    }
66}