1mod invoke_sse;
5mod resources;
6mod stream;
7
8use crate::{
9 HttpBody, WireResponse,
10 adapt::{JsonInvocation, invoke_json},
11 capability::Upstream,
12 codec::CodecLimits,
13 transform::{
14 TransformError,
15 guardian::{
16 self, GuardianClassifyResult, GuardianDialectRequest, GuardianExtraction,
17 GuardianOperation, GuardianPreparedRequest, GuardianReviewResult,
18 },
19 },
20};
21
22pub use invoke_sse::{classify_sse, review_sse};
23pub use resources::{GuardianResourceLimits, prepare_with_resources};
24pub use stream::{GuardianStreamContext, GuardianStreamInvocation, GuardianStreamLimits};
25
26#[derive(Debug, Clone, Copy)]
27pub struct GuardianLimits {
28 pub codec: CodecLimits,
29 pub max_request_bytes: u64,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct GuardianPreflight {
34 pub operation: GuardianOperation,
35 pub model: String,
36 pub max_bytes: u64,
37}
38
39#[derive(Debug)]
40#[allow(clippy::large_enum_variant)]
41pub enum GuardianReviewInvocation {
42 Success {
43 result: GuardianReviewResult,
44 native: GuardianNativeResponse,
45 },
46 Rejected(WireResponse<HttpBody>),
47 Invalid {
48 error: TransformError,
49 native: GuardianNativeResponse,
50 },
51}
52
53#[derive(Debug)]
54#[allow(clippy::large_enum_variant)]
55pub enum GuardianClassifyInvocation {
56 Success {
57 result: GuardianClassifyResult,
58 native: GuardianNativeResponse,
59 },
60 Rejected(WireResponse<HttpBody>),
61 Invalid {
62 error: TransformError,
63 native: GuardianNativeResponse,
64 },
65}
66
67#[derive(Debug)]
72#[allow(clippy::large_enum_variant)]
73pub enum GuardianNativeResponse {
74 Claude(WireResponse<crate::wire::claude::generate_content::GenerateContentResponseBody>),
75 Gemini(WireResponse<crate::wire::gemini::GenerateContentResponseBody>),
76 Chat(WireResponse<crate::wire::openai::chat::response::GenerateContentResponseBody>),
77 Responses(WireResponse<crate::wire::openai::responses::GenerateContentResponseBody>),
78}
79
80pub async fn review<U: Upstream>(
81 upstream: &U,
82 target: &U::Target,
83 request: GuardianPreparedRequest,
84 limits: GuardianLimits,
85) -> Result<GuardianReviewInvocation, TransformError> {
86 if request.operation() != GuardianOperation::Review {
87 return Err(TransformError::new(
88 crate::transform::TransformErrorKind::Conflict,
89 "guardian.operation",
90 "review received a classify request",
91 ));
92 }
93 invoke(
94 upstream,
95 target,
96 request.into_request(),
97 GuardianOperation::Review,
98 limits,
99 )
100 .await
101 .map(map_review)
102}
103
104pub async fn classify<U: Upstream>(
105 upstream: &U,
106 target: &U::Target,
107 request: GuardianPreparedRequest,
108 limits: GuardianLimits,
109) -> Result<GuardianClassifyInvocation, TransformError> {
110 if request.operation() != GuardianOperation::Classify {
111 return Err(TransformError::new(
112 crate::transform::TransformErrorKind::Conflict,
113 "guardian.operation",
114 "classify received a review request",
115 ));
116 }
117 invoke(
118 upstream,
119 target,
120 request.into_request(),
121 GuardianOperation::Classify,
122 limits,
123 )
124 .await
125 .map(map_classify)
126}
127
128async fn invoke<U: Upstream>(
129 upstream: &U,
130 target: &U::Target,
131 request: GuardianDialectRequest,
132 operation: GuardianOperation,
133 limits: GuardianLimits,
134) -> Result<Extracted, TransformError> {
135 if limits.max_request_bytes == 0 {
136 return Err(TransformError::new(
137 crate::transform::TransformErrorKind::Limit,
138 "guardian.request",
139 "preflight byte cap must be positive",
140 ));
141 }
142 guardian::preflight_request(&request, limits.codec, limits.max_request_bytes)?;
143 match request {
144 GuardianDialectRequest::Claude(request) => {
145 match invoke_json(upstream, target, request, limits.codec).await? {
146 JsonInvocation::Success(response) => {
147 Ok(GuardianExtractionInvocation::Claude(response))
148 }
149 JsonInvocation::Rejected(response) => {
150 Ok(GuardianExtractionInvocation::Rejected(response))
151 }
152 }
153 }
154 GuardianDialectRequest::Gemini(request) => {
155 match invoke_json(upstream, target, request, limits.codec).await? {
156 JsonInvocation::Success(response) => {
157 Ok(GuardianExtractionInvocation::Gemini(response))
158 }
159 JsonInvocation::Rejected(response) => {
160 Ok(GuardianExtractionInvocation::Rejected(response))
161 }
162 }
163 }
164 GuardianDialectRequest::OpenAiChat(request) => {
165 match invoke_json(upstream, target, request, limits.codec).await? {
166 JsonInvocation::Success(response) => {
167 Ok(GuardianExtractionInvocation::Chat(response))
168 }
169 JsonInvocation::Rejected(response) => {
170 Ok(GuardianExtractionInvocation::Rejected(response))
171 }
172 }
173 }
174 GuardianDialectRequest::OpenAiResponses(request) => {
175 match invoke_json(upstream, target, request, limits.codec).await? {
176 JsonInvocation::Success(response) => {
177 Ok(GuardianExtractionInvocation::Responses(response))
178 }
179 JsonInvocation::Rejected(response) => {
180 Ok(GuardianExtractionInvocation::Rejected(response))
181 }
182 }
183 }
184 }
185 .and_then(|result| result.extract(operation))
186}
187
188#[allow(clippy::large_enum_variant)]
189enum GuardianExtractionInvocation {
190 Claude(WireResponse<crate::wire::claude::generate_content::GenerateContentResponseBody>),
191 Gemini(WireResponse<crate::wire::gemini::GenerateContentResponseBody>),
192 Chat(WireResponse<crate::wire::openai::chat::response::GenerateContentResponseBody>),
193 Responses(WireResponse<crate::wire::openai::responses::GenerateContentResponseBody>),
194 Rejected(WireResponse<HttpBody>),
195}
196
197enum Extracted {
198 Review {
199 result: GuardianReviewResult,
200 native: GuardianNativeResponse,
201 },
202 Classify {
203 result: GuardianClassifyResult,
204 native: GuardianNativeResponse,
205 },
206 Rejected(WireResponse<HttpBody>),
207 Invalid {
208 error: TransformError,
209 native: GuardianNativeResponse,
210 },
211}
212
213impl GuardianExtractionInvocation {
214 fn extract(self, operation: GuardianOperation) -> Result<Extracted, TransformError> {
215 match self {
216 Self::Claude(response) => map_response(
217 response,
218 operation,
219 guardian::extract_claude,
220 GuardianNativeResponse::Claude,
221 ),
222 Self::Gemini(response) => map_response(
223 response,
224 operation,
225 guardian::extract_gemini,
226 GuardianNativeResponse::Gemini,
227 ),
228 Self::Chat(response) => map_response(
229 response,
230 operation,
231 guardian::extract_chat,
232 GuardianNativeResponse::Chat,
233 ),
234 Self::Responses(response) => map_response(
235 response,
236 operation,
237 guardian::extract_responses,
238 GuardianNativeResponse::Responses,
239 ),
240 Self::Rejected(response) => Ok(Extracted::Rejected(response)),
241 }
242 }
243}
244
245fn map_response<T, F>(
246 response: WireResponse<T>,
247 operation: GuardianOperation,
248 extract: F,
249 native: impl FnOnce(WireResponse<T>) -> GuardianNativeResponse,
250) -> Result<Extracted, TransformError>
251where
252 F: FnOnce(T, GuardianOperation) -> Result<GuardianExtraction, TransformError>,
253 T: Clone,
254{
255 let WireResponse {
256 status,
257 headers,
258 body,
259 } = response;
260 let parsed = extract(body.clone(), operation);
261 let native = native(WireResponse {
262 status,
263 headers,
264 body,
265 });
266 match parsed {
267 Err(error) => Ok(Extracted::Invalid { error, native }),
268 Ok(GuardianExtraction::Review(value)) => Ok(Extracted::Review {
269 result: value,
270 native,
271 }),
272 Ok(GuardianExtraction::Classify(value)) => Ok(Extracted::Classify {
273 result: value,
274 native,
275 }),
276 }
277}
278
279fn map_review(value: Extracted) -> GuardianReviewInvocation {
280 match value {
281 Extracted::Review { result, native } => {
282 GuardianReviewInvocation::Success { result, native }
283 }
284 Extracted::Rejected(value) => GuardianReviewInvocation::Rejected(value),
285 Extracted::Invalid { error, native } => GuardianReviewInvocation::Invalid { error, native },
286 Extracted::Classify { .. } => unreachable!("review extraction returned classify result"),
287 }
288}
289
290fn map_classify(value: Extracted) -> GuardianClassifyInvocation {
291 match value {
292 Extracted::Classify { result, native } => {
293 GuardianClassifyInvocation::Success { result, native }
294 }
295 Extracted::Rejected(value) => GuardianClassifyInvocation::Rejected(value),
296 Extracted::Invalid { error, native } => {
297 GuardianClassifyInvocation::Invalid { error, native }
298 }
299 Extracted::Review { .. } => unreachable!("classify extraction returned review result"),
300 }
301}