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}