use super::BindingValue;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(tag = "service", rename_all = "lowercase")]
pub enum AiBinding {
Bedrock(BedrockAiBinding),
Vertex(VertexAiBinding),
Foundry(FoundryAiBinding),
#[serde(rename = "external-ai")]
External(ExternalAiBinding),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(rename_all = "camelCase")]
pub struct BedrockAiBinding {
pub region: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(rename_all = "camelCase")]
pub struct VertexAiBinding {
pub project: String,
pub location: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(rename_all = "camelCase")]
pub struct FoundryAiBinding {
pub endpoint: String,
pub account: String,
}
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(rename_all = "camelCase")]
pub struct ExternalAiBinding {
pub provider: String,
pub api_key: BindingValue<String>,
}
impl std::fmt::Debug for ExternalAiBinding {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExternalAiBinding")
.field("provider", &self.provider)
.field("api_key", &"<redacted>")
.finish()
}
}
impl AiBinding {
pub const LOCAL_API_KEY_ENV: &'static str = "OPENAI_API_KEY";
pub const LOCAL_DEFAULT_PROVIDER: &'static str = "openai";
pub fn bedrock(region: impl Into<String>) -> Self {
Self::Bedrock(BedrockAiBinding {
region: region.into(),
})
}
pub fn vertex(project: impl Into<String>, location: impl Into<String>) -> Self {
Self::Vertex(VertexAiBinding {
project: project.into(),
location: location.into(),
})
}
pub fn foundry(endpoint: impl Into<String>, account: impl Into<String>) -> Self {
Self::Foundry(FoundryAiBinding {
endpoint: endpoint.into(),
account: account.into(),
})
}
pub fn external(
provider: impl Into<String>,
api_key: impl Into<BindingValue<String>>,
) -> Self {
Self::External(ExternalAiBinding {
provider: provider.into(),
api_key: api_key.into(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bedrock_binding_roundtrip() {
let binding = AiBinding::bedrock("us-east-1");
let json = serde_json::to_string(&binding).unwrap();
assert!(json.contains(r#""service":"bedrock""#));
let deserialized: AiBinding = serde_json::from_str(&json).unwrap();
assert_eq!(binding, deserialized);
}
#[test]
fn test_vertex_binding_roundtrip() {
let binding = AiBinding::vertex("my-project", "us-central1");
let json = serde_json::to_string(&binding).unwrap();
assert!(json.contains(r#""service":"vertex""#));
let deserialized: AiBinding = serde_json::from_str(&json).unwrap();
assert_eq!(binding, deserialized);
}
#[test]
fn test_foundry_binding_roundtrip() {
let binding = AiBinding::foundry("https://my-foundry.openai.azure.com", "my-subscription");
let json = serde_json::to_value(&binding).unwrap();
let json_str = json.to_string();
assert!(json_str.contains(r#""service":"foundry""#));
assert!(
!json_str.contains("instantAccess"),
"foundry binding must not serialize instant_access"
);
let deserialized: AiBinding = serde_json::from_value(json).unwrap();
assert_eq!(binding, deserialized);
}
#[test]
fn test_external_binding_roundtrip() {
let binding = AiBinding::external("openai", "sk-test-key");
let json = serde_json::to_value(&binding).unwrap();
let json_str = json.to_string();
assert!(json_str.contains(r#""apiKey""#), "external binding must serialize apiKey in camelCase");
assert_eq!(
json,
serde_json::json!({
"service": "external-ai",
"provider": "openai",
"apiKey": "sk-test-key",
})
);
let deserialized: AiBinding =
serde_json::from_value(json).expect("external binding should round-trip");
assert_eq!(binding, deserialized);
}
}