use crate::core::{GenericProvider, HttpClient, Protocol};
use crate::error::LlmConnectorError;
use std::collections::HashMap;
pub struct ProviderBuilder<P: Protocol> {
protocol: P,
base_url: String,
timeout_secs: Option<u64>,
proxy: Option<String>,
extra_headers: HashMap<String, String>,
}
impl<P: Protocol> ProviderBuilder<P> {
pub fn new(protocol: P, base_url: &str) -> Self {
Self {
protocol,
base_url: base_url.to_string(),
timeout_secs: None,
proxy: None,
extra_headers: HashMap::new(),
}
}
pub fn timeout(mut self, secs: u64) -> Self {
self.timeout_secs = Some(secs);
self
}
pub fn proxy(mut self, proxy: &str) -> Self {
self.proxy = Some(proxy.to_string());
self
}
pub fn header(mut self, key: &str, value: &str) -> Self {
self.extra_headers
.insert(key.to_string(), value.to_string());
self
}
pub fn build(self) -> Result<GenericProvider<P>, LlmConnectorError> {
let client =
HttpClient::with_config(&self.base_url, self.timeout_secs, self.proxy.as_deref())?;
let mut headers: HashMap<String, String> =
self.protocol.auth_headers().into_iter().collect();
headers.extend(self.extra_headers);
let client = client.with_headers(headers);
Ok(GenericProvider::new(self.protocol, client))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::protocols::OpenAIProtocol;
#[test]
fn test_builder_basic() {
let provider =
ProviderBuilder::new(OpenAIProtocol::new("sk-test"), "https://api.openai.com").build();
assert!(provider.is_ok());
}
#[test]
fn test_builder_with_timeout() {
let provider =
ProviderBuilder::new(OpenAIProtocol::new("sk-test"), "https://api.openai.com")
.timeout(60)
.build();
assert!(provider.is_ok());
}
#[test]
fn test_builder_with_proxy() {
let provider =
ProviderBuilder::new(OpenAIProtocol::new("sk-test"), "https://api.openai.com")
.proxy("http://proxy:8080")
.build();
assert!(provider.is_ok());
}
#[test]
fn test_builder_with_headers() {
let provider =
ProviderBuilder::new(OpenAIProtocol::new("sk-test"), "https://api.openai.com")
.header("X-Custom-Header", "value")
.header("X-Another-Header", "value2")
.build();
assert!(provider.is_ok());
}
#[test]
fn test_builder_chain() {
let provider =
ProviderBuilder::new(OpenAIProtocol::new("sk-test"), "https://api.openai.com")
.timeout(60)
.proxy("http://proxy:8080")
.header("X-Custom-Header", "value")
.build();
assert!(provider.is_ok());
}
}