agentplane 0.10.0

//! Embedding drivers, for the seam semantic retrieval needs.
//!
//! Three, chosen by which wire a provider speaks:
//!
//! * `OpenAiEmbedder` — `POST /v1/embeddings`, which reaches `OpenAI`, Voyage
//!   AI, Ollama, vLLM, TGI, LM Studio and Hugging Face's router. One driver, as
//!   `chat-completions` is one driver for completions.
//! * `GeminiEmbedder` — Google's `:embedContent`, a different wire.
//! * `BedrockEmbedder` — Titan and Cohere on Bedrock, for a plane whose data
//!   may not leave one AWS account.
//!
//! Anthropic has no embeddings API and recommends Voyage, which the first driver
//! reaches by base URL alone.
//!
//! # What is deliberately absent
//!
//! **Batching.** The seam embeds one text, because the vector goes into an
//! effect key and one effect is one observation. Batching would have to decide
//! how a partial failure maps onto several effects.
//!
//! **A default model.** [`Embedder::revision`](crate::memory::Embedder::revision)
//! is in the effect key, so a guessed default would put a guess in the identity
//! of every vector.

use crate::core::Secret;
use crate::core::StoreError;
use crate::memory::Embedder;

/// One embedding component, as `f32`.
///
/// Vectors are `f32` by the seam's own contract — that is what an index stores
/// and what cosine is computed in. JSON carries `f64`, so narrowing is not a
/// loss of information the caller had: it is the wire's precision meeting the
/// type the whole retrieval path already uses. A value no `f32` can hold is not
/// an embedding component, and the caller's length check refuses the reply
/// rather than silently shortening it.
#[allow(clippy::cast_possible_truncation)]
fn json_f32(value: &serde_json::Value) -> Option<f32> {
    value.as_f64().map(|f| f as f32)
}

/// Embeddings over the `OpenAI`-compatible wire.
#[derive(Debug, Clone)]
pub struct OpenAiEmbedder {
    http: reqwest::Client,
    key: Option<Secret>,
    base: String,
    model: String,
    dimensions: Option<u32>,
    input_type: Option<String>,
    egress: Option<crate::core::Egress>,
    timeout: std::time::Duration,
}

impl OpenAiEmbedder {
    /// Five minutes, matching the model drivers.
    pub const DEFAULT_TIMEOUT: std::time::Duration = std::time::Duration::from_mins(5);

    /// A driver for one embedding model.
    ///
    /// The model is required and has no default: it is `revision()`, which sits
    /// in the retrieval effect's key, and a guessed default would put a guess
    /// into the identity of every vector this ever produced.
    ///
    /// # Errors
    ///
    /// If the HTTP client cannot be built.
    pub fn new(model: impl Into<String>) -> Result<Self, StoreError> {
        let http = reqwest::Client::builder()
            .build()
            .map_err(|e| StoreError::Backend(format!("could not build an HTTP client: {e}")))?;
        Ok(Self {
            http,
            key: None,
            base: "https://api.openai.com".to_owned(),
            model: model.into(),
            dimensions: None,
            input_type: None,
            egress: None,
            timeout: Self::DEFAULT_TIMEOUT,
        })
    }

    /// The bearer token, when the server wants one.
    ///
    /// Optional because the common local server does not, which is the same
    /// shape `chat-completions` has.
    #[must_use]
    pub fn key(mut self, key: impl Into<String>) -> Self {
        self.key = Some(Secret::new(key));
        self
    }

    /// Point at a different host — a local server, a gateway, or a test double.
    #[must_use]
    pub fn base(mut self, base: impl Into<String>) -> Self {
        self.base = base.into();
        self
    }

    /// Ask for a shortened vector, where the model supports it.
    ///
    /// In [`revision`](Embedder::revision), because a 1536-dimension vector and
    /// a 256-dimension one from the same model are not comparable and must not
    /// share an index — nor an effect identity.
    #[must_use]
    pub const fn dimensions(mut self, dimensions: u32) -> Self {
        self.dimensions = Some(dimensions);
        self
    }

