gproxy_protocol/adapt/generate/stream/prepare/
chat_gemini.rs

1use std::future::Future;
2
3use super::super::super::{
4    Endpoint, GenerationIdentity, GenerationResources, GenerationStateAccess,
5    chat_gemini::{ChatViaGemini, GeminiViaChat},
6};
7use super::super::{StreamInvocation, StreamSettings, StreamTarget};
8use super::{RequestMode, start};
9use crate::{
10    capability::{ResourceAccess, StateStore},
11    transform::{TransformError, generate::gemini_chat::stream as p},
12    wire::{DeclaredFields, gemini as g, openai::chat as h},
13};
14
15pub struct ChatViaGeminiStreamFacts {
16    pub function_names: std::collections::BTreeMap<String, String>,
17    pub response: p::GeminiToChatContext,
18}
19
20impl ChatViaGemini {
21    pub async fn prepare_stream<S: StateStore>(
22        input: h::GenerateContentRequestBody,
23        target: StreamTarget,
24        context: ChatViaGeminiStreamFacts,
25        settings: StreamSettings,
26        state: &GenerationStateAccess<'_, S>,
27    ) -> Result<StreamInvocation<p::GeminiToChatStream>, TransformError> {
28        let ChatViaGeminiStreamFacts {
29            function_names,
30            response,
31        } = context;
32        Self::stream(
33            input,
34            target,
35            response,
36            settings,
37            state,
38            |input, endpoint, ids| {
39                Self::prepare_with_state(input, endpoint, ids, state, &function_names)
40            },
41        )
42        .await
43    }
44    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
45        input: h::GenerateContentRequestBody,
46        target: StreamTarget,
47        context: ChatViaGeminiStreamFacts,
48        settings: StreamSettings,
49        state: &GenerationStateAccess<'_, S>,
50        resources: &GenerationResources<'_, R>,
51    ) -> Result<StreamInvocation<p::GeminiToChatStream>, TransformError> {
52        let ChatViaGeminiStreamFacts {
53            function_names,
54            response,
55        } = context;
56        Self::stream(
57            input,
58            target,
59            response,
60            settings,
61            state,
62            |input, endpoint, ids| {
63                Self::prepare_with_capabilities(
64                    input,
65                    endpoint,
66                    ids,
67                    state,
68                    resources,
69                    &function_names,
70                )
71            },
72        )
73        .await
74    }
75    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
76        input: h::GenerateContentRequestBody,
77        target: StreamTarget,
78        response: p::GeminiToChatContext,
79        settings: StreamSettings,
80        state: &GenerationStateAccess<'_, S>,
81        prepare: impl FnOnce(h::GenerateContentRequestBody, Endpoint, GenerationIdentity) -> Fut,
82    ) -> Result<StreamInvocation<p::GeminiToChatStream>, TransformError> {
83        let original = input.into_declared();
84        let prepared = prepare(
85            original.clone().buffered(),
86            target.endpoint.clone(),
87            target.identities.clone(),
88        )
89        .await?;
90        start(
91            original,
92            prepared,
93            target,
94            settings,
95            state,
96            |_, target, settings| {
97                p::GeminiToChatStream::new_with_policy(
98                    response,
99                    target.identities.response.clone(),
100                    settings.events.into(),
101                    target.identities.response_policy.clone(),
102                )
103            },
104        )
105        .await
106    }
107}
108
109impl GeminiViaChat {
110    pub async fn prepare_stream<S: StateStore>(
111        input: g::GenerateContentRequestBody,
112        target: StreamTarget,
113        settings: StreamSettings,
114        state: &GenerationStateAccess<'_, S>,
115    ) -> Result<StreamInvocation<p::ChatToGeminiStream>, TransformError> {
116        Self::stream(input, target, settings, state, |input, endpoint, ids| {
117            Self::prepare_with_state(input, endpoint, ids, state)
118        })
119        .await
120    }
121    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
122        input: g::GenerateContentRequestBody,
123        target: StreamTarget,
124        settings: StreamSettings,
125        state: &GenerationStateAccess<'_, S>,
126        resources: &GenerationResources<'_, R>,
127    ) -> Result<StreamInvocation<p::ChatToGeminiStream>, TransformError> {
128        Self::stream(input, target, settings, state, |input, endpoint, ids| {
129            Self::prepare_with_capabilities(input, endpoint, ids, state, resources)
130        })
131        .await
132    }
133    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
134        input: g::GenerateContentRequestBody,
135        target: StreamTarget,
136        settings: StreamSettings,
137        state: &GenerationStateAccess<'_, S>,
138        prepare: impl FnOnce(g::GenerateContentRequestBody, Endpoint, GenerationIdentity) -> Fut,
139    ) -> Result<StreamInvocation<p::ChatToGeminiStream>, TransformError> {
140        let original = input.into_declared();
141        let prepared = prepare(
142            original.clone().buffered(),
143            target.endpoint.clone(),
144            target.identities.clone(),
145        )
146        .await?;
147        start(
148            original,
149            prepared,
150            target,
151            settings,
152            state,
153            |_, target, settings| {
154                p::ChatToGeminiStream::new_with_policy(
155                    target.identities.response.clone(),
156                    settings.events.into(),
157                    target.identities.response_policy.clone(),
158                )
159            },
160        )
161        .await
162    }
163}
gproxy_protocol/adapt/generate/stream/prepare/chat_gemini.rs

gproxy_protocol/adapt/generate/stream/prepare/
chat_gemini.rs
