Skip to main content

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

1use std::future::Future;
2
3use super::super::super::{
4    Endpoint, GenerationIdentity, GenerationResources, GenerationStateAccess,
5    claude_gemini::{ClaudeViaGemini, GeminiViaClaude},
6};
7use super::super::{StreamInvocation, StreamSettings, StreamTarget};
8use super::{RequestMode, start};
9use crate::{
10    capability::{ResourceAccess, StateStore},
11    transform::{
12        TransformError,
13        generate::{claude_gemini as pair, claude_gemini::stream as p},
14    },
15    wire::{DeclaredFields, claude::generate_content as c, gemini as g},
16};
17
18pub struct ClaudeViaGeminiStreamFacts {
19    pub request: pair::ClaudeGeminiRequestContext,
20    pub response: p::GeminiToClaudeContext,
21}
22
23pub struct GeminiViaClaudeStreamFacts {
24    pub max_tokens: Option<i64>,
25    pub response: p::ClaudeToGeminiContext,
26}
27
28impl ClaudeViaGemini {
29    pub async fn prepare_stream<S: StateStore>(
30        input: c::GenerateContentRequestBody,
31        target: StreamTarget,
32        context: ClaudeViaGeminiStreamFacts,
33        settings: StreamSettings,
34        state: &GenerationStateAccess<'_, S>,
35    ) -> Result<StreamInvocation<p::GeminiToClaudeStream>, TransformError> {
36        Self::stream(
37            input,
38            target,
39            context,
40            settings,
41            state,
42            |input, endpoint, ids, request| {
43                Self::prepare_with_state(input, endpoint, ids, state, request)
44            },
45        )
46        .await
47    }
48    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
49        input: c::GenerateContentRequestBody,
50        target: StreamTarget,
51        context: ClaudeViaGeminiStreamFacts,
52        settings: StreamSettings,
53        state: &GenerationStateAccess<'_, S>,
54        resources: &GenerationResources<'_, R>,
55    ) -> Result<StreamInvocation<p::GeminiToClaudeStream>, TransformError> {
56        Self::stream(
57            input,
58            target,
59            context,
60            settings,
61            state,
62            |input, endpoint, ids, request| {
63                Self::prepare_with_capabilities(input, endpoint, ids, state, resources, request)
64            },
65        )
66        .await
67    }
68    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
69        input: c::GenerateContentRequestBody,
70        target: StreamTarget,
71        context: ClaudeViaGeminiStreamFacts,
72        settings: StreamSettings,
73        state: &GenerationStateAccess<'_, S>,
74        prepare: impl FnOnce(
75            c::GenerateContentRequestBody,
76            Endpoint,
77            GenerationIdentity,
78            pair::ClaudeGeminiRequestContext,
79        ) -> Fut,
80    ) -> Result<StreamInvocation<p::GeminiToClaudeStream>, TransformError> {
81        let original = input.into_declared();
82        let prepared = prepare(
83            original.clone().buffered(),
84            target.endpoint.clone(),
85            target.identities.clone(),
86            context.request,
87        )
88        .await?;
89        start(
90            original,
91            prepared,
92            target,
93            settings,
94            state,
95            |_, target, settings| {
96                p::GeminiToClaudeStream::new_with_policy(
97                    context.response,
98                    target.identities.response.clone(),
99                    settings.events.into(),
100                    target.identities.response_policy.clone(),
101                )
102            },
103        )
104        .await
105    }
106}
107
108impl GeminiViaClaude {
109    pub async fn prepare_stream<S: StateStore>(
110        input: g::GenerateContentRequestBody,
111        target: StreamTarget,
112        context: GeminiViaClaudeStreamFacts,
113        settings: StreamSettings,
114        state: &GenerationStateAccess<'_, S>,
115    ) -> Result<StreamInvocation<p::ClaudeToGeminiStream>, TransformError> {
116        Self::stream(
117            input,
118            target,
119            context,
120            settings,
121            state,
122            |input, endpoint, ids, max| Self::prepare_with_state(input, endpoint, ids, state, max),
123        )
124        .await
125    }
126    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
127        input: g::GenerateContentRequestBody,
128        target: StreamTarget,
129        context: GeminiViaClaudeStreamFacts,
130        settings: StreamSettings,
131        state: &GenerationStateAccess<'_, S>,
132        resources: &GenerationResources<'_, R>,
133    ) -> Result<StreamInvocation<p::ClaudeToGeminiStream>, TransformError> {
134        Self::stream(
135            input,
136            target,
137            context,
138            settings,
139            state,
140            |input, endpoint, ids, max| {
141                Self::prepare_with_capabilities(input, endpoint, ids, state, resources, max)
142            },
143        )
144        .await
145    }
146    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
147        input: g::GenerateContentRequestBody,
148        target: StreamTarget,
149        mut context: GeminiViaClaudeStreamFacts,
150        settings: StreamSettings,
151        state: &GenerationStateAccess<'_, S>,
152        prepare: impl FnOnce(
153            g::GenerateContentRequestBody,
154            Endpoint,
155            GenerationIdentity,
156            Option<i64>,
157        ) -> Fut,
158    ) -> Result<StreamInvocation<p::ClaudeToGeminiStream>, TransformError> {
159        let original = input.into_declared();
160        let prepared = prepare(
161            original.clone().buffered(),
162            target.endpoint.clone(),
163            target.identities.clone(),
164            context.max_tokens,
165        )
166        .await?;
167        start(
168            original,
169            prepared,
170            target,
171            settings,
172            state,
173            |prepared, target, settings| {
174                context.response.usage = prepared.usage_facts(context.response.usage);
175                p::ClaudeToGeminiStream::new_with_policy(
176                    context.response,
177                    target.identities.response.clone(),
178                    settings.events.into(),
179                    target.identities.response_policy.clone(),
180                )
181            },
182        )
183        .await
184    }
185}