rig_core/
image_generation.rs

1//! Image-generation requests, normalized responses, and model interfaces.
2//!
3//! ```no_run
4//! use rig_core::DynModel;
5//! use rig_core::image_generation::ImageGenerationRequestBuilder;
6//! use rig_core::operation::ImageGeneration;
7//!
8//! # async fn example(model: DynModel<ImageGeneration>) -> Result<(), Box<dyn std::error::Error>> {
9//! let request = ImageGenerationRequestBuilder::new("A mountain lake")
10//!     .width(1024)
11//!     .height(1024)
12//!     .build();
13//! let response = model.call(request).await?;
14//! # let _ = response;
15//! # Ok(())
16//! # }
17//! ```
18use crate::completion::Usage;
19use crate::message::Image;
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22
23/// Generated image bytes and normalized provider metadata.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct ImageGenerationResponse {
26    /// The generated image, decoded to bytes.
27    pub image: Vec<u8>,
28    /// Usage as the provider reported it; every counter is `None` when the
29    /// provider reported none (see [`Usage`]).
30    #[serde(default)]
31    pub usage: Usage,
32    /// Stable descriptor name of the provider that produced this response,
33    /// for example `"openai"`. Always populated.
34    pub provider: String,
35    /// Provider-reported model identifier, when the wire response named one.
36    #[serde(default)]
37    pub model: Option<String>,
38    /// Provider-assigned response-scoped identifier, when reported.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub response_id: Option<String>,
41    /// Transport request ID from HTTP headers, or `None` when unreported.
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub provider_request_id: Option<String>,
44    /// Provider response metadata. May be null for byte-only responses or
45    /// responses constructed without metadata; image bytes remain in [`Self::image`].
46    #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
47    pub raw: serde_json::Value,
48}
49
50impl ImageGenerationResponse {
51    /// A response carrying `image`. The driver writes the provider, the
52    /// transport request id and the reply document; decoders set what the
53    /// provider reported.
54    pub fn new(image: Vec<u8>) -> Self {
55        Self {
56            image,
57            usage: Usage::default(),
58            provider: String::new(),
59            model: None,
60            response_id: None,
61            provider_request_id: None,
62            raw: serde_json::Value::Null,
63        }
64    }
65}
66
67pub struct ImageGenerationRequest {
68    pub prompt: String,
69    pub width: u32,
70    pub height: u32,
71    /// Reference images the model draws from or edits. A provider that
72    /// takes none refuses a request that names one.
73    pub images: Vec<Image>,
74    pub additional_params: Option<Value>,
75}
76
77/// Builds an image request for a prompt. Defaults to 256 by 256 pixels;
78/// supported dimensions depend on the provider.
79pub struct ImageGenerationRequestBuilder {
80    request: ImageGenerationRequest,
81}
82
83impl ImageGenerationRequestBuilder {
84    /// A request for `prompt`.
85    pub fn new(prompt: impl Into<String>) -> Self {
86        Self {
87            request: ImageGenerationRequest {
88                prompt: prompt.into(),
89                width: 256,
90                height: 256,
91                images: Vec::new(),
92                additional_params: None,
93            },
94        }
95    }
96
97    /// The width of the generated image.
98    pub fn width(mut self, width: u32) -> Self {
99        self.request.width = width;
100        self
101    }
102
103    /// The height of the generated image.
104    pub fn height(mut self, height: u32) -> Self {
105        self.request.height = height;
106        self
107    }
108
109    /// Add a reference image the model draws from or edits.
110    pub fn image(mut self, image: Image) -> Self {
111        self.request.images.push(image);
112        self
113    }
114
115    /// Merges provider-specific parameters over earlier ones, key by key for
116    /// JSON objects; `None` clears existing parameters.
117    pub fn additional_params(mut self, params: impl Into<Option<Value>>) -> Self {
118        self.request.additional_params =
119            crate::json_utils::merge_params(self.request.additional_params.take(), params.into());
120        self
121    }
122
123    /// Builds the image generation request.
124    pub fn build(self) -> ImageGenerationRequest {
125        self.request
126    }
127}
128
129#[cfg(test)]
130mod builder_tests;
131#[cfg(test)]
132mod provider_response_tests;
rig_core/image_generation.rs

rig_core/
image_generation.rs
