Skip to main content

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

1use std::future::Future;
2
3use super::super::super::{
4    Endpoint, GenerationIdentity, GenerationResources, GenerationStateAccess,
5    chat_responses::{ChatViaResponses, ResponsesViaChat},
6};
7use super::super::{ResponsesHistoryCache, StreamInvocation, StreamSettings, StreamTarget};
8use super::{RequestMode, start};
9use crate::{
10    capability::{ResourceAccess, StateStore},
11    transform::{TransformError, generate::chat_responses::stream as p},
12    wire::{
13        DeclaredFields,
14        openai::{chat as h, responses as r},
15    },
16};
17
18impl ChatViaResponses {
19    pub async fn prepare_stream<S: StateStore>(
20        input: h::GenerateContentRequestBody,
21        target: StreamTarget,
22        settings: StreamSettings,
23        state: &GenerationStateAccess<'_, S>,
24    ) -> Result<StreamInvocation<p::ResponsesToChatStream>, TransformError> {
25        Self::stream(input, target, settings, state, |input, endpoint, ids| {
26            Self::prepare_with_state(input, endpoint, ids, state)
27        })
28        .await
29    }
30    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
31        input: h::GenerateContentRequestBody,
32        target: StreamTarget,
33        settings: StreamSettings,
34        state: &GenerationStateAccess<'_, S>,
35        resources: &GenerationResources<'_, R>,
36    ) -> Result<StreamInvocation<p::ResponsesToChatStream>, TransformError> {
37        Self::stream(input, target, settings, state, |input, endpoint, ids| {
38            Self::prepare_with_capabilities(input, endpoint, ids, state, resources)
39        })
40        .await
41    }
42    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
43        input: h::GenerateContentRequestBody,
44        target: StreamTarget,
45        settings: StreamSettings,
46        state: &GenerationStateAccess<'_, S>,
47        prepare: impl FnOnce(h::GenerateContentRequestBody, Endpoint, GenerationIdentity) -> Fut,
48    ) -> Result<StreamInvocation<p::ResponsesToChatStream>, TransformError> {
49        let original = input.into_declared();
50        let prepared = prepare(
51            original.clone().buffered(),
52            target.endpoint.clone(),
53            target.identities.clone(),
54        )
55        .await?;
56        start(
57            original,
58            prepared,
59            target,
60            settings,
61            state,
62            |_, target, settings| {
63                p::ResponsesToChatStream::new_with_policy(
64                    target.identities.response.clone(),
65                    settings.events.into(),
66                    target.identities.response_policy.clone(),
67                )
68            },
69        )
70        .await
71    }
72}
73
74impl ResponsesViaChat {
75    pub async fn prepare_stream<S: StateStore>(
76        input: r::GenerateContentRequestBody,
77        target: StreamTarget,
78        context: p::ChatToResponsesContext,
79        settings: StreamSettings,
80        state: &GenerationStateAccess<'_, S>,
81    ) -> Result<StreamInvocation<p::ChatToResponsesStream>, TransformError> {
82        Self::stream(
83            input,
84            target,
85            context,
86            settings,
87            state,
88            None,
89            |input, endpoint, ids| Self::prepare_with_state(input, endpoint, ids, state),
90        )
91        .await
92    }
93    pub async fn prepare_stream_with_history_cache<S: StateStore>(
94        input: r::GenerateContentRequestBody,
95        target: StreamTarget,
96        context: p::ChatToResponsesContext,
97        settings: StreamSettings,
98        state: &GenerationStateAccess<'_, S>,
99        history_cache: &ResponsesHistoryCache,
100    ) -> Result<StreamInvocation<p::ChatToResponsesStream>, TransformError> {
101        let cache = Some(history_cache);
102        Self::stream(
103            input,
104            target,
105            context,
106            settings,
107            state,
108            cache,
109            |input, endpoint, ids| Self::prepare_with_state(input, endpoint, ids, state),
110        )
111        .await
112    }
113    pub async fn prepare_stream_with_capabilities<S: StateStore, R: ResourceAccess>(
114        input: r::GenerateContentRequestBody,
115        target: StreamTarget,
116        context: p::ChatToResponsesContext,
117        settings: StreamSettings,
118        state: &GenerationStateAccess<'_, S>,
119        resources: &GenerationResources<'_, R>,
120    ) -> Result<StreamInvocation<p::ChatToResponsesStream>, TransformError> {
121        Self::stream(
122            input,
123            target,
124            context,
125            settings,
126            state,
127            None,
128            |input, endpoint, ids| {
129                Self::prepare_with_capabilities(input, endpoint, ids, state, resources)
130            },
131        )
132        .await
133    }
134    pub async fn prepare_stream_with_capabilities_and_history_cache<
135        S: StateStore,
136        R: ResourceAccess,
137    >(
138        input: r::GenerateContentRequestBody,
139        target: StreamTarget,
140        context: p::ChatToResponsesContext,
141        settings: StreamSettings,
142        state: &GenerationStateAccess<'_, S>,
143        resources: &GenerationResources<'_, R>,
144        history_cache: &ResponsesHistoryCache,
145    ) -> Result<StreamInvocation<p::ChatToResponsesStream>, TransformError> {
146        let cache = Some(history_cache);
147        Self::stream(
148            input,
149            target,
150            context,
151            settings,
152            state,
153            cache,
154            |input, endpoint, ids| {
155                Self::prepare_with_capabilities(input, endpoint, ids, state, resources)
156            },
157        )
158        .await
159    }
160    async fn stream<S: StateStore, Fut: Future<Output = Result<Self, TransformError>>>(
161        input: r::GenerateContentRequestBody,
162        target: StreamTarget,
163        mut context: p::ChatToResponsesContext,
164        settings: StreamSettings,
165        state: &GenerationStateAccess<'_, S>,
166        history_cache: Option<&ResponsesHistoryCache>,
167        prepare: impl FnOnce(r::GenerateContentRequestBody, Endpoint, GenerationIdentity) -> Fut,
168    ) -> Result<StreamInvocation<p::ChatToResponsesStream>, TransformError> {
169        let original = input.into_declared();
170        let (history, expanded) = super::super::history::History::prepare_with_cache(
171            &original,
172            state,
173            settings.codec,
174            history_cache,
175        )
176        .await?;
177        context.response.request = original.clone();
178        context.response.request.input = expanded.input.clone();
179        let prepared = prepare(
180            expanded.buffered(),
181            target.endpoint.clone(),
182            target.identities.clone(),
183        )
184        .await?;
185        let mut invocation = start(
186            original,
187            prepared,
188            target,
189            settings,
190            state,
191            |_, target, settings| {
192                history.claim_response_id(&mut target.identities, crate::Dialect::OpenAiChat)?;
193                p::ChatToResponsesStream::new_with_policy(
194                    context,
195                    target.identities.response.clone(),
196                    settings.events.into(),
197                    target.identities.response_policy.clone(),
198                )
199            },
200        )
201        .await?;
202        history.bind(&invocation.binding, state)?;
203        invocation.history = Some(history);
204        Ok(invocation)
205    }
206}