1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

//! Rust Client for interacting with the [Replicate API](https://replicate.com/docs/api/).
//!
//! ### Example
//! In this example we will run a model that generates a caption for an image using the [Stable Diffusion](
//! https://replicate.ai/stability-ai/stable-diffusion) model.
//!
//! ```
//! use replicate_rust::Replicate;
//!
//! // Create a new Replicate client.
//! let replicate = Replicate::new()
//!
//! // Creating the inputs
//! let mut inputs = std::collections::HashMap::new();
//! inputs.insert("prompt", "a  19th century portrait of a wombat gentleman");
//!
//! let version = String::from("stability-ai/stable-diffusion:27b93a2413e7f36cd83da926f3656280b2931564ff050bf9575f1fdf9bcd7478");
//!
//! // Run the model.
//! let result = replicate.run(version, inputs);
//!
//! // Print the result
//! match result {
//!     Ok(result) => println!("Success : {:?}", result.output),
//!     Err(e) => println!("Error : {}", e),
//! }
//!
//! ```

use std::collections::HashMap;

use api_definitions::GetPrediction;
use collection::Collection;
use model::Model;
use prediction::Prediction;
use training::Training;

pub mod client;
pub mod collection;
pub mod model;
pub mod prediction;
pub mod training;
pub mod version;

pub mod api_definitions;
pub mod prediction_client;
pub mod retry;

pub struct Replicate {
    client: client::Client,
    pub predictions: Prediction,
    pub models: Model,
    pub training: Training,
    pub collection: Collection,
}

/// Rust Client for interacting with the [Replicate API](https://replicate.com/docs/api/).
///
impl Replicate {
    /// Create a new Replicate client.
    /// 
    /// # Example
    /// ```
    /// use replicate_rust::Replicate;
    /// let replicate = Replicate::new();
    /// ```
    pub fn new() -> Self {
        let client = client::Client::new();

        // TODO : Maybe reference instead of clone
        let predictions = Prediction::new(client.clone());
        let models = Model::new(client.clone());
        let training = Training::new(client.clone());
        let collection = Collection::new(client.clone());

        Self {
            client,
            predictions,
            models,
            training,
            collection,
        }
    }

    /// Run a model with the given inputs in a blocking manner.
    /// # Arguments
    /// * `version` - The version of the model to run.
    /// * `inputs` - The inputs to the model in the form of a HashMap.
    /// # Example
    /// ```
    /// use replicate_rust::Replicate;
    /// let replicate = Replicate::new();
    ///
    /// // Construct the inputs.
    /// let mut inputs = std::collections::HashMap::new();
    /// inputs.insert("prompt", "a  19th century portrait of a wombat gentleman");
    ///
    /// let version = String::from("stability-ai/stable-diffusion:27b93a2413e7f36cd83da926f3656280b2931564ff050bf9575f1fdf9bcd7478");
    ///
    /// // Run the model.
    /// let result = replicate.run(version, inputs);
    ///
    /// // Print the result.
    /// match result {
    ///    Ok(result) => println!("Success : {:?}", result.output),
    ///   Err(e) => println!("Error : {}", e),
    /// }
    /// ```
    /// # Errors
    ///
    /// TODO : Add errors
    ///
    pub fn run<K: serde::Serialize, V: serde::Serialize>(
        &self,
        version: String,
        inputs: HashMap<K, V>,
        // TODO : Perhaps not Box<dyn std::error::Error> but something more specific?
    ) -> Result<GetPrediction, Box<dyn std::error::Error>> {
        let prediction = Prediction::new(self.client.clone()).create(version, inputs);

        prediction.wait()
    }
}

// mod tests {
//     use super::*;

//     #[test]
//     fn it_works() {
//         let result = add(2, 2);
//         assert_eq!(result, 4);
//     }
// }