diff --git a/crates/omnigraph-cluster/src/types.rs b/crates/omnigraph-cluster/src/types.rs index e44e2f4..9560673 100644 --- a/crates/omnigraph-cluster/src/types.rs +++ b/crates/omnigraph-cluster/src/types.rs @@ -400,6 +400,54 @@ pub(crate) struct GraphConfig { pub(crate) schema: PathBuf, #[serde(default)] pub(crate) queries: QueriesDecl, + /// Optional per-graph embedding provider profile (RFC-012 Phase 5). + #[serde(default)] + pub(crate) embeddings: Option, +} + +/// A graph's embedding provider profile (RFC-012 Phase 5). `provider`/`base_url`/ +/// `model` default exactly as the engine's `EmbeddingConfig::from_env` does; +/// `api_key` is a `${NAME}` env reference resolved at serving boot, never an +/// inline secret. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct EmbeddingProfile { + #[serde(default)] + pub provider: Option, + #[serde(default)] + pub base_url: Option, + #[serde(default)] + pub model: Option, + pub api_key: String, +} + +impl EmbeddingProfile { + /// Resolve into an engine `EmbeddingConfig`, reading the `${NAME}` api-key + /// reference from process env. Errors if `api_key` is not a `${NAME}` + /// reference or the named var is unset. + pub fn resolve(&self) -> Result { + let api_key = resolve_secret_ref(&self.api_key)?; + omnigraph::embedding::EmbeddingConfig::from_parts( + self.provider.as_deref(), + self.base_url.clone(), + self.model.clone(), + api_key, + ) + .map_err(|e| e.to_string()) + } +} + +/// Resolve a `${NAME}` secret reference from process env. Rejects an inline value +/// (anything not wrapped in `${…}`) so secrets never sit in the cluster config. +fn resolve_secret_ref(value: &str) -> Result { + let name = value + .trim() + .strip_prefix("${") + .and_then(|s| s.strip_suffix('}')) + .ok_or_else(|| { + format!("embedding api_key must be a ${{NAME}} env reference, got '{}'", value.trim()) + })?; + std::env::var(name).map_err(|_| format!("embedding api_key env var '{name}' is not set")) } /// How a graph declares its stored queries. Terraform-style: the `.gq` @@ -518,3 +566,62 @@ pub(crate) struct SweepOutcome { /// files are rewritten with consumed_at only after the state write lands. pub(crate) consumed_approvals: Vec, } + +#[cfg(test)] +mod embedding_profile_tests { + use super::EmbeddingProfile; + + #[test] + fn resolves_secret_from_env_and_applies_defaults() { + // SAFETY: a unique var name, no concurrent reader. + unsafe { std::env::set_var("OG_TEST_EMBED_KEY_A", "secret-x") }; + let profile = EmbeddingProfile { + provider: Some("openai-compatible".to_string()), + base_url: None, + model: Some("m".to_string()), + api_key: "${OG_TEST_EMBED_KEY_A}".to_string(), + }; + let config = profile.resolve().unwrap(); + assert_eq!(config.api_key, "secret-x"); + assert_eq!(config.model, "m"); + unsafe { std::env::remove_var("OG_TEST_EMBED_KEY_A") }; + } + + #[test] + fn rejects_inline_api_key() { + let profile = EmbeddingProfile { + provider: None, + base_url: None, + model: None, + api_key: "sk-inline".to_string(), + }; + let err = profile.resolve().unwrap_err(); + assert!(err.contains("${NAME}"), "got: {err}"); + } + + #[test] + fn errors_on_unset_secret() { + let profile = EmbeddingProfile { + provider: None, + base_url: None, + model: None, + api_key: "${OG_TEST_DEFINITELY_UNSET_VAR}".to_string(), + }; + let err = profile.resolve().unwrap_err(); + assert!(err.contains("not set"), "got: {err}"); + } + + #[test] + fn rejects_unknown_provider() { + unsafe { std::env::set_var("OG_TEST_EMBED_KEY_B", "x") }; + let profile = EmbeddingProfile { + provider: Some("cohere".to_string()), + base_url: None, + model: None, + api_key: "${OG_TEST_EMBED_KEY_B}".to_string(), + }; + let err = profile.resolve().unwrap_err(); + assert!(err.contains("unknown embedding provider"), "got: {err}"); + unsafe { std::env::remove_var("OG_TEST_EMBED_KEY_B") }; + } +} diff --git a/crates/omnigraph/src/embedding.rs b/crates/omnigraph/src/embedding.rs index abb321c..0210fe2 100644 --- a/crates/omnigraph/src/embedding.rs +++ b/crates/omnigraph/src/embedding.rs @@ -73,45 +73,12 @@ impl EmbeddingConfig { return Ok(Self::mock()); } - // The default base URL and model depend on the provider *alias*, not just - // the wire shape: `openai-compatible` (and the unset default) point at the - // OpenRouter gateway, while `openai` points at OpenAI's own host. - // Each provider alias yields its full default profile — base URL, model, - // and the ordered api-key envs to try — so every alias-dependent default - // lives in one place. (Dispatching the key on the `Provider` enum would - // collapse `openai` and `openai-compatible` and send an OpenRouter key to - // OpenAI's host.) `openai` targets OpenAI directly and takes only - // `OPENAI_API_KEY`; the OpenRouter gateway default prefers - // `OPENROUTER_API_KEY`, falling back to `OPENAI_API_KEY`. - let (provider, default_base, default_model, key_envs): (Provider, &str, &str, &[&str]) = - match env_string("OMNIGRAPH_EMBED_PROVIDER").as_deref() { - None | Some("openai-compatible") => ( - Provider::OpenAiCompatible, - DEFAULT_OPENROUTER_BASE_URL, - DEFAULT_OPENROUTER_MODEL, - &["OPENROUTER_API_KEY", "OPENAI_API_KEY"], - ), - Some("openai") => ( - Provider::OpenAiCompatible, - DEFAULT_OPENAI_BASE_URL, - DEFAULT_OPENAI_MODEL, - &["OPENAI_API_KEY"], - ), - Some("gemini") => ( - Provider::Gemini, - DEFAULT_GEMINI_BASE_URL, - DEFAULT_GEMINI_MODEL, - &["GEMINI_API_KEY"], - ), - Some("mock") => return Ok(Self::mock()), - Some(other) => { - return Err(OmniError::manifest_internal(format!( - "unknown OMNIGRAPH_EMBED_PROVIDER '{}' (expected openai-compatible|openai|gemini|mock)", - other - ))); - } - }; + let alias = env_string("OMNIGRAPH_EMBED_PROVIDER"); + if alias.as_deref() == Some("mock") { + return Ok(Self::mock()); + } + let (provider, default_base, default_model, key_envs) = provider_profile(alias.as_deref())?; let base_url = env_string("OMNIGRAPH_EMBED_BASE_URL") .unwrap_or_else(|| default_base.to_string()) .trim_end_matches('/') @@ -123,7 +90,7 @@ impl EmbeddingConfig { OmniError::manifest_internal(format!( "{} is required for the {} embedding provider", key_envs.join(" or "), - env_string("OMNIGRAPH_EMBED_PROVIDER").as_deref().unwrap_or("openai-compatible") + alias.as_deref().unwrap_or("openai-compatible") )) })?; @@ -135,6 +102,33 @@ impl EmbeddingConfig { }) } + /// Build a config from explicit parts — the cluster `embeddings` profile path + /// (RFC-012 Phase 5). `provider`/`base_url`/`model` default exactly as + /// `from_env` does (shared `provider_profile`); `api_key` is already resolved + /// (the cluster path resolves a `${NAME}` ref before calling this). + pub fn from_parts( + provider: Option<&str>, + base_url: Option, + model: Option, + api_key: String, + ) -> Result { + if provider == Some("mock") { + return Ok(Self::mock()); + } + let (provider, default_base, default_model, _key_envs) = provider_profile(provider)?; + let base_url = base_url + .unwrap_or_else(|| default_base.to_string()) + .trim_end_matches('/') + .to_string(); + let model = model.unwrap_or_else(|| default_model.to_string()); + Ok(Self { + provider, + model, + base_url, + api_key, + }) + } + fn mock() -> Self { Self { provider: Provider::Mock, @@ -570,6 +564,43 @@ fn parse_openai_error_message(body: &str) -> Option { .filter(|msg| !msg.trim().is_empty()) } +/// Map a provider alias to `(provider, default base URL, default model, ordered +/// api-key envs)`. Shared by `from_env` and `from_parts` so both apply identical +/// defaults: `openai-compatible`/unset → the OpenRouter gateway, `openai` → +/// OpenAI's own host. `mock` is handled by callers before this is reached. The +/// `Provider` enum alone would collapse the two openai aliases, so the alias +/// (not the enum) determines the key-env order here. +fn provider_profile( + alias: Option<&str>, +) -> Result<(Provider, &'static str, &'static str, &'static [&'static str])> { + Ok(match alias { + None | Some("openai-compatible") => ( + Provider::OpenAiCompatible, + DEFAULT_OPENROUTER_BASE_URL, + DEFAULT_OPENROUTER_MODEL, + &["OPENROUTER_API_KEY", "OPENAI_API_KEY"], + ), + Some("openai") => ( + Provider::OpenAiCompatible, + DEFAULT_OPENAI_BASE_URL, + DEFAULT_OPENAI_MODEL, + &["OPENAI_API_KEY"], + ), + Some("gemini") => ( + Provider::Gemini, + DEFAULT_GEMINI_BASE_URL, + DEFAULT_GEMINI_MODEL, + &["GEMINI_API_KEY"], + ), + Some(other) => { + return Err(OmniError::manifest_internal(format!( + "unknown embedding provider '{}' (expected openai-compatible|openai|gemini|mock)", + other + ))); + } + }) +} + fn env_string(name: &str) -> Option { std::env::var(name) .ok() @@ -877,6 +908,37 @@ mod tests { assert!(err.to_string().contains("OPENAI_API_KEY"), "got: {err}"); } + #[test] + fn from_parts_applies_provider_defaults_and_overrides() { + let openrouter = EmbeddingConfig::from_parts(None, None, None, "k".to_string()).unwrap(); + assert_eq!(openrouter.provider, Provider::OpenAiCompatible); + assert_eq!(openrouter.base_url, DEFAULT_OPENROUTER_BASE_URL); + assert_eq!(openrouter.model, DEFAULT_OPENROUTER_MODEL); + assert_eq!(openrouter.api_key, "k"); + + let gemini = + EmbeddingConfig::from_parts(Some("gemini"), None, None, "g".to_string()).unwrap(); + assert_eq!(gemini.provider, Provider::Gemini); + assert_eq!(gemini.base_url, DEFAULT_GEMINI_BASE_URL); + + let overridden = EmbeddingConfig::from_parts( + Some("openai"), + Some("https://x/v1/".to_string()), + Some("custom".to_string()), + "k".to_string(), + ) + .unwrap(); + assert_eq!(overridden.base_url, "https://x/v1"); // trailing slash trimmed + assert_eq!(overridden.model, "custom"); + + let err = + EmbeddingConfig::from_parts(Some("cohere"), None, None, "k".to_string()).unwrap_err(); + assert!( + err.to_string().contains("unknown embedding provider"), + "got: {err}" + ); + } + #[test] #[serial] fn from_env_openai_compatible_prefers_openrouter_key() { @@ -921,7 +983,7 @@ mod tests { fn from_env_unknown_provider_errors() { let _guard = cleared_env(&[("OMNIGRAPH_EMBED_PROVIDER", Some("cohere"))]); let err = EmbeddingConfig::from_env().unwrap_err(); - assert!(err.to_string().contains("unknown OMNIGRAPH_EMBED_PROVIDER")); + assert!(err.to_string().contains("unknown embedding provider")); } #[test]