use crate::core::{GenericProvider, HttpClient, Protocol};
use crate::protocols::ZhipuProtocol;
use crate::error::LlmConnectorError;
use std::collections::HashMap;
pub type ZhipuProvider = GenericProvider<ZhipuProtocol>;
pub fn zhipu(api_key: &str) -> Result<ZhipuProvider, LlmConnectorError> {
zhipu_with_config(api_key, false, None, None, None)
}
pub fn zhipu_openai_compatible(api_key: &str) -> Result<ZhipuProvider, LlmConnectorError> {
zhipu_with_config(api_key, true, None, None, None)
}
pub fn zhipu_with_config(
api_key: &str,
openai_compatible: bool,
base_url: Option<&str>,
timeout_secs: Option<u64>,
proxy: Option<&str>,
) -> Result<ZhipuProvider, LlmConnectorError> {
let protocol = if openai_compatible {
ZhipuProtocol::new_openai_compatible(api_key)
} else {
ZhipuProtocol::new(api_key)
};
let client = HttpClient::with_config(
base_url.unwrap_or("https://open.bigmodel.cn"),
timeout_secs,
proxy,
)?;
let auth_headers: HashMap<String, String> = protocol.auth_headers().into_iter().collect();
let client = client.with_headers(auth_headers);
Ok(GenericProvider::new(protocol, client))
}
pub fn zhipu_default(api_key: &str) -> Result<ZhipuProvider, LlmConnectorError> {
zhipu_openai_compatible(api_key)
}
pub fn zhipu_with_timeout(
api_key: &str,
timeout_secs: u64,
) -> Result<ZhipuProvider, LlmConnectorError> {
zhipu_with_config(api_key, true, None, Some(timeout_secs), None)
}
pub fn zhipu_enterprise(
api_key: &str,
enterprise_endpoint: &str,
) -> Result<ZhipuProvider, LlmConnectorError> {
zhipu_with_config(api_key, true, Some(enterprise_endpoint), None, None)
}
pub fn validate_zhipu_key(api_key: &str) -> bool {
!api_key.is_empty() && api_key.len() > 10
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_zhipu_provider_creation() {
let provider = zhipu("test-key");
assert!(provider.is_ok());
let provider = provider.unwrap();
assert_eq!(provider.protocol().name(), "zhipu");
assert_eq!(provider.protocol().api_key(), "test-key");
assert!(!provider.protocol().is_openai_compatible());
}
#[test]
fn test_zhipu_openai_compatible() {
let provider = zhipu_openai_compatible("test-key");
assert!(provider.is_ok());
let provider = provider.unwrap();
assert_eq!(provider.protocol().name(), "zhipu");
assert!(provider.protocol().is_openai_compatible());
}
#[test]
fn test_zhipu_with_config() {
let provider = zhipu_with_config(
"test-key",
true,
Some("https://custom.bigmodel.cn"),
Some(60),
None
);
assert!(provider.is_ok());
let provider = provider.unwrap();
assert_eq!(provider.client().base_url(), "https://custom.bigmodel.cn");
assert!(provider.protocol().is_openai_compatible());
}
#[test]
fn test_zhipu_default() {
let provider = zhipu_default("test-key");
assert!(provider.is_ok());
let provider = provider.unwrap();
assert!(provider.protocol().is_openai_compatible());
}
#[test]
fn test_zhipu_with_timeout() {
let provider = zhipu_with_timeout("test-key", 120);
assert!(provider.is_ok());
}
#[test]
fn test_zhipu_enterprise() {
let provider = zhipu_enterprise("test-key", "https://enterprise.bigmodel.cn");
assert!(provider.is_ok());
let provider = provider.unwrap();
assert_eq!(provider.client().base_url(), "https://enterprise.bigmodel.cn");
}
#[test]
fn test_validate_zhipu_key() {
assert!(validate_zhipu_key("valid-test-key"));
assert!(validate_zhipu_key("another-valid-key-12345"));
assert!(!validate_zhipu_key("short"));
assert!(!validate_zhipu_key(""));
}
}