swiftide_integrations/open_router/
mod.rs

1//! This module provides integration with `OpenRouter`'s API, enabling the use of language models
2//! and embeddings within the Swiftide project. It includes the `OpenRouter` struct for managing API
3//! clients and default options for embedding and prompt models. The module is conditionally
4//! compiled based on the "openrouter" feature flag.
5
6use config::OpenRouterConfig;
7use derive_builder::Builder;
8use std::sync::Arc;
9
10pub mod chat_completion;
11pub mod config;
12pub mod simple_prompt;
13
14/// The `OpenRouter` struct encapsulates an `OpenRouter` client and default options for embedding
15/// and prompt models. It uses the `Builder` pattern for flexible and customizable instantiation.
16///
17/// By default it will look for a `OPENROUTER_API_KEY` environment variable. Note that either a
18/// prompt model or embedding model always need to be set, either with
19/// [`OpenRouter::with_default_prompt_model`] or [`OpenRouter::with_default_embed_model`] or via the
20/// builder. You can find available models in the `OpenRouter` documentation.
21///
22/// Under the hood it uses [`async_openai`], with the `OpenRouter` openai compatible api. This means
23/// some features might not work as expected. See the `OpenRouter` documentation for details.
24#[derive(Debug, Builder, Clone)]
25#[builder(setter(into, strip_option))]
26pub struct OpenRouter {
27    /// The `OpenRouter` client, wrapped in an `Arc` for thread-safe reference counting.
28    #[builder(default = "default_client()", setter(custom))]
29    client: Arc<async_openai::Client<OpenRouterConfig>>,
30    /// Default options for the embedding and prompt models.
31    #[builder(default)]
32    default_options: Options,
33}
34
35impl Default for OpenRouter {
36    fn default() -> Self {
37        Self {
38            client: default_client(),
39            default_options: Options::default(),
40        }
41    }
42}
43
44/// The `Options` struct holds configuration options for the `OpenRouter` client.
45/// It includes optional fields for specifying the embedding and prompt models.
46#[derive(Debug, Default, Clone, Builder)]
47#[builder(setter(into, strip_option))]
48pub struct Options {
49    /// The default prompt model to use, if specified.
50    #[builder(default)]
51    pub prompt_model: Option<String>,
52}
53
54impl Options {
55    /// Creates a new `OptionsBuilder` for constructing `Options` instances.
56    pub fn builder() -> OptionsBuilder {
57        OptionsBuilder::default()
58    }
59}
60
61impl OpenRouter {
62    /// Creates a new `OpenRouterBuilder` for constructing `OpenRouter` instances.
63    pub fn builder() -> OpenRouterBuilder {
64        OpenRouterBuilder::default()
65    }
66
67    /// Sets a default prompt model to use when prompting
68    pub fn with_default_prompt_model(&mut self, model: impl Into<String>) -> &mut Self {
69        self.default_options = Options {
70            prompt_model: Some(model.into()),
71        };
72        self
73    }
74}
75
76impl OpenRouterBuilder {
77    /// Sets the `OpenRouter` client for the `OpenRouter` instance.
78    ///
79    /// # Parameters
80    /// - `client`: The `OpenRouter` client to set.
81    ///
82    /// # Returns
83    /// A mutable reference to the `OpenRouterBuilder`.
84    pub fn client(&mut self, client: async_openai::Client<OpenRouterConfig>) -> &mut Self {
85        self.client = Some(Arc::new(client));
86        self
87    }
88
89    /// Sets the default prompt model for the `OpenRouter` instance.
90    ///
91    /// # Parameters
92    /// - `model`: The prompt model to set.
93    ///
94    /// # Returns
95    /// A mutable reference to the `OpenRouterBuilder`.
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            });
103        }
104        self
105    }
106}
107
108fn default_client() -> Arc<async_openai::Client<OpenRouterConfig>> {
109    Arc::new(async_openai::Client::with_config(
110        OpenRouterConfig::default(),
111    ))
112}
113
114#[cfg(test)]
115mod test {
116    use super::*;
117
118    #[test]
119    fn test_default_prompt_model() {
120        let openai = OpenRouter::builder()
121            .default_prompt_model("llama3.1")
122            .build()
123            .unwrap();
124        assert_eq!(
125            openai.default_options.prompt_model,
126            Some("llama3.1".to_string())
127        );
128    }
129
130    #[test]
131    fn test_default_models() {
132        let openrouter = OpenRouter::builder()
133            .default_prompt_model("llama3.1")
134            .build()
135            .unwrap();
136        assert_eq!(
137            openrouter.default_options.prompt_model,
138            Some("llama3.1".to_string())
139        );
140    }
141
142    #[test]
143    fn test_building_via_default_prompt_model() {
144        let mut client = OpenRouter::default();
145
146        assert!(client.default_options.prompt_model.is_none());
147
148        client.with_default_prompt_model("llama3.1");
149        assert_eq!(
150            client.default_options.prompt_model,
151            Some("llama3.1".to_string())
152        );
153    }
154}
swiftide_integrations/open_router/mod.rs

swiftide_integrations/open_router/
mod.rs
