mirror of
https://github.com/ModernRelay/omnigraph.git
synced 2026-06-18 02:24:27 +02:00
feat(cluster): per-graph embedding profile config (RFC-012 Phase 5)
Engine: factor the alias->(provider, base_url, model, key-envs) defaults into a shared provider_profile, used by both from_env and a new pub from_parts(provider, base_url, model, api_key) so the cluster path applies identical defaults (DRY).
Cluster: GraphConfig gains an optional per-graph embeddings: profile (EmbeddingProfile { provider?, base_url?, model?, api_key }). EmbeddingProfile::resolve() resolves the api_key from a ${NAME} env ref (resolve_secret_ref rejects inline values and unset vars) and builds an engine EmbeddingConfig via from_parts.
Parse + resolve + validate only; wiring the profile through the applied cluster state into the served handle (the config-free s3:// boot reads applied state, not cluster.yaml) follows. 21 engine + 4 cluster tests pass.
This commit is contained in:
Ragnor Comerford
parent
3de93ae7ab
commit
7d412988b5
2 changed files with 209 additions and 40 deletions
|
|
@ -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<EmbeddingProfile>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
#[serde(default)]
|
||||
pub base_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
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<omnigraph::embedding::EmbeddingConfig, String> {
|
||||
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<String, String> {
|
||||
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<String>,
|
||||
}
|
||||
|
||||
#[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") };
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
model: Option<String>,
|
||||
api_key: String,
|
||||
) -> Result<Self> {
|
||||
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<String> {
|
|||
.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<String> {
|
||||
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]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue