use crate::providers::{
anthropic::AnthropicConfig, gemini::GeminiConfig, openai::OpenAIConfig,
stakpak::StakpakProviderConfig,
};
#[cfg(feature = "bedrock")]
use crate::providers::bedrock::BedrockConfig;
#[derive(Debug, Clone, Default)]
pub struct ClientConfig {
pub default_temperature: Option<f32>,
pub default_max_tokens: Option<u32>,
pub timeout_seconds: Option<u64>,
}
impl ClientConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.default_temperature = Some(temperature);
self
}
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
self.default_max_tokens = Some(max_tokens);
self
}
pub fn with_timeout(mut self, seconds: u64) -> Self {
self.timeout_seconds = Some(seconds);
self
}
}
#[derive(Debug, Default)]
pub struct InferenceConfig {
pub(crate) openai_config: Option<OpenAIConfig>,
pub(crate) anthropic_config: Option<AnthropicConfig>,
pub(crate) gemini_config: Option<GeminiConfig>,
pub(crate) stakpak_config: Option<StakpakProviderConfig>,
#[cfg(feature = "bedrock")]
pub(crate) bedrock_config: Option<BedrockConfig>,
pub(crate) client_config: ClientConfig,
}
impl InferenceConfig {
pub fn new() -> Self {
Self::default()
}
pub fn openai(mut self, api_key: impl Into<String>, base_url: Option<String>) -> Self {
let mut config = OpenAIConfig::new(api_key);
if let Some(url) = base_url {
config = config.with_base_url(url);
}
self.openai_config = Some(config);
self
}
pub fn openai_config(mut self, config: OpenAIConfig) -> Self {
self.openai_config = Some(config);
self
}
pub fn anthropic(mut self, api_key: impl Into<String>, base_url: Option<String>) -> Self {
let mut config = AnthropicConfig::new(api_key);
if let Some(url) = base_url {
config = config.with_base_url(url);
}
self.anthropic_config = Some(config);
self
}
pub fn anthropic_config(mut self, config: AnthropicConfig) -> Self {
self.anthropic_config = Some(config);
self
}
pub fn gemini(mut self, api_key: impl Into<String>, base_url: Option<String>) -> Self {
let mut config = GeminiConfig::new(api_key);
if let Some(url) = base_url {
config = config.with_base_url(url);
}
self.gemini_config = Some(config);
self
}
pub fn gemini_config(mut self, config: GeminiConfig) -> Self {
self.gemini_config = Some(config);
self
}
pub fn temperature(mut self, temperature: f32) -> Self {
self.client_config.default_temperature = Some(temperature);
self
}
pub fn max_tokens(mut self, max_tokens: u32) -> Self {
self.client_config.default_max_tokens = Some(max_tokens);
self
}
pub fn timeout(mut self, seconds: u64) -> Self {
self.client_config.timeout_seconds = Some(seconds);
self
}
pub fn stakpak(mut self, api_key: impl Into<String>, base_url: Option<String>) -> Self {
let mut config = StakpakProviderConfig::new(api_key);
if let Some(url) = base_url {
config = config.with_base_url(url);
}
self.stakpak_config = Some(config);
self
}
pub fn stakpak_config(mut self, config: StakpakProviderConfig) -> Self {
self.stakpak_config = Some(config);
self
}
#[cfg(feature = "bedrock")]
pub fn bedrock(mut self, region: impl Into<String>) -> Self {
self.bedrock_config = Some(BedrockConfig::new(region));
self
}
#[cfg(feature = "bedrock")]
pub fn bedrock_from_env(mut self) -> Self {
self.bedrock_config = Some(BedrockConfig::from_env());
self
}
#[cfg(feature = "bedrock")]
pub fn bedrock_config(mut self, config: BedrockConfig) -> Self {
self.bedrock_config = Some(config);
self
}
}