    /// Tell an asymmetric model that this text is a **query**.
    ///
    /// Some models embed questions and documents into deliberately different
    /// regions, and rank badly when the two are swapped. Voyage prepends a
    /// prompt when told (`"query"`); Cohere's direct API takes the same field
    /// (`"search_query"`). `OpenAI`'s own models are symmetric and take no such
    /// parameter, which is why this is opt-in rather than defaulted — sending an
    /// unknown field to a server that does not want it is a 400.
    ///
    /// **Nothing in a reply says you got this wrong.** The vectors come back the
    /// right shape and rank slightly worse forever, which is why it is a typed
    /// option here rather than something a caller is left to discover.
    ///
    /// It goes into [`revision`](Embedder::revision): a query embedded with the
    /// hint and one without are not comparable, so they must not share an
    /// effect identity.
    #[must_use]
    pub fn input_type(mut self, input_type: impl Into<String>) -> Self {
        self.input_type = Some(input_type.into());
        self
    }

    /// Restrict where this driver may connect.
    ///
    /// Deny-by-default once set, refused **before the request is built**, so
    /// nothing leaves. An embedding call sends the query text to a third party;
    /// that it is short does not make it uninteresting.
    #[must_use]
    pub fn egress(mut self, egress: crate::core::Egress) -> Self {
        self.egress = Some(egress);
        self
    }

    /// Bound connection and response as one operation.
    #[must_use]
    pub const fn timeout(mut self, timeout: std::time::Duration) -> Self {
        self.timeout = timeout;
        self
    }

    fn check_egress(&self) -> Result<(), StoreError> {
        let Some(egress) = &self.egress else {
            return Ok(());
        };
        let host = reqwest::Url::parse(&self.base)
            .ok()
            .and_then(|u| u.host_str().map(ToOwned::to_owned));
        egress
            .permits(host.as_deref())
            .map_err(|e| StoreError::Backend(e.to_string()))
    }
}

#[derive(serde::Deserialize)]
struct EmbeddingsReply {
    data: Vec<EmbeddingDatum>,
}

#[derive(serde::Deserialize)]
struct EmbeddingDatum {
    embedding: Vec<f32>,
}

#[async_trait::async_trait]
impl Embedder for OpenAiEmbedder {
    /// Model **and** dimension count.
    ///
    /// Both, because both change the vector. Two runs whose queries were
    /// embedded at different widths would rank against different geometry, and
    /// the effect key is what stops a replay reading one as the other.
    fn revision(&self) -> String {
        use std::fmt::Write as _;
        let mut revision = self.model.clone();
        if let Some(d) = self.dimensions {
            let _ = write!(revision, "@{d}");
        }
        if let Some(input_type) = &self.input_type {
            let _ = write!(revision, "/{input_type}");
        }
        revision
    }

    async fn embed(&self, text: &str) -> Result<Vec<f32>, StoreError> {
        self.check_egress()?;

        let mut body = serde_json::json!({ "model": self.model, "input": text });
        if let Some(dimensions) = self.dimensions {
            body["dimensions"] = serde_json::json!(dimensions);
        }
        if let Some(input_type) = &self.input_type {
            body["input_type"] = serde_json::json!(input_type);
        }

        let url = format!("{}/v1/embeddings", self.base.trim_end_matches('/'));
        let mut request = self.http.post(&url).timeout(self.timeout).json(&body);
        if let Some(key) = &self.key {
            request = request.bearer_auth(key.expose());
        }

        let response = request
            .send()
            .await
            .map_err(|e| StoreError::Backend(format!("{url}: {e}")))?;
        let status = response.status();
        let text_body = response
            .text()
            .await
            .map_err(|e| StoreError::Backend(format!("{url}: unreadable reply: {e}")))?;
        if !status.is_success() {
            // The body, not only the code: an embeddings 400 is almost always
            // "that model does not exist" or "input too long", and a bare status
            // sends somebody to check their key.
            return Err(StoreError::Backend(format!(
                "{url}: embeddings returned {status}: {text_body}"
            )));
        }

        let reply: EmbeddingsReply = serde_json::from_str(&text_body)
            .map_err(|e| StoreError::Backend(format!("{url}: unreadable reply: {e}")))?;

        // Exactly one, because exactly one text was sent. A server answering
        // with more — or none — is not speaking this wire, and returning the
        // first of several would silently pick a vector for somebody else's
        // input.
        let [datum] = reply.data.as_slice() else {
            return Err(StoreError::Backend(format!(
                "{url}: one input was sent and {} embeddings came back",
                reply.data.len()
            )));
        };
        if datum.embedding.is_empty() {
            return Err(StoreError::Backend(format!(
                "{url}: the embedding is empty, which no index can rank against"
            )));
        }
        Ok(datum.embedding.clone())
    }
}

/// How a Bedrock embedding family spells a request.
///
/// Bedrock's `InvokeModel` is a **passthrough**: the body is the vendor's own,
/// and the vendors do not agree. Titan takes `{"inputText": …}` and answers
/// `{"embedding": […]}`; Cohere takes `{"texts": […], "input_type": …}` and
/// answers `{"embeddings": [[…]]}`. There is no shape this driver can derive
/// from the model id, so it is **declared, never sniffed** — the same decision
/// [`ReasoningDialect`](super::bedrock::ReasoningDialect) makes, and for the
/// same reason: cross-region inference profiles prefix the id
/// (`us.amazon.titan-…`), so a substring match would bind behaviour to a naming
/// convention AWS owns and can change.
///
/// An undeclared dialect refuses rather than guessing. Sending Titan's body to
/// Cohere does not fail cleanly — it is a 400 whose text blames the caller — and
/// guessing wrong on the *response* would be worse: `embeddings[0]` read as
/// `embedding` yields a vector of the wrong rank that ranks against nothing.
#[cfg(feature = "bedrock")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum EmbeddingDialect {
    /// Amazon Titan Embed: `inputText`, answering `embedding`.
    Titan,
    /// Cohere Embed: `texts` plus `input_type`, answering `embeddings`.
    Cohere,
}

/// The request body for one Bedrock embedding dialect.
///
/// A free function rather than a method, so the vendor mapping — the part that
/// is easy to get silently wrong and impossible to test against AWS without
/// credentials — is exercisable on its own.
///
/// # Errors
///
/// If a parameter is set that the dialect has no place for.
#[cfg(feature = "bedrock")]
fn bedrock_body(
    dialect: EmbeddingDialect,
    text: &str,
    dimensions: Option<u32>,
) -> Result<serde_json::Value, StoreError> {
    match dialect {
        EmbeddingDialect::Titan => {
            let mut body = serde_json::json!({ "inputText": text });
            if let Some(dimensions) = dimensions {
                body["dimensions"] = serde_json::json!(dimensions);
            }
            Ok(body)
        }
        EmbeddingDialect::Cohere => {
            if dimensions.is_some() {
                return Err(StoreError::Backend(
                    "Cohere Embed takes no `dimensions`; the width is the model's. \
                     Drop `.dimensions(..)` or choose a Titan model — a knob that \
                     silently did nothing would put a width in the effect key that \
                     never reached the wire"
                        .to_owned(),
                ));
            }
            // `search_query`, not `search_document`: this embeds the thing being
            // looked *for*. Cohere's asymmetric models rank badly when the two
            // are swapped, and nothing in the reply says so.
            Ok(serde_json::json!({ "texts": [text], "input_type": "search_query" }))
        }
    }
}

/// Embeddings through Amazon Bedrock's `InvokeModel`.
///
/// # Why this exists beside the `OpenAI`-compatible one
///
/// Bedrock is the deployment that cannot reach the other driver: a plane chosen
/// for Bedrock was usually chosen because its data may not leave a particular
/// account, and telling it to call `api.openai.com` for the *query text* is the
/// one thing it cannot do. Without this, semantic retrieval was unavailable to
/// exactly the regulated deployments this crate is built for — the AWS SDK
/// already being paid for by the `bedrock` feature.
#[cfg(feature = "bedrock")]
#[derive(Debug, Clone)]
pub struct BedrockEmbedder {
    client: aws_sdk_bedrockruntime::Client,
    model: String,
    region: String,
    dialect: EmbeddingDialect,
    dimensions: Option<u32>,
}

#[cfg(feature = "bedrock")]
impl BedrockEmbedder {
    /// Load the standard AWS credential chain for this region.
    ///
    /// # Errors
    ///
    /// If the region is blank.
    pub async fn from_env(
        region: impl Into<String>,
        model: impl Into<String>,
        dialect: EmbeddingDialect,
    ) -> Result<Self, StoreError> {
        let region = region.into();
        if region.trim().is_empty() {
            return Err(StoreError::Backend(
                "an AWS region is required for Bedrock embeddings".to_owned(),
            ));
        }
        let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
            .region(aws_config::Region::new(region.clone()))
            .load()
            .await;
        Ok(Self::from_client(
            aws_sdk_bedrockruntime::Client::new(&config),
            region,
            model,
            dialect,
        ))
    }

    /// Build from an already configured AWS client.
    #[must_use]
    pub fn from_client(
        client: aws_sdk_bedrockruntime::Client,
        region: impl Into<String>,
        model: impl Into<String>,
        dialect: EmbeddingDialect,
    ) -> Self {
        Self {
            client,
            model: model.into(),
            region: region.into(),
            dialect,
            dimensions: None,
        }
    }

    /// Ask Titan for a shortened vector.
    ///
    /// Titan v2 accepts 256, 512 or 1024. Cohere does not take the parameter at
    /// all, so setting it there is refused rather than ignored — a ceiling that
    /// silently does nothing is the shape this crate refuses everywhere else.
    #[must_use]
    pub const fn dimensions(mut self, dimensions: u32) -> Self {
        self.dimensions = Some(dimensions);
        self
    }
}

#[cfg(feature = "bedrock")]
#[async_trait::async_trait]
impl Embedder for BedrockEmbedder {
    /// Region, model and width.
    ///
    /// The **region** is in it because a Bedrock model id names a model, not a
    /// deployment: the same id in two regions is two services, and a vector from
    /// one has no standing in an index built from the other. It is also the fact
    /// a compliance reader most wants on the record.
    fn revision(&self) -> String {
        let base = format!("bedrock:{}/{}", self.region, self.model);
        self.dimensions
            .map_or(base.clone(), |d| format!("{base}@{d}"))
    }

    async fn embed(&self, text: &str) -> Result<Vec<f32>, StoreError> {
        let body = bedrock_body(self.dialect, text, self.dimensions)?;

        let reply = self
            .client
            .invoke_model()
            .model_id(&self.model)
            .content_type("application/json")
            .accept("application/json")
            // `canon`, not `serde_json::to_vec`: the layering guard funnels every
            // byte-producing path through the canonical writer, and canonical
            // JSON is still JSON — sorting the keys costs nothing on the wire
            // and makes this driver's request bytes reproducible, which is the
            // property the rest of the crate already depends on.
            .body(aws_smithy_types::Blob::new(
                crate::core::canon::value_bytes(&body),
            ))
            .send()
            .await
            .map_err(|e| StoreError::Backend(format!("bedrock embeddings: {e}")))?;

        let parsed: serde_json::Value = serde_json::from_slice(reply.body().as_ref())
            .map_err(|e| StoreError::Backend(format!("bedrock embeddings: unreadable: {e}")))?;

        let vector = match self.dialect {
            EmbeddingDialect::Titan => parsed.get("embedding").cloned(),
            // One text was sent, so one row comes back. Taking the first of
            // several would rank this query against somebody else's vector.
            EmbeddingDialect::Cohere => {
                match parsed
                    .get("embeddings")
                    .and_then(|e| e.as_array())
                    .map(Vec::as_slice)
                {
                    Some([only]) => Some(only.clone()),
                    _ => None,
                }
            }
        };
        let Some(serde_json::Value::Array(values)) = vector else {
            return Err(StoreError::Backend(format!(
                "bedrock embeddings: the reply carried no single vector in the \
                 {:?} shape — a declared dialect that does not match the model is \
                 the usual cause: {parsed}",
                self.dialect
            )));
        };
        let vector: Vec<f32> = values.iter().filter_map(json_f32).collect();
        if vector.len() != values.len() || vector.is_empty() {
            return Err(StoreError::Backend(
                "bedrock embeddings: the vector is empty or not all numbers".to_owned(),
            ));
        }
        Ok(vector)
    }
}

/// Embeddings through Google's Gemini API.
///
/// `POST /v1beta/models/{model}:embedContent`, which is not the `OpenAI` wire —
/// different path, different body, different reply — so it gets a driver rather
/// than a base-URL override.
///
/// # Two details that are easy to get silently wrong
///
/// **`taskType` is fixed to `RETRIEVAL_QUERY`, not offered as a knob.** This
/// seam only ever embeds the thing being looked *for*: `SemanticQuery::embedding`
/// is a query by construction, and the documents were embedded by whoever built
/// the index, which is outside this crate. Exposing the choice would let a
/// caller embed a query as a document, which ranks worse and reports nothing.
///
/// **A truncated vector is re-normalised.** `outputDimensionality` uses
/// Matryoshka truncation — the tail is cut off — and the result is no longer
/// unit length. Cosine similarity against a normalised index would then be
/// scaled by whatever magnitude survived, quietly biasing every score. The
/// native width is already normalised and is left alone.
#[cfg(feature = "providers")]
#[derive(Debug, Clone)]
pub struct GeminiEmbedder {
    http: reqwest::Client,
    key: Secret,
    base: String,
    model: String,
    dimensions: Option<u32>,
    egress: Option<crate::core::Egress>,
    timeout: std::time::Duration,
}

#[cfg(feature = "providers")]
impl GeminiEmbedder {
    /// Google's public endpoint.
    pub const DEFAULT_BASE: &'static str = "https://generativelanguage.googleapis.com";

    /// A driver for one Gemini embedding model.
    ///
    /// # Errors
    ///
    /// If the HTTP client cannot be built.
    pub fn new(key: impl Into<String>, model: impl Into<String>) -> Result<Self, StoreError> {
        let http = reqwest::Client::builder()
            .build()
            .map_err(|e| StoreError::Backend(format!("could not build an HTTP client: {e}")))?;
        Ok(Self {
            http,
            key: Secret::new(key),
            base: Self::DEFAULT_BASE.to_owned(),
            model: model.into(),
            dimensions: None,
            egress: None,
            timeout: OpenAiEmbedder::DEFAULT_TIMEOUT,
        })
    }

    /// `GEMINI_API_KEY`, falling back to `GOOGLE_API_KEY`.
    ///
    /// Both are in wide use, and a deployment that exported the other one would
    /// otherwise meet an authentication failure naming neither.
    ///
    /// # Errors
    ///
    /// If neither variable is set.
    pub fn from_env(model: impl Into<String>) -> Result<Self, StoreError> {
        let key = std::env::var("GEMINI_API_KEY")
            .or_else(|_| std::env::var("GOOGLE_API_KEY"))
            .map_err(|_| {
                StoreError::Backend("neither GEMINI_API_KEY nor GOOGLE_API_KEY is set".to_owned())
            })?;
        Self::new(key, model)
    }

    /// Point at a different host — a gateway, or a test server.
    #[must_use]
    pub fn base(mut self, base: impl Into<String>) -> Self {
        self.base = base.into();
        self
    }

    /// Ask for a shortened vector, re-normalised.
    #[must_use]
    pub const fn dimensions(mut self, dimensions: u32) -> Self {
        self.dimensions = Some(dimensions);
        self
    }

    /// Restrict where this driver may connect.
    #[must_use]
    pub fn egress(mut self, egress: crate::core::Egress) -> Self {
        self.egress = Some(egress);
        self
    }

    /// Bound connection and response as one operation.
    #[must_use]
    pub const fn timeout(mut self, timeout: std::time::Duration) -> Self {
        self.timeout = timeout;
        self
    }
}

#[cfg(feature = "providers")]
#[async_trait::async_trait]
impl Embedder for GeminiEmbedder {
    fn revision(&self) -> String {
        let base = format!("gemini:{}", self.model);
        self.dimensions
            .map_or(base.clone(), |d| format!("{base}@{d}"))
    }

    async fn embed(&self, text: &str) -> Result<Vec<f32>, StoreError> {
        if let Some(egress) = &self.egress {
            let host = reqwest::Url::parse(&self.base)
                .ok()
                .and_then(|u| u.host_str().map(ToOwned::to_owned));
            egress
                .permits(host.as_deref())
                .map_err(|e| StoreError::Backend(e.to_string()))?;
        }

        let mut body = serde_json::json!({
            "content": { "parts": [{ "text": text }] },
            "taskType": "RETRIEVAL_QUERY",
        });
        if let Some(dimensions) = self.dimensions {
            body["outputDimensionality"] = serde_json::json!(dimensions);
        }

        let url = format!(
            "{}/v1beta/models/{}:embedContent",
            self.base.trim_end_matches('/'),
            self.model
        );
        // The header form rather than a `?key=` query parameter: a URL reaches
        // proxy logs and crash reports, and a key in one is a key disclosed.
        let response = self
            .http
            .post(&url)
            .timeout(self.timeout)
            .header("x-goog-api-key", self.key.expose())
            .json(&body)
            .send()
            .await
            .map_err(|e| StoreError::Backend(format!("{url}: {e}")))?;
        let status = response.status();
        let text_body = response
            .text()
            .await
            .map_err(|e| StoreError::Backend(format!("{url}: unreadable reply: {e}")))?;
        if !status.is_success() {
            return Err(StoreError::Backend(format!(
                "{url}: embedContent returned {status}: {text_body}"
            )));
        }

        let parsed: serde_json::Value = serde_json::from_str(&text_body)
            .map_err(|e| StoreError::Backend(format!("{url}: unreadable reply: {e}")))?;
        let Some(serde_json::Value::Array(values)) = parsed
            .get("embedding")
            .and_then(|e| e.get("values"))
            .cloned()
        else {
            return Err(StoreError::Backend(format!(
                "{url}: the reply carried no embedding.values: {parsed}"
            )));
        };
        let mut vector: Vec<f32> = values.iter().filter_map(json_f32).collect();
        if vector.len() != values.len() || vector.is_empty() {
            return Err(StoreError::Backend(format!(
                "{url}: the vector is empty or not all numbers"
            )));
        }

        // Matryoshka truncation cuts the tail, so a shortened vector is no
        // longer unit length and cosine against a normalised index would be
        // scaled by whatever magnitude happened to survive. The native width
        // arrives normalised and is left exactly as it came.
        if self.dimensions.is_some() {
            let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
            if norm > 0.0 {
                for v in &mut vector {
                    *v /= norm;
                }
            }
        }
        Ok(vector)
    }
}

#[cfg(all(test, feature = "bedrock"))]
mod bedrock_dialect_tests {
    use super::{EmbeddingDialect, bedrock_body};

    /// Each dialect sends the body its vendor actually takes.
    ///
    /// The mapping is declared rather than sniffed from the model id, so nothing
    /// downstream can notice it being wrong: Titan's `inputText` sent to Cohere
    /// is a 400 blaming the caller, and — worse — Cohere's `embeddings[0]` read
    /// as Titan's `embedding` would yield a vector of the wrong rank that ranks
    /// against nothing.
    #[test]
    fn each_dialect_sends_its_own_shape() {
        let titan = bedrock_body(EmbeddingDialect::Titan, "refund policy", Some(512))
            .expect("titan takes a width");
        assert_eq!(titan["inputText"], "refund policy");
        assert_eq!(titan["dimensions"], 512);

        let cohere = bedrock_body(EmbeddingDialect::Cohere, "refund policy", None).expect("cohere");
        assert_eq!(cohere["texts"][0], "refund policy");
        assert_eq!(
            cohere["input_type"], "search_query",
            "this seam embeds the thing being looked *for*; embedding it as a \
             document ranks worse and reports nothing"
        );
    }

    /// A width Cohere has no place for is refused, not dropped.
    ///
    /// It would otherwise reach `revision()` — and therefore the effect key —
    /// while never reaching the wire, so two runs with different declared widths
    /// would have different identities and identical vectors.
    #[test]
    fn cohere_refuses_a_width_it_cannot_send() {
        let err = bedrock_body(EmbeddingDialect::Cohere, "x", Some(512))
            .expect_err("a width Cohere cannot send was accepted");
        assert!(err.to_string().contains("no `dimensions`"), "{err}");
    }
}