use crate::{
ZaiError, ZaiResult,
client::{
endpoints::{ApiBase, EndpointConfig},
http::HttpClientConfig,
},
};
#[derive(Debug, Clone, Default)]
pub struct ZaiConfig {
pub api_key: String,
pub endpoints: EndpointConfig,
pub http: HttpClientConfig,
pub reqwest: Option<reqwest::Client>,
}
impl ZaiConfig {
pub fn builder() -> ZaiConfigBuilder {
ZaiConfigBuilder {
config: ZaiConfig::default(),
}
}
pub fn from_env() -> ZaiResult<Self> {
let api_key = std::env::var("ZHIPU_API_KEY").map_err(|_| ZaiError::AuthError {
code: 1001,
message: "ZHIPU_API_KEY environment variable not set".to_string(),
})?;
Ok(Self {
api_key,
..ZaiConfig::default()
})
}
pub fn realtime_url(&self) -> String {
self.endpoints.url(&ApiBase::Realtime, "")
}
pub fn paas_v4_url(&self, path: &str) -> String {
self.endpoints.url(&ApiBase::PaasV4, path)
}
pub fn llm_application_url(&self, path: &str) -> String {
self.endpoints.url(&ApiBase::LlmApplication, path)
}
pub fn monitor_url(&self, path: &str) -> String {
self.endpoints.url(&ApiBase::Monitor, path)
}
}
#[derive(Debug, Clone, Default)]
pub struct ZaiConfigBuilder {
config: ZaiConfig,
}
impl ZaiConfigBuilder {
pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
self.config.api_key = api_key.into();
self
}
pub fn endpoint_config(mut self, endpoints: EndpointConfig) -> Self {
self.config.endpoints = endpoints;
self
}
pub fn realtime_base(mut self, base: impl Into<String>) -> Self {
self.config.endpoints = self.config.endpoints.with_realtime_base(base);
self
}
pub fn monitor_base(mut self, base: impl Into<String>) -> Self {
self.config.endpoints = self.config.endpoints.with_monitor_base(base);
self
}
pub fn http_config(mut self, http: HttpClientConfig) -> Self {
self.config.http = http;
self
}
pub fn reqwest_client(mut self, client: reqwest::Client) -> Self {
self.config.reqwest = Some(client);
self
}
pub fn build(self) -> ZaiResult<ZaiConfig> {
if self.config.api_key.is_empty() {
return Err(ZaiError::ApiError {
code: 1200,
message: "ZaiConfig requires an api_key".to_string(),
});
}
Ok(self.config)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::endpoints::REALTIME_BASE;
#[test]
fn builder_requires_api_key() {
assert!(ZaiConfig::builder().build().is_err());
let cfg = ZaiConfig::builder()
.api_key("abcdefghij.0123456789abcdef")
.build()
.unwrap();
assert_eq!(cfg.api_key, "abcdefghij.0123456789abcdef");
}
#[test]
fn realtime_url_uses_official_endpoint() {
let cfg = ZaiConfig::builder()
.api_key("abcdefghij.0123456789abcdef")
.build()
.unwrap();
assert_eq!(cfg.realtime_url(), REALTIME_BASE);
}
#[test]
fn custom_realtime_base_overrides() {
let cfg = ZaiConfig::builder()
.api_key("abcdefghij.0123456789abcdef")
.realtime_base("wss://custom.example.com/realtime")
.build()
.unwrap();
assert_eq!(cfg.realtime_url(), "wss://custom.example.com/realtime");
}
}