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

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

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