use super::*;
use crate::{
HttpBody,
capability::{ResourceAccess, ResourceReference},
codec,
transform::images::inspect_image,
wire::DeclaredFields,
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use futures_util::StreamExt;
impl GeminiImageReads {
pub async fn materialize<R: ResourceAccess>(
&mut self,
input: &g::GenerateContentResponseBody,
resources: &super::super::GenerationResources<'_, R>,
) -> Result<g::GenerateContentResponseBody, TransformError> {
if self.failed {
return Err(invalid("image read progress failed"));
}
let result = self.materialize_inner(input, resources).await;
if result.as_ref().err().is_some_and(|error| {
matches!(
error.kind(),
TransformErrorKind::InvalidInput
| TransformErrorKind::InvalidResult
| TransformErrorKind::Unsupported
| TransformErrorKind::Limit
)
}) {
self.pending = None;
self.failed = true;
}
result
}
async fn materialize_inner<R: ResourceAccess>(
&mut self,
input: &g::GenerateContentResponseBody,
resources: &super::super::GenerationResources<'_, R>,
) -> Result<g::GenerateContentResponseBody, TransformError> {
let input = input.clone().into_declared();
codec::encode_json(&input, resources.limits).map_err(codec_error)?;
if let Some(original) = &self.original {
if original != &input {
return Err(conflict("native image response changed during recovery"));
}
} else {
self.original = Some(input.clone());
}
let mut output = input;
for (candidate_index, candidate) in output.candidates.iter_mut().flatten().enumerate() {
for (part_index, part) in candidate
.content
.iter_mut()
.flat_map(|content| content.parts.iter_mut().flatten())
.enumerate()
{
let Some(file) = &part.file_data else {
continue;
};
if part.inline_data.is_some()
|| part.text.is_some()
|| part.thought == Some(true)
|| part.function_call.is_some()
|| part.function_response.is_some()
|| part.executable_code.is_some()
|| part.code_execution_result.is_some()
|| part.tool_call.is_some()
|| part.tool_response.is_some()
|| part.video_metadata.is_some()
|| part.media_processing.is_some()
|| part.audio_transcription.is_some()
|| part.speech_metadata.is_some()
|| part.part_metadata.is_some()
|| part.media_resolution.is_some()
|| file.file_uri.trim().is_empty()
{
return Err(invalid(
"generated file image must be an unambiguous standalone Part",
));
}
let key = (candidate_index, part_index);
if !self.images.contains_key(&key) {
self.read_image(key, file, resources).await?;
let pending = self.pending.take().expect("completed retained read");
let actual = inspect_image(&pending.buffer)?.mime();
if pending
.metadata
.mime
.as_deref()
.is_some_and(|mime| mime != actual)
|| file.mime_type.as_deref().is_some_and(|mime| mime != actual)
{
return Err(invalid("generated image MIME contradicts actual container"));
}
self.images.insert(
key,
MaterializedImage {
candidate: candidate_index,
part: part_index,
original: part.clone(),
resource_expiry: pending.metadata.expires_at,
materialized: g::Blob::builder(
actual.to_owned(),
STANDARD.encode(pending.buffer),
)
.build(),
},
);
}
let image = self.images.get(&key).expect("materialized image retained");
if &image.original != part {
return Err(conflict("native image Part changed"));
}
part.file_data = None;
part.inline_data = Some(image.materialized.clone());
}
}
codec::encode_json(&output, resources.limits).map_err(codec_error)?;
Ok(output)
}
async fn read_image<R: ResourceAccess>(
&mut self,
key: (usize, usize),
file: &g::FileData,
resources: &super::super::GenerationResources<'_, R>,
) -> Result<(), TransformError> {
let reference = ResourceReference::Url(file.file_uri.clone());
let total = resources
.limits
.max_body_bytes
.min(resources.access.limits().read_bytes);
if let Some(pending) = &self.pending {
if pending.key != key || pending.reference != reference {
return Err(conflict("pending image read changed"));
}
} else {
if self.attempted >= resources.max_references || self.bytes >= total {
return Err(limit());
}
self.attempted += 1;
let read = resources.access.read(resources.scope, &reference).await?;
self.pending = Some(PendingImageRead {
key,
reference,
metadata: read.metadata,
body: read.body,
buffer: Vec::new(),
limit: total
.saturating_sub(self.bytes)
.min(resources.limits.max_buffer_bytes)
.min(resources.limits.max_part_bytes),
});
}
let pending = self.pending.as_mut().expect("retained resource body");
if pending
.metadata
.expires_at
.is_some_and(|expiry| expiry <= resources.now)
{
return Err(missing("generated image resource expired"));
}
let available = total
.saturating_sub(self.bytes)
.saturating_add(pending.buffer.len() as u64);
pending.limit = pending
.limit
.min(available)
.min(resources.limits.max_buffer_bytes)
.min(resources.limits.max_part_bytes);
if pending.metadata.length.is_some_and(|n| n > pending.limit) || self.bytes > total {
return Err(limit());
}
let mut chunks_since_yield = 0usize;
loop {
let next = match &mut self.pending.as_mut().expect("retained body").body {
HttpBody::Bytes(bytes) => {
if bytes.is_empty() {
None
} else {
Some(Ok(std::mem::take(bytes)))
}
}
HttpBody::Stream(stream) => stream.next().await,
};
let Some(chunk) = next else {
break;
};
let chunk = chunk.map_err(|error| {
TransformError::with_source(
TransformErrorKind::Host,
"generation.image_resources.read",
error.to_string(),
error,
)
})?;
self.bytes = self
.bytes
.checked_add(chunk.len() as u64)
.ok_or_else(limit)?;
let pending = self.pending.as_mut().expect("retained body");
if self.bytes > total
|| (pending.buffer.len() as u64).saturating_add(chunk.len() as u64) > pending.limit
{
return Err(limit());
}
pending.buffer.extend_from_slice(&chunk);
chunks_since_yield += 1;
if chunks_since_yield == 1024 {
chunks_since_yield = 0;
let mut yielded = false;
std::future::poll_fn(|cx| {
if yielded {
std::task::Poll::Ready(())
} else {
yielded = true;
cx.waker().wake_by_ref();
std::task::Poll::Pending
}
})
.await;
}
}
let pending = self.pending.as_ref().expect("complete retained body");
if pending
.metadata
.length
.is_some_and(|n| n != pending.buffer.len() as u64)
{
return Err(invalid(
"generated image length contradicts resource metadata",
));
}
Ok(())
}
}