rig/client/
image_generation.rs

1#[cfg(feature = "image")]
2mod image {
3    use crate::client::{AsImageGeneration, ProviderClient};
4    use crate::image_generation::{
5        ImageGenerationError, ImageGenerationModel, ImageGenerationModelDyn,
6        ImageGenerationRequest, ImageGenerationResponse,
7    };
8    use std::future::Future;
9    use std::sync::Arc;
10
11    /// A provider client with image generation capabilities.
12    /// Clone is required for conversions between client types.
13    pub trait ImageGenerationClient: ProviderClient + Clone {
14        /// The ImageGenerationModel used by the Client
15        type ImageGenerationModel: ImageGenerationModel;
16
17        /// Create an image generation model with the given name.
18        ///
19        /// # Example with OpenAI
20        /// ```
21        /// use rig::prelude::*;
22        /// use rig::providers::openai::{Client, self};
23        ///
24        /// // Initialize the OpenAI client
25        /// let openai = Client::new("your-open-ai-api-key");
26        ///
27        /// let gpt4 = openai.image_generation_model(openai::DALL_E_3);
28        /// ```
29        fn image_generation_model(&self, model: &str) -> Self::ImageGenerationModel;
30    }
31
32    pub trait ImageGenerationClientDyn: ProviderClient {
33        /// Create an image generation model with the given name.
34        fn image_generation_model<'a>(&self, model: &str) -> Box<dyn ImageGenerationModelDyn + 'a>;
35    }
36
37    impl<T: ImageGenerationClient<ImageGenerationModel = M>, M: ImageGenerationModel + 'static>
38        ImageGenerationClientDyn for T
39    {
40        fn image_generation_model<'a>(&self, model: &str) -> Box<dyn ImageGenerationModelDyn + 'a> {
41            Box::new(self.image_generation_model(model))
42        }
43    }
44
45    impl<T: ImageGenerationClientDyn + Clone + 'static> AsImageGeneration for T {
46        fn as_image_generation(&self) -> Option<Box<dyn ImageGenerationClientDyn>> {
47            Some(Box::new(self.clone()))
48        }
49    }
50
51    /// Wraps a ImageGenerationModel in a dyn-compatible way for ImageGenerationRequestBuilder.
52    #[derive(Clone)]
53    pub struct ImageGenerationModelHandle<'a> {
54        pub(crate) inner: Arc<dyn ImageGenerationModelDyn + 'a>,
55    }
56    impl ImageGenerationModel for ImageGenerationModelHandle<'_> {
57        type Response = ();
58
59        fn image_generation(
60            &self,
61            request: ImageGenerationRequest,
62        ) -> impl Future<
63            Output = Result<ImageGenerationResponse<Self::Response>, ImageGenerationError>,
64        > + Send {
65            self.inner.image_generation(request)
66        }
67    }
68}
69
70#[cfg(feature = "image")]
71pub use image::*;
rig/client/image_generation.rs

rig/client/
image_generation.rs
