swiftide_integrations/dashscope/
mod.rs

1use std::sync::Arc;
2
3use config::DashscopeConfig;
4use derive_builder::Builder;
5
6mod config;
7mod embed;
8mod simple_prompt;
9
10#[derive(Debug, Builder, Clone)]
11#[builder(setter(into, strip_option))]
12pub struct Dashscope {
13    #[builder(default = "default_client()", setter(custom))]
14    client: Arc<async_openai::Client<DashscopeConfig>>,
15    /// Default options for prompt models.
16    #[builder(default)]
17    default_options: Options,
18}
19
20impl Default for Dashscope {
21    fn default() -> Self {
22        Self {
23            client: default_client(),
24            default_options: Options::default(),
25        }
26    }
27}
28
29fn default_client() -> Arc<async_openai::Client<DashscopeConfig>> {
30    async_openai::Client::with_config(DashscopeConfig::default()).into()
31}
32
33#[derive(Debug, Default, Clone, Builder)]
34#[builder(setter(into, strip_option))]
35pub struct Options {
36    /// The default prompt model to use, if specified.
37    #[builder(default)]
38    pub prompt_model: Option<String>,
39    #[builder(default)]
40    pub embed_model: Option<String>,
41    #[builder(default)]
42    pub dimensions: u16,
43}
44
45impl Options {
46    /// Creates a new `OptionsBuilder` for constructing `Options` instances.
47    pub fn builder() -> OptionsBuilder {
48        OptionsBuilder::default()
49    }
50}
51
52impl Dashscope {
53    /// Creates a new `DashscopeBuilder` for constructing `Dashscope` instances.
54    pub fn builder() -> DashscopeBuilder {
55        DashscopeBuilder::default()
56    }
57
58    /// Sets a default prompt model to use when prompting
59    pub fn with_default_prompt_model(&mut self, model: impl Into<String>) -> &mut Self {
60        self.default_options = Options {
61            prompt_model: Some(model.into()),
62            ..Default::default()
63        };
64        self
65    }
66
67    pub fn with_default_embed_model(&mut self, model: impl Into<String>) -> &mut Self {
68        self.default_options = Options {
69            embed_model: Some(model.into()),
70            ..Default::default()
71        };
72        self
73    }
74}
75
76impl DashscopeBuilder {
77    /// Sets the `Dashscope` client for the `Dashscope` instance.
78    ///
79    /// # Parameters
80    /// - `client`: The `Dashscope` client to set.
81    ///
82    /// # Returns
83    /// A mutable reference to the `DashscopeBuilder`.
84    pub fn client(&mut self, client: async_openai::Client<DashscopeConfig>) -> &mut Self {
85        self.client = Some(Arc::new(client));
86        self
87    }
88
89    /// Sets the default prompt model for the `Dashscope` instance.
90    ///
91    /// # Parameters
92    /// - `model`: The prompt model to set.
93    ///
94    /// # Returns
95    /// A mutable reference to the `DashscopeBuilder`.
96    pub fn default_prompt_model(&mut self, model: impl Into<String>) -> &mut Self {
97        if let Some(options) = self.default_options.as_mut() {
98            options.prompt_model = Some(model.into());
99        } else {
100            self.default_options = Some(Options {
101                prompt_model: Some(model.into()),
102                ..Default::default()
103            });
104        }
105        self
106    }
107
108    pub fn default_embed_model(&mut self, model: impl Into<String>) -> &mut Self {
109        if let Some(options) = self.default_options.as_mut() {
110            options.embed_model = Some(model.into());
111        } else {
112            self.default_options = Some(Options {
113                embed_model: Some(model.into()),
114                ..Default::default()
115            });
116        }
117        self
118    }
119
120    /// Sets the default dimensions for the `Dashscope` instance.
121    ///
122    /// # Parameters
123    /// - `dimensions`: The dimensions to set.
124    ///
125    /// # Returns
126    /// A mutable reference to the `DashscopeBuilder`.
127    pub fn default_dimensions(&mut self, dimensions: u16) -> &mut Self {
128        if let Some(options) = self.default_options.as_mut() {
129            options.dimensions = dimensions;
130        } else {
131            self.default_options = Some(Options {
132                dimensions,
133                ..Default::default()
134            });
135        }
136        self
137    }
138}
139
140#[cfg(test)]
141mod test {
142    use super::*;
143    #[test]
144    fn test_default_prompt_model() {
145        let openai = Dashscope::builder()
146            .default_prompt_model("qwen-long")
147            .build()
148            .unwrap();
149        assert_eq!(
150            openai.default_options.prompt_model,
151            Some("qwen-long".to_string())
152        );
153
154        let openai = Dashscope::builder()
155            .default_prompt_model("qwen-turbo")
156            .build()
157            .unwrap();
158        assert_eq!(
159            openai.default_options.prompt_model,
160            Some("qwen-turbo".to_string())
161        );
162    }
163}