apunkt/omnigraph
1
0
Fork 0
mirror of https://github.com/ModernRelay/omnigraph.git synced 2026-06-18 02:24:27 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(engine)!: provider-independent embedding client (RFC-012 Phase 2)

Browse source 
Replace the Gemini-only EmbeddingClient with one resolved EmbeddingConfig { provider, model, base_url, api_key } behind a sealed Provider enum (OpenAiCompatible | Gemini | Mock). OpenAiCompatible (POST {base}/embeddings, bearer, {model, input, dimensions}) covers OpenRouter — the new default gateway — OpenAI direct, and self-hosted endpoints; Gemini keeps its RETRIEVAL_QUERY/RETRIEVAL_DOCUMENT task types; Mock is offline/deterministic. EmbedRole replaces the task-type string.

from_env() resolves provider via OMNIGRAPH_EMBED_PROVIDER (default openai-compatible), base/model via OMNIGRAPH_EMBED_BASE_URL/_MODEL, and the api key from OPENROUTER_API_KEY/OPENAI_API_KEY or GEMINI_API_KEY. BREAKING (pre-release, no back-compat): the default provider is now OpenRouter, OMNIGRAPH_GEMINI_BASE_URL is dropped, and Gemini-only users set OMNIGRAPH_EMBED_PROVIDER=gemini.

Folds in RFC-012 Phase 1 NFR floor: a total-operation OMNIGRAPH_EMBED_QUERY_DEADLINE_MS deadline (default 60s; 0=unbounded) bounds the ~121s worst case, and tracing spans (target omnigraph::embedding) record provider/model/dim/attempt/elapsed/outcome. The offline 'omnigraph embed' CLI follows the resolved provider (its hardcoded gemini-only bail removed). 17 engine embedding unit tests, 4 CLI embed tests, and the search integration suite (22) pass.

Cross-query client reuse and the docs refresh land in follow-up commits.
This commit is contained in:
Ragnor Comerford 2026-06-15 17:15:11 +02:00
parent 7c916f5b98
commit b999ae3753
No known key found for this signature in database
4 changed files with 518 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

15
crates/omnigraph-cli/src/embed.rs
View file

@ -9,8 +9,6 @@ use omnigraph::embedding::EmbeddingClient;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
const DEFAULT_EMBED_MODEL: &str = "gemini-embedding-2-preview";
#[derive(Debug, Args, Clone)]
pub(crate) struct EmbedArgs {
/// Seed manifest path
@ -85,7 +83,7 @@ impl EmbedMode {
#[derive(Debug, Clone, Deserialize)]
struct EmbedSpec {
#[serde(default = "default_embed_model")]
#[serde(default)]
model: String,
dimension: usize,
types: BTreeMap<String, EmbedTypeSpec>,
@ -180,13 +178,6 @@ pub(crate) fn resolve_embed_job(args: &EmbedArgs) -> Result<EmbedJob> {
(input, output, spec)
};
if spec.model != DEFAULT_EMBED_MODEL {
bail!(
"only {} is supported for explicit seed embeddings right now",
DEFAULT_EMBED_MODEL
);
}
Ok(EmbedJob {
input,
output,
@ -315,10 +306,6 @@ fn temp_output_path(output: &Path) -> PathBuf {
PathBuf::from(temp)
}
fn default_embed_model() -> String {
DEFAULT_EMBED_MODEL.to_string()
}
fn load_embed_spec(path: &Path) -> Result<EmbedSpec> {
Ok(serde_json::from_str(&fs::read_to_string(path)?)?)
}