ferrin_google/batch/
results.rs

1//! Conversion of batch result lines to specification items.
2//!
3//! Derived from the Vercel AI SDK (Apache-2.0, Copyright 2023 Vercel, Inc.),
4//! translated from TypeScript to Rust and modified; see `NOTICE`.
5
6use ferrin_provider_util::tool_name_mapping::ToolNameMapping;
7use ferrin_spec::Content;
8use ferrin_spec::JsonObject;
9use ferrin_spec::JsonValue;
10use ferrin_spec::ModelId;
11use ferrin_spec::batch::BatchError;
12use ferrin_spec::batch::BatchItem;
13use ferrin_spec::batch::BatchItemResult;
14use ferrin_spec::language_model::GenerateResult;
15use serde::Deserialize;
16use serde_json::json;
17
18use super::rpc_error;
19use crate::api_types::GenerateContentResponse;
20use crate::api_types::RpcStatus;
21use crate::config::SharedConfig;
22use crate::image::image_result;
23use crate::language_model::convert_generate_content_response;
24use crate::output::OutputMapper;
25
26/// One line of a batch results file (`{key, response?, error?}`).
27#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
28pub struct BatchResultLine {
29    /// Caller-assigned request id.
30    pub key: String,
31    /// `generateContent` response.
32    #[serde(default)]
33    pub response: Option<JsonValue>,
34    /// Error of a failed request.
35    #[serde(default)]
36    pub error: Option<RpcStatus>,
37}
38
39/// One inlined response (`{metadata: {key}, response?, error?}`).
40#[derive(Debug, Clone, Deserialize)]
41pub(super) struct InlinedResponse {
42    metadata: InlinedMetadata,
43    #[serde(default)]
44    response: Option<JsonValue>,
45    #[serde(default)]
46    error: Option<RpcStatus>,
47}
48
49#[derive(Debug, Clone, Deserialize)]
50struct InlinedMetadata {
51    key: String,
52}
53
54impl InlinedResponse {
55    pub(super) fn into_line(self) -> BatchResultLine {
56        BatchResultLine {
57            key: self.metadata.key,
58            response: self.response,
59            error: self.error,
60        }
61    }
62}
63
64fn failed(id: String, message: impl Into<String>, code: &str) -> BatchItem<GenerateResult> {
65    BatchItem::Failed {
66        id,
67        error: BatchError {
68            message: message.into(),
69            error_type: None,
70            code: Some(code.to_owned()),
71            status_code: None,
72        },
73        provider_metadata: None,
74    }
75}
76
77fn text(item: BatchItem<GenerateResult>) -> BatchItemResult {
78    BatchItemResult::Text(Box::new(item))
79}
80
81fn is_image_file(content: &Content) -> bool {
82    matches!(
83        content,
84        Content::File { media_type, .. } if media_type.as_str().starts_with("image/")
85    )
86}
87
88fn unsupported_kind(content: &Content) -> Option<&'static str> {
89    match content {
90        Content::Text { .. }
91        | Content::Reasoning { .. }
92        | Content::Source(_)
93        | Content::ToolCall(_)
94        | Content::ToolResult(_) => None,
95        Content::File { .. } => Some("file"),
96        Content::ReasoningFile { .. } => Some("reasoning-file"),
97        Content::Custom { .. } => Some("custom"),
98        Content::ToolApprovalRequest { .. } => Some("tool-approval-request"),
99        #[allow(unreachable_patterns, reason = "Content is non-exhaustive")]
100        _ => Some("unknown"),
101    }
102}
103
104fn blocked(config: &SharedConfig, id: String, response: &JsonValue) -> BatchItemResult {
105    let prompt_feedback = response
106        .get("promptFeedback")
107        .and_then(JsonValue::as_object);
108    let block_reason = prompt_feedback
109        .and_then(|feedback| feedback.get("blockReason"))
110        .and_then(JsonValue::as_str);
111    let (message, code) = match block_reason {
112        Some(reason) => (
113            format!("Google blocked the batch request ({reason})"),
114            "prompt_blocked",
115        ),
116        None => (
117            "Google returned a batch response without any candidates".to_owned(),
118            "invalid_response",
119        ),
120    };
121    let provider_metadata = prompt_feedback.map(|_| {
122        let mut object = JsonObject::new();
123        object.insert(
124            "promptFeedback".to_owned(),
125            json!({"blockReason": block_reason}),
126        );
127        OutputMapper::new(config.clone(), ToolNameMapping::default()).metadata(object)
128    });
129    text(BatchItem::Failed {
130        id,
131        error: BatchError {
132            message,
133            error_type: block_reason.map(str::to_owned),
134            code: Some(code.to_owned()),
135            status_code: None,
136        },
137        provider_metadata,
138    })
139}
140
141/// Converts one result line.
142#[must_use]
143pub fn convert_line(config: &SharedConfig, line: BatchResultLine) -> BatchItemResult {
144    let id = line.key;
145    if let Some(status) = line.error {
146        let cancelled = status.status.as_deref() == Some("CANCELLED") || status.code == Some(1);
147        let error = rpc_error(&status, "Google batch request failed");
148        return text(if cancelled {
149            BatchItem::Cancelled {
150                id,
151                error: Some(error),
152                provider_metadata: None,
153            }
154        } else {
155            BatchItem::Failed {
156                id,
157                error,
158                provider_metadata: None,
159            }
160        });
161    }
162    let Some(response) = line.response else {
163        return text(failed(
164            id,
165            "Google returned a batch result without a response or error",
166            "invalid_batch_result",
167        ));
168    };
169    if response
170        .get("candidates")
171        .and_then(JsonValue::as_array)
172        .is_none_or(Vec::is_empty)
173    {
174        return blocked(config, id, &response);
175    }
176    let Ok(body) = serde_json::from_value::<GenerateContentResponse>(response.clone()) else {
177        return text(failed(
178            id,
179            "Google returned an invalid GenerateContent batch result",
180            "invalid_response",
181        ));
182    };
183    let result = match convert_generate_content_response(
184        config,
185        ToolNameMapping::default(),
186        Vec::new(),
187        &body,
188        Some(&response),
189    ) {
190        Ok(result) => result,
191        Err(error) => return text(failed(id, error.to_string(), "invalid_response")),
192    };
193    if result.content.iter().any(is_image_file) {
194        let model_id = result
195            .response
196            .model_id
197            .clone()
198            .unwrap_or_else(|| ModelId::new(""));
199        return BatchItemResult::Image(Box::new(BatchItem::Succeeded {
200            id,
201            result: image_result(config, model_id, result, Vec::new()),
202        }));
203    }
204    if let Some(kind) = result.content.iter().find_map(unsupported_kind) {
205        return text(failed(
206            id,
207            format!(
208                "Google returned a \"{kind}\" content block, but that content is not supported in text batches"
209            ),
210            "unsupported_content",
211        ));
212    }
213    text(BatchItem::Succeeded { id, result })
214}
ferrin_google/batch/results.rs

ferrin_google/batch/
results.rs
