gproxy_protocol/adapt/generate/image_resources/
mod.rs

1//! Explicit image reads and publication with caller-owned cancellation progress.
2//! The native Gemini response is kept apart from the view its images are read or
3//! published into; an image signature travels with the view to the client.
4
5mod publish;
6mod read;
7mod stream;
8mod thread;
9use super::GenerationProgress;
10use crate::{
11    capability::PublishedResource,
12    transform::{TransformError, TransformErrorKind},
13    wire::gemini as g,
14};
15use std::collections::BTreeMap;
16pub use stream::ImageStreamProgress;
17pub use thread::{ResourceSend, ResourceSync};
18
19#[derive(Debug, Clone)]
20pub struct MaterializedImage {
21    pub candidate: usize,
22    pub part: usize,
23    pub original: g::Part,
24    pub materialized: g::Blob,
25    /// Expiry observed while reading the original resource, if the host knows it.
26    pub resource_expiry: Option<std::time::SystemTime>,
27}
28
29/// Retains completed reads and an active resource body. Bytes are counted as
30/// each chunk arrives; cancellation resumes that same body without a new read.
31#[derive(Default)]
32pub struct GeminiImageReads {
33    original: Option<g::GenerateContentResponseBody>,
34    images: BTreeMap<(usize, usize), MaterializedImage>,
35    attempted: usize,
36    bytes: u64,
37    pending: Option<PendingImageRead>,
38    failed: bool,
39}
40
41impl GeminiImageReads {
42    pub fn bytes_read(&self) -> u64 {
43        self.bytes
44    }
45    pub fn attempted_reads(&self) -> usize {
46        self.attempted
47    }
48    pub fn images(&self) -> impl Iterator<Item = &MaterializedImage> {
49        self.images.values()
50    }
51}
52
53pub(super) struct Publication {
54    source: g::Blob,
55    operation: String,
56    expiry: std::time::SystemTime,
57    receipt: Option<usize>,
58}
59
60/// Receipts, including malformed host replies, remain inspectable for host
61/// compensation. An uncertain publication is queried, never automatically sent
62/// again. This object and its resource scope belong to one invocation.
63pub struct GeminiImagePublications<H> {
64    original: Option<g::GenerateContentResponseBody>,
65    publications: BTreeMap<(usize, usize), Publication>,
66    receipts: Vec<PublishedResource<H>>,
67    failed: bool,
68}
69
70impl<H> Default for GeminiImagePublications<H> {
71    fn default() -> Self {
72        Self {
73            original: None,
74            publications: BTreeMap::new(),
75            receipts: Vec::new(),
76            failed: false,
77        }
78    }
79}
80
81impl<H> GeminiImagePublications<H> {
82    pub fn operation_ids(&self) -> impl Iterator<Item = &str> {
83        self.publications
84            .values()
85            .map(|value| value.operation.as_str())
86    }
87    pub fn receipts(&self) -> &[PublishedResource<H>] {
88        &self.receipts
89    }
90}
91
92/// Generation and resource work have independent retained progress.
93/// Re-enter through `recover_with_image_resources`, never repeat an uncertain POST.
94pub struct ImageResourceProgress<N, H> {
95    pub generation: GenerationProgress<N>,
96    pub reads: GeminiImageReads,
97    pub publications: GeminiImagePublications<H>,
98}
99
100impl<N, H> Default for ImageResourceProgress<N, H> {
101    fn default() -> Self {
102        Self {
103            generation: Default::default(),
104            reads: Default::default(),
105            publications: Default::default(),
106        }
107    }
108}
109
110fn invalid(message: impl Into<String>) -> TransformError {
111    TransformError::invalid_result("generation.image_resources", message)
112}
113
114fn missing(message: impl Into<String>) -> TransformError {
115    TransformError::new(
116        TransformErrorKind::MissingState,
117        "generation.image_resources",
118        message,
119    )
120}
121
122fn conflict(message: impl Into<String>) -> TransformError {
123    TransformError::new(
124        TransformErrorKind::Conflict,
125        "generation.image_resources",
126        message,
127    )
128}
129
130fn limit() -> TransformError {
131    TransformError::new(
132        TransformErrorKind::Limit,
133        "generation.image_resources",
134        "image resource budget exceeded",
135    )
136}
137
138pub(crate) fn wants_uri(input: &g::GenerateContentRequestBody) -> bool {
139    input
140        .generation_config
141        .as_ref()
142        .and_then(|config| config.response_format.as_ref())
143        .and_then(|format| format.image.as_ref())
144        .and_then(|image| image.delivery.as_ref())
145        == Some(&g::Delivery::Uri)
146}
147
148struct PendingImageRead {
149    key: (usize, usize),
150    reference: crate::capability::ResourceReference,
151    metadata: crate::capability::ResourceMetadata,
152    body: crate::HttpBody,
153    buffer: Vec<u8>,
154    limit: u64,
155}
156
157fn codec_error(error: crate::codec::CodecError) -> TransformError {
158    let kind = if error.kind() == crate::codec::CodecErrorKind::Limit {
159        TransformErrorKind::Limit
160    } else {
161        TransformErrorKind::InvalidResult
162    };
163    TransformError::with_source(kind, "generation.image_resources", error.to_string(), error)
164}
gproxy_protocol/adapt/generate/image_resources/mod.rs

gproxy_protocol/adapt/generate/image_resources/
mod.rs
