llm/providers/openai_compatible/
generic.rs

1use async_openai::{Client, config::OpenAIConfig};
2use schemars::Schema;
3
4use crate::provider::get_context_window;
5use crate::tool_schema::normalize_for_moonshot;
6use crate::{
7    Context, LlmError, LlmModel, LlmResponseStream, ProviderAuthMode, ProviderConnectionConfig, Result,
8    StreamingModelProvider,
9};
10
11use super::{AetherOpenAiConfig, build_chat_request, create_custom_stream_generic};
12
13/// Configuration for an OpenAI-compatible provider.
14///
15/// Each provider that uses the standard `build_chat_request → create_custom_stream_generic`
16/// flow differs only in these constants.
17pub struct ProviderConfig {
18    pub api_base: Option<&'static str>,
19    pub env_var: &'static str,
20    pub default_model: &'static str,
21    pub prefix: &'static str,
22    pub display_name: &'static str,
23    pub tool_schema_transform: Option<fn(&mut Schema)>,
24}
25
26pub const DEEPSEEK: ProviderConfig = ProviderConfig {
27    api_base: Some("https://api.deepseek.com"),
28    env_var: "DEEPSEEK_API_KEY",
29    default_model: "deepseek-chat",
30    prefix: "deepseek",
31    display_name: "DeepSeek",
32    tool_schema_transform: None,
33};
34
35pub const MOONSHOT: ProviderConfig = ProviderConfig {
36    api_base: Some("https://api.moonshot.ai/v1"),
37    env_var: "MOONSHOT_API_KEY",
38    default_model: "moonshot-v1-8k",
39    prefix: "moonshot",
40    display_name: "Moonshot",
41    tool_schema_transform: Some(normalize_for_moonshot),
42};
43
44pub const ZAI: ProviderConfig = ProviderConfig {
45    api_base: Some("https://api.z.ai/api/coding/paas/v4"),
46    env_var: "ZAI_API_KEY",
47    default_model: "GLM-4.6",
48    prefix: "zai",
49    display_name: "Z.ai",
50    tool_schema_transform: None,
51};
52
53pub const AZURE_FOUNDRY: ProviderConfig = ProviderConfig {
54    api_base: None,
55    env_var: "AZURE_OPENAI_API_KEY",
56    default_model: "gpt-5.5",
57    prefix: "azure-foundry",
58    display_name: "Microsoft Foundry",
59    tool_schema_transform: None,
60};
61
62pub const FIREWORKS: ProviderConfig = ProviderConfig {
63    api_base: Some("https://api.fireworks.ai/inference/v1"),
64    env_var: "FIREWORKS_API_KEY",
65    default_model: "accounts/fireworks/models/glm-5p1",
66    prefix: "fireworks",
67    display_name: "Fireworks AI",
68    tool_schema_transform: None,
69};
70
71/// A generic provider for APIs that are fully OpenAI-compatible.
72pub struct GenericOpenAiProvider {
73    client: Client<AetherOpenAiConfig>,
74    model: String,
75    request_model: Option<String>,
76    config: &'static ProviderConfig,
77}
78
79impl GenericOpenAiProvider {
80    pub fn from_env(config: &'static ProviderConfig) -> Result<Self> {
81        Self::from_env_with_connection(config, ProviderConnectionConfig::default())
82    }
83
84    pub fn from_env_with_connection(
85        config: &'static ProviderConfig,
86        connection: ProviderConnectionConfig,
87    ) -> Result<Self> {
88        let api_key = match connection.auth_mode {
89            ProviderAuthMode::Default => {
90                std::env::var(config.env_var).map_err(|_| LlmError::MissingApiKey(config.env_var.to_string()))?
91            }
92            ProviderAuthMode::None => String::new(),
93        };
94        Self::new_with_connection(api_key, config, connection)
95    }
96
97    pub fn new(api_key: String, config: &'static ProviderConfig) -> Result<Self> {
98        Self::new_with_connection(api_key, config, ProviderConnectionConfig::default())
99    }
100
101    pub fn new_with_connection(
102        api_key: String,
103        config: &'static ProviderConfig,
104        connection: ProviderConnectionConfig,
105    ) -> Result<Self> {
106        let api_base = connection
107            .base_url
108            .or_else(|| config.api_base.map(str::to_string))
109            .ok_or_else(|| LlmError::MissingProviderUrl { provider: config.prefix.to_string() })?
110            .trim_end_matches('/')
111            .to_string();
112        let openai_config = OpenAIConfig::new().with_api_key(api_key).with_api_base(api_base);
113        let openai_config = AetherOpenAiConfig::new(openai_config, connection.auth_mode);
114
115        Ok(Self {
116            client: Client::with_config(openai_config),
117            model: config.default_model.to_string(),
118            request_model: connection.request_model,
119            config,
120        })
121    }
122
123    pub fn with_model(mut self, model: &str) -> Self {
124        self.model = model.to_string();
125        self
126    }
127}
128
129impl StreamingModelProvider for GenericOpenAiProvider {
130    fn model(&self) -> Option<LlmModel> {
131        format!("{}:{}", self.config.prefix, self.model).parse().ok()
132    }
133
134    fn context_window(&self) -> Option<u32> {
135        get_context_window(self.config.prefix, &self.model)
136    }
137
138    fn stream_response(&self, context: &Context) -> LlmResponseStream {
139        let request = match build_chat_request(
140            self.request_model.as_deref().unwrap_or(&self.model),
141            context,
142            self.config.tool_schema_transform,
143        ) {
144            Ok(req) => req,
145            Err(e) => return Box::pin(async_stream::stream! { yield Err(e); }),
146        };
147        create_custom_stream_generic(&self.client, request)
148    }
149
150    fn display_name(&self) -> String {
151        format!("{} ({})", self.config.display_name, self.model)
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use futures::StreamExt;
158
159    use super::*;
160    use crate::providers::test_capture_server::CaptureServer;
161    use crate::types::IsoString;
162    use crate::{ChatMessage, ContentBlock};
163
164    #[test]
165    fn azure_foundry_requires_a_configured_url() {
166        let Err(error) = GenericOpenAiProvider::new("key".to_string(), &AZURE_FOUNDRY) else {
167            panic!("Azure Foundry must require a URL");
168        };
169        assert!(matches!(error, LlmError::MissingProviderUrl { provider } if provider == "azure-foundry"));
170    }
171
172    #[tokio::test]
173    async fn request_model_routes_the_request_without_changing_catalog_identity() {
174        let mut server = CaptureServer::start().await;
175        let provider = GenericOpenAiProvider::new_with_connection(
176            "key".to_string(),
177            &AZURE_FOUNDRY,
178            ProviderConnectionConfig {
179                base_url: Some(format!("{}/", server.base_url)),
180                auth_mode: ProviderAuthMode::None,
181                request_model: Some("production-coding".to_string()),
182                ..Default::default()
183            },
184        )
185        .unwrap()
186        .with_model("gpt-5.5");
187        let context = Context::new(
188            vec![ChatMessage::User { content: vec![ContentBlock::text("Hello")], timestamp: IsoString::now() }],
189            vec![],
190        );
191
192        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
193        let captured = server.captured().await;
194
195        assert!(!responses.is_empty());
196        assert_eq!(captured.path, "/chat/completions");
197        assert_eq!(captured.body["model"], "production-coding");
198        assert_eq!(captured.body["stream"], true);
199        assert_eq!(captured.body["stream_options"]["include_usage"], true);
200        assert!(captured.headers.get("authorization").is_none());
201        assert_eq!(provider.model().unwrap().to_string(), "azure-foundry:gpt-5.5");
202        assert_eq!(provider.display_name(), "Microsoft Foundry (gpt-5.5)");
203    }
204}
llm/providers/openai_compatible/generic.rs

llm/providers/openai_compatible/
generic.rs
