kernelx_core/providers/remote/openai/
mod.rs

1mod capabilities;
2mod models;
3mod utils;
4
5pub use models::OpenAIModels;
6pub use utils::IntoOpenAIChatMessage;
7
8use crate::{
9    capabilities::*,
10    models::{ModelConfig, ModelInfo},
11    providers::{Provider, ProviderBuilder},
12    Error, Result,
13};
14use async_openai::{config::OpenAIConfig, Client};
15use std::fmt;
16
17#[derive(Clone)]
18pub struct OpenAI {
19    client: Client<OpenAIConfig>,
20    models: Vec<ModelInfo>,
21    config: ModelConfig,
22}
23
24impl fmt::Debug for OpenAI {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        f.debug_struct("OpenAI")
27            .field("models", &self.models)
28            .field("config", &self.config)
29            .finish_non_exhaustive()
30    }
31}
32
33impl OpenAI {
34    pub fn builder() -> ProviderBuilder<Self> {
35        ProviderBuilder::new()
36    }
37
38    pub fn from_env() -> Result<Self> {
39        Self::builder()
40            .api_key(
41                std::env::var("OPENAI_API_KEY")
42                    .map_err(|_| Error::Config("OPENAI_API_KEY not set".into()))?,
43            )
44            .build()
45    }
46
47    fn default_models() -> Vec<ModelInfo> {
48        crate::models! {
49            OpenAIModels::Gpt4o.as_str() => [Capability::Complete, Capability::Chat, Capability::Structured],
50            OpenAIModels::Gpt4oMini.as_str() => [Capability::Chat, Capability::Complete, Capability::Structured],
51            OpenAIModels::Dalle3.as_str() => [Capability::ImageGen],
52            OpenAIModels::TextEmbedding3Small.as_str() => [Capability::Embed],
53        }
54    }
55}
56
57impl Provider for OpenAI {
58    fn models(&self) -> &[ModelInfo] {
59        &self.models
60    }
61
62    fn from_builder(mut builder: ProviderBuilder<Self>) -> Result<Self> {
63        let api_base = builder
64            .get_api_base()
65            .unwrap_or("https://api.openai.com/v1");
66
67        let config = OpenAIConfig::new()
68            .with_api_key(builder.get_api_key().unwrap())
69            .with_api_base(api_base);
70
71        Ok(Self {
72            client: Client::with_config(config),
73            models: builder.take_models().unwrap_or_else(Self::default_models),
74            config: builder.take_config().unwrap_or_default(),
75        })
76    }
77}
78
79impl HasCapability for OpenAI {
80    fn model_id(&self) -> &str {
81        "openai"
82    }
83
84    fn config(&self) -> &ModelConfig {
85        &self.config
86    }
87}