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;
7
8use crate::openai;
9
10pub mod config;
11
12/// The `OpenRouter` struct encapsulates an `OpenRouter` client and default options for embedding
13/// and prompt models. It uses the `Builder` pattern for flexible and customizable instantiation.
14///
15/// By default it will look for a `OPENROUTER_API_KEY` environment variable. Note that either a
16/// prompt model or embedding model always need to be set, either with
17/// [`OpenRouter::with_default_prompt_model`] or [`OpenRouter::with_default_embed_model`] or via the
18/// builder. You can find available models in the `OpenRouter` documentation.
19///
20/// Under the hood it uses [`async_openai`], with the `OpenRouter` openai compatible api. This means
21/// some features might not work as expected. See the `OpenRouter` documentation for details.
22pub type OpenRouter = openai::GenericOpenAI<OpenRouterConfig>;
23pub type OpenRouterBuilder = openai::GenericOpenAIBuilder<OpenRouterConfig>;
24pub type OpenRouterBuilderError = openai::GenericOpenAIBuilderError;
25pub use openai::{Options, OptionsBuilder, OptionsBuilderError};
26
27impl OpenRouter {
28    /// Creates a new `OpenRouterBuilder` for constructing `OpenRouter` instances.
29    pub fn builder() -> OpenRouterBuilder {
30        OpenRouterBuilder::default()
31    }
32}
33
34impl Default for OpenRouter {
35    fn default() -> Self {
36        Self::builder().build().unwrap()
37    }
38}
39
40#[cfg(test)]
41mod test {
42    use super::*;
43
44    #[test]
45    fn test_default_prompt_model() {
46        let openai = OpenRouter::builder()
47            .default_prompt_model("llama3.1")
48            .build()
49            .unwrap();
50        assert_eq!(
51            openai.default_options.prompt_model,
52            Some("llama3.1".to_string())
53        );
54    }
55
56    #[test]
57    fn test_default_models() {
58        let openrouter = OpenRouter::builder()
59            .default_prompt_model("llama3.1")
60            .build()
61            .unwrap();
62        assert_eq!(
63            openrouter.default_options.prompt_model,
64            Some("llama3.1".to_string())
65        );
66    }
67
68    #[test]
69    fn test_building_via_default_prompt_model() {
70        let mut client = OpenRouter::default();
71
72        assert!(client.default_options.prompt_model.is_none());
73
74        client.with_default_prompt_model("llama3.1");
75        assert_eq!(
76            client.default_options.prompt_model,
77            Some("llama3.1".to_string())
78        );
79    }
80}