use std::time::Duration;
use async_trait::async_trait;
use graph_storage_sdk::models::EmbeddingSpaceId;
use graph_storage_sdk::plugin_api::{
EmbedRequest, EmbedResponse, EmbeddingProviderError, EmbeddingProviderV1,
};
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::{debug, warn};
use url::Url;
#[derive(Clone, Debug)]
pub struct RemoteProviderConfig {
pub base_url: String,
pub model: String,
pub api_key: Option<SecretString>,
pub dimension: u32,
pub request_dimensions: bool,
pub normalize: bool,
pub batch_size: usize,
pub timeout: Duration,
pub max_retries: u32,
}
impl RemoteProviderConfig {
#[must_use]
pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
Self {
base_url: base_url.into(),
model: model.into(),
api_key: None,
dimension: 384,
request_dimensions: true,
normalize: true,
batch_size: 64,
timeout: Duration::from_mins(1),
max_retries: 2,
}
}
#[must_use]
pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
self.api_key = Some(SecretString::from(api_key.into()));
self
}
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum RemoteConfigError {
#[error("base_url {0:?} is not an absolute http(s) URL")]
BaseUrl(String),
#[error(
"base_url carries credentials in its userinfo; put the key in the environment variable \
`embedding_remote_api_key_env` names instead, where it is not part of a URL that gets \
logged, echoed in an error, or copied into a bug report"
)]
CredentialsInUrl,
#[error(
"an api_key is configured and base_url is plain http to {host}, which puts the bearer \
token on the wire in clear for every proxy and log in between; use https, or a \
loopback host if this is a local endpoint"
)]
CredentialOverPlainHttp { host: String },
#[error("model must not be empty")]
Model,
#[error("model must be at most {max} characters, and this one is {got}")]
ModelTooLong { max: usize, got: usize },
#[error(
"model must not carry control characters: it becomes part of the embedding-space \
identity, which is compared for equality and written to the operator log"
)]
ModelControlCharacters,
#[error("dimension must be positive")]
Dimension,
#[error("batch_size must be positive")]
BatchSize,
#[error("the HTTP client could not be built: {0}")]
Client(String),
}
fn is_loopback(url: &Url) -> bool {
match url.host() {
Some(url::Host::Ipv4(address)) => address.is_loopback(),
Some(url::Host::Ipv6(address)) => address.is_loopback(),
Some(url::Host::Domain(name)) => name.eq_ignore_ascii_case("localhost"),
None => false,
}
}
const MAX_MODEL_LEN: usize = 200;
const EMPTY_INPUT_PLACEHOLDER: &str = "(empty)";
pub struct RemoteEmbeddingProvider {
http: reqwest::Client,
endpoint: Url,
config: RemoteProviderConfig,
space: EmbeddingSpaceId,
}
impl RemoteEmbeddingProvider {
pub fn new(config: RemoteProviderConfig) -> Result<Self, RemoteConfigError> {
let base = Url::parse(config.base_url.trim_end_matches('/'))
.ok()
.filter(|url| matches!(url.scheme(), "http" | "https") && url.host_str().is_some())
.ok_or_else(|| RemoteConfigError::BaseUrl(config.base_url.clone()))?;
if !base.username().is_empty() || base.password().is_some() {
return Err(RemoteConfigError::CredentialsInUrl);
}
if config.api_key.is_some() && base.scheme() == "http" && !is_loopback(&base) {
return Err(RemoteConfigError::CredentialOverPlainHttp {
host: base.host_str().unwrap_or_default().to_owned(),
});
}
let model = config.model.trim();
if model.is_empty() {
return Err(RemoteConfigError::Model);
}
if model.chars().count() > MAX_MODEL_LEN {
return Err(RemoteConfigError::ModelTooLong {
max: MAX_MODEL_LEN,
got: model.chars().count(),
});
}
if model.chars().any(char::is_control) {
return Err(RemoteConfigError::ModelControlCharacters);
}
if config.dimension == 0 {
return Err(RemoteConfigError::Dimension);
}
if config.batch_size == 0 {
return Err(RemoteConfigError::BatchSize);
}
let mut endpoint = base.clone();
endpoint.set_path(&format!("{}/embeddings", base.path().trim_end_matches('/')));
endpoint.set_query(None);
let http = reqwest::Client::builder()
.timeout(config.timeout)
.build()
.map_err(|error| RemoteConfigError::Client(error.to_string()))?;
let space = EmbeddingSpaceId::new(
format!("{}@{}", config.model.trim(), endpoint_name(&endpoint)),
"provider-managed",
serde_json::json!({
"protocol": "openai-embeddings-v1",
"requested_dimensions": config.request_dimensions.then_some(config.dimension),
"empty_input": EMPTY_INPUT_PLACEHOLDER,
}),
serde_json::json!({ "strategy": "provider" }),
serde_json::json!({ "l2": config.normalize }),
config.dimension,
);
Ok(Self {
http,
endpoint,
config,
space,
})
}
#[must_use]
pub fn endpoint(&self) -> &Url {
&self.endpoint
}
}
fn endpoint_name(endpoint: &Url) -> String {
let scheme = endpoint.scheme();
let host = endpoint.host_str().unwrap_or("unknown-host");
match endpoint.port() {
Some(port) => format!("{scheme}://{host}:{port}{}", endpoint.path()),
None => format!("{scheme}://{host}{}", endpoint.path()),
}
}
#[derive(Serialize)]
struct EmbeddingsRequest<'a> {
model: &'a str,
input: Vec<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
dimensions: Option<u32>,
}
#[derive(Deserialize)]
struct EmbeddingsResponse {
#[serde(default)]
data: Vec<EmbeddingDatum>,
}
#[derive(Deserialize)]
struct EmbeddingDatum {
index: usize,
embedding: Vec<f32>,
}
#[async_trait]
impl EmbeddingProviderV1 for RemoteEmbeddingProvider {
fn embedding_space(&self) -> &EmbeddingSpaceId {
&self.space
}
fn dimension(&self) -> u32 {
self.config.dimension
}
async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, EmbeddingProviderError> {
if req.cancel.is_cancelled() {
return Err(EmbeddingProviderError::Cancelled);
}
if req.budget.is_exhausted() {
return Err(EmbeddingProviderError::Deadline);
}
let mut vectors: Vec<Vec<f32>> = Vec::with_capacity(req.inputs.len());
for chunk in req.inputs.chunks(self.config.batch_size) {
if req.cancel.is_cancelled() {
return Err(EmbeddingProviderError::Cancelled);
}
let remaining = req.budget.remaining();
if remaining.is_zero() {
return Err(EmbeddingProviderError::Deadline);
}
let batch = self.embed_chunk_with_retries(chunk, &req).await?;
vectors.extend(batch);
}
Ok(EmbedResponse {
vectors,
space: self.space.clone(),
})
}
async fn health(&self) -> Result<(), EmbeddingProviderError> {
self.embed_chunk(
&["health".to_owned()],
AttemptWindow::provider_bound(self.config.timeout.min(Duration::from_secs(10))),
)
.await
.map(drop)
.map_err(|refusal| refusal.error)
}
}
#[derive(Clone, Copy)]
struct AttemptWindow {
timeout: Duration,
caller_bound: bool,
}
impl AttemptWindow {
fn for_attempt(remaining: Duration, configured: Duration) -> Self {
Self {
timeout: remaining.min(configured),
caller_bound: remaining <= configured,
}
}
fn provider_bound(timeout: Duration) -> Self {
Self {
timeout,
caller_bound: false,
}
}
}
impl RemoteEmbeddingProvider {
async fn embed_chunk_with_retries(
&self,
chunk: &[String],
req: &EmbedRequest,
) -> Result<Vec<Vec<f32>>, EmbeddingProviderError> {
let mut backoff = Duration::from_millis(200);
for attempt in 0..=self.config.max_retries {
let remaining = req.budget.remaining();
if remaining.is_zero() {
return Err(EmbeddingProviderError::Deadline);
}
let call = self.embed_chunk(
chunk,
AttemptWindow::for_attempt(remaining, self.config.timeout),
);
let outcome = tokio::select! {
() = req.cancel.cancelled() => return Err(EmbeddingProviderError::Cancelled),
result = call => result,
};
let refusal = match outcome {
Ok(vectors) => return Ok(vectors),
Err(refusal) => refusal,
};
if attempt == self.config.max_retries || !refusal.retryable {
return Err(refusal.error);
}
let wait = backoff.min(req.budget.remaining());
if wait.is_zero() {
return Err(refusal.error);
}
warn!(
endpoint = %endpoint_name(&self.endpoint),
attempt = attempt + 1,
wait_ms = wait.as_millis(),
"the embeddings endpoint answered transiently; retrying"
);
tokio::select! {
() = req.cancel.cancelled() => return Err(EmbeddingProviderError::Cancelled),
() = tokio::time::sleep(wait) => {}
}
backoff = backoff.saturating_mul(2);
}
Err(EmbeddingProviderError::Internal(
"the retry loop ended without an attempt".to_owned(),
))
}
async fn embed_chunk(
&self,
inputs: &[String],
window: AttemptWindow,
) -> Result<Vec<Vec<f32>>, Refusal> {
let body = EmbeddingsRequest {
model: self.config.model.trim(),
input: inputs
.iter()
.map(|text| {
if text.trim().is_empty() {
EMPTY_INPUT_PLACEHOLDER
} else {
text.as_str()
}
})
.collect(),
dimensions: self
.config
.request_dimensions
.then_some(self.config.dimension),
};
let mut request = self
.http
.post(self.endpoint.clone())
.timeout(window.timeout)
.json(&body);
if let Some(key) = &self.config.api_key {
request = request.bearer_auth(key.expose_secret());
}
let response = request.send().await.map_err(|error| {
if error.is_timeout() && window.caller_bound {
Refusal::permanent(EmbeddingProviderError::Deadline)
} else {
Refusal::transient(EmbeddingProviderError::Unavailable {
reason: format!("{}: {error}", endpoint_name(&self.endpoint)),
})
}
})?;
let status = response.status();
if !status.is_success() {
let body_bytes = response.bytes().await.map_or(0, |b| b.len());
warn!(
endpoint = %endpoint_name(&self.endpoint),
status = status.as_u16(),
body_bytes,
"the embeddings endpoint refused the request"
);
return Err(classify_status(status));
}
let parsed: EmbeddingsResponse = response.json().await.map_err(|error| {
Refusal::permanent(EmbeddingProviderError::Internal(format!(
"unparseable embeddings response: {error}"
)))
})?;
self.align(inputs.len(), parsed.data)
.map_err(Refusal::permanent)
}
fn align(
&self,
expected: usize,
data: Vec<EmbeddingDatum>,
) -> Result<Vec<Vec<f32>>, EmbeddingProviderError> {
let width = self.config.dimension as usize;
let mut slots: Vec<Option<Vec<f32>>> = vec![None; expected];
for datum in data {
let Some(slot) = slots.get_mut(datum.index) else {
return Err(EmbeddingProviderError::Internal(format!(
"the endpoint returned index {} for a batch of {expected}",
datum.index
)));
};
if slot.is_some() {
return Err(EmbeddingProviderError::Internal(format!(
"the endpoint returned index {} twice",
datum.index
)));
}
if datum.embedding.len() != width {
debug!(
got = datum.embedding.len(),
want = width,
"the endpoint returned a vector of another width"
);
return Err(EmbeddingProviderError::SpaceMismatch);
}
if !datum.embedding.iter().all(|lane| lane.is_finite()) {
return Err(EmbeddingProviderError::Internal(format!(
"vector {} carries a non-finite lane",
datum.index
)));
}
*slot = Some(if self.config.normalize {
normalize(datum.embedding)
} else {
datum.embedding
});
}
slots
.into_iter()
.enumerate()
.map(|(index, slot)| {
slot.ok_or_else(|| {
EmbeddingProviderError::Internal(format!(
"the endpoint returned no vector for input {index} of {expected}"
))
})
})
.collect()
}
}
struct Refusal {
error: EmbeddingProviderError,
retryable: bool,
}
impl Refusal {
const fn transient(error: EmbeddingProviderError) -> Self {
Self {
error,
retryable: true,
}
}
const fn permanent(error: EmbeddingProviderError) -> Self {
Self {
error,
retryable: false,
}
}
}
fn classify_status(status: reqwest::StatusCode) -> Refusal {
match status.as_u16() {
401 | 403 => Refusal::permanent(EmbeddingProviderError::Unavailable {
reason: format!("the endpoint refused the credential (HTTP {status})"),
}),
408 | 429 | 500 | 502 | 503 | 504 => {
Refusal::transient(EmbeddingProviderError::Unavailable {
reason: format!("HTTP {status}"),
})
}
501 | 505 => Refusal::permanent(EmbeddingProviderError::Unavailable {
reason: format!("the endpoint cannot serve this request at all (HTTP {status})"),
}),
_ => Refusal::permanent(EmbeddingProviderError::Internal(format!(
"the endpoint answered HTTP {status}"
))),
}
}
#[expect(
clippy::cast_possible_truncation,
reason = "narrowing to f32 is the destination: pgvector stores single precision"
)]
fn normalize(vector: Vec<f32>) -> Vec<f32> {
let norm = vector
.iter()
.map(|lane| f64::from(*lane).powi(2))
.sum::<f64>()
.sqrt();
if norm == 0.0 {
return vector;
}
vector
.into_iter()
.map(|lane| (f64::from(lane) / norm) as f32)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn provider(base_url: &str, model: &str) -> RemoteEmbeddingProvider {
RemoteEmbeddingProvider::new(RemoteProviderConfig::new(base_url, model))
.unwrap_or_else(|error| panic!("a valid configuration must build: {error}"))
}
#[test]
fn the_endpoint_is_the_base_url_plus_embeddings() {
assert_eq!(
provider("https://api.openai.com/v1", "m")
.endpoint()
.as_str(),
"https://api.openai.com/v1/embeddings"
);
assert_eq!(
provider("https://api.openai.com/v1/", "m")
.endpoint()
.as_str(),
"https://api.openai.com/v1/embeddings"
);
assert_eq!(
provider("http://ollama:11434/v1?x=1", "m")
.endpoint()
.as_str(),
"http://ollama:11434/v1/embeddings"
);
}
#[test]
fn a_trailing_slash_or_query_does_not_change_the_identity() {
let one = provider("https://api.openai.com/v1", "text-embedding-3-small");
let other = provider(
"https://api.openai.com/v1/?trace=1",
"text-embedding-3-small",
);
assert_eq!(
one.embedding_space().identity_hash,
other.embedding_space().identity_hash
);
}
#[test]
fn model_endpoint_and_width_are_each_part_of_the_identity() {
let base = provider("https://api.openai.com/v1", "text-embedding-3-small");
let other_model = provider("https://api.openai.com/v1", "text-embedding-3-large");
let other_host = provider("https://eu.api.example.com/v1", "text-embedding-3-small");
let mut narrow =
RemoteProviderConfig::new("https://api.openai.com/v1", "text-embedding-3-small");
narrow.dimension = 256;
let narrow = RemoteEmbeddingProvider::new(narrow)
.unwrap_or_else(|error| panic!("a valid configuration must build: {error}"));
let hash = |p: &RemoteEmbeddingProvider| p.embedding_space().identity_hash.clone();
assert_ne!(hash(&base), hash(&other_model));
assert_ne!(hash(&base), hash(&other_host));
assert_ne!(hash(&base), hash(&narrow));
}
#[test]
fn a_base_url_carrying_credentials_is_refused() {
for url in [
"https://user:secret@embeddings.test/v1",
"https://tokenonly@embeddings.test/v1",
] {
let error = RemoteEmbeddingProvider::new(RemoteProviderConfig::new(url, "m"))
.err()
.expect("a URL with userinfo is refused");
assert!(
matches!(error, RemoteConfigError::CredentialsInUrl),
"{url}: {error}"
);
}
let ok = provider("https://embeddings.test/v1", "m");
assert_eq!(ok.endpoint().username(), "");
}
#[test]
fn the_transport_is_part_of_the_identity() {
let secure = provider("https://embeddings.test/v1", "m");
let plain = provider("http://embeddings.test/v1", "m");
assert_ne!(
secure.embedding_space().identity_hash,
plain.embedding_space().identity_hash,
"http and https must not share an embedding space"
);
}
#[test]
fn a_relative_or_non_http_base_url_is_refused() {
for bad in ["api.openai.com/v1", "ftp://x/v1", "", "https://"] {
let error = RemoteEmbeddingProvider::new(RemoteProviderConfig::new(bad, "m"))
.err()
.unwrap_or_else(|| panic!("{bad:?} must be refused"));
assert!(
matches!(error, RemoteConfigError::BaseUrl(_)),
"{bad:?}: {error}"
);
}
}
#[test]
fn the_credential_does_not_render_in_debug() {
let config = RemoteProviderConfig::new("https://api.openai.com/v1", "m")
.with_api_key("sk-this-must-not-leak");
let rendered = format!("{config:?}");
assert!(!rendered.contains("sk-this-must-not-leak"), "{rendered}");
}
#[test]
fn normalization_yields_unit_length_and_leaves_zero_alone() {
let unit = normalize(vec![3.0, 4.0]);
let norm: f64 = unit
.iter()
.map(|x| f64::from(*x).powi(2))
.sum::<f64>()
.sqrt();
assert!((norm - 1.0).abs() < 1e-6, "{norm}");
assert_eq!(normalize(vec![0.0, 0.0]), vec![0.0, 0.0]);
}
fn classify(code: u16) -> Refusal {
classify_status(reqwest::StatusCode::from_u16(code).unwrap_or_default())
}
#[test]
fn a_credential_is_not_sent_in_clear_to_another_host() {
let build = |url: &str, key: Option<&str>| {
let mut config = RemoteProviderConfig::new(url, "text-embedding-3-small");
if let Some(key) = key {
config = config.with_api_key(key);
}
RemoteEmbeddingProvider::new(config)
};
for url in [
"http://embeddings.internal/v1",
"http://10.0.0.7:8080/v1",
"http://example.test/v1",
] {
let refused = build(url, Some("sk-secret"))
.err()
.expect("a credential over plain http to another host must be refused");
assert!(
matches!(refused, RemoteConfigError::CredentialOverPlainHttp { .. }),
"{url} must be refused for the scheme, got {refused}"
);
assert!(
!refused.to_string().contains("sk-secret"),
"and the refusal must not quote the credential: {refused}"
);
}
for url in ["http://embeddings.internal/v1", "http://10.0.0.7:8080/v1"] {
assert!(
build(url, None).is_ok(),
"{url} carries no credential and stays allowed"
);
}
for url in [
"http://127.0.0.1:11434/v1",
"http://localhost:11434/v1",
"http://[::1]:11434/v1",
] {
assert!(
build(url, Some("sk-secret")).is_ok(),
"{url} is the local machine and stays allowed"
);
}
assert!(
build("https://api.openai.com/v1", Some("sk-secret")).is_ok(),
"https with a credential is the point of the feature"
);
}
#[test]
fn a_model_that_is_not_a_model_name_is_refused() {
let refused = |model: &str| {
RemoteEmbeddingProvider::new(RemoteProviderConfig::new(
"https://example.test/v1",
model,
))
.err()
.expect("this model must be refused")
};
assert!(
matches!(refused(" "), RemoteConfigError::Model),
"a blank model keeps its own error"
);
assert!(
matches!(
refused(&"m".repeat(MAX_MODEL_LEN + 1)),
RemoteConfigError::ModelTooLong { .. }
),
"a model past the bound must be refused"
);
for model in ["text-embedding\n3-small", "text\u{0}embedding", "a\tb"] {
assert!(
matches!(refused(model), RemoteConfigError::ModelControlCharacters),
"{model:?} must be refused"
);
}
}
#[test]
fn a_real_model_name_is_accepted() {
for model in [
"text-embedding-3-small",
" text-embedding-3-small ",
"sentence-transformers/all-MiniLM-L6-v2",
&"m".repeat(MAX_MODEL_LEN),
] {
assert!(
RemoteEmbeddingProvider::new(RemoteProviderConfig::new(
"https://example.test/v1",
model,
))
.is_ok(),
"{model:?} must be accepted"
);
}
}
#[test]
fn statuses_split_into_unavailable_and_internal() {
let unavailable = |code: u16| {
matches!(
classify(code).error,
EmbeddingProviderError::Unavailable { .. }
)
};
assert!(unavailable(401));
assert!(unavailable(429));
assert!(unavailable(503));
assert!(!unavailable(400));
assert!(!unavailable(404));
}
#[test]
fn a_refused_credential_is_unavailable_but_not_retried() {
for code in [401, 403] {
let refusal = classify(code);
assert!(
matches!(refusal.error, EmbeddingProviderError::Unavailable { .. }),
"HTTP {code} is still an unavailable provider"
);
assert!(
!refusal.retryable,
"HTTP {code} must not be retried: the credential will be refused again, and \
retrying every chunk of every batch is how a rotated key becomes a rate limit"
);
}
}
#[test]
fn a_permanent_5xx_is_not_retried() {
for code in [501, 505] {
let refusal = classify(code);
assert!(
!refusal.retryable,
"HTTP {code} describes a fixed capability, so repeating the request repeats \
the answer"
);
assert!(
matches!(refusal.error, EmbeddingProviderError::Unavailable { .. }),
"HTTP {code} still leaves the vector arm down, so it is still unavailable"
);
}
}
#[test]
fn the_statuses_a_retry_can_fix_are_still_retried() {
for code in [408, 429, 500, 502, 503, 504] {
assert!(
classify(code).retryable,
"HTTP {code} is the endpoint saying `not now`, which is what retrying is for"
);
}
for code in [400, 404, 422] {
assert!(
!classify(code).retryable,
"HTTP {code} is about the request, and repeating it repeats the request"
);
}
}
}