agent-framework-openai 0.5.0

//! OpenAI embeddings client.
//!
//! Rust equivalent of upstream's `OpenAIEmbeddingClient`
//! (`agent_framework_openai/_embedding_client.py`): the
//! [`POST /v1/embeddings`](https://platform.openai.com/docs/api-reference/embeddings)
//! endpoint, batching all input values into one request. Works against any
//! OpenAI-compatible server via [`OpenAIEmbeddingClient::with_base_url`].

use std::sync::Arc;

use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine as _;

use agent_framework_core::client::EmbeddingClient;
use agent_framework_core::error::{Error, Result};
use agent_framework_core::types::{
    Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails,
};
use serde_json::{json, Map, Value};

use crate::{classify_service_error, parse_retry_after, DEFAULT_BASE_URL};

/// An OpenAI (or OpenAI-compatible) embeddings client.
///
/// ```no_run
/// # use agent_framework_openai::OpenAIEmbeddingClient;
/// # use agent_framework_core::client::EmbeddingClient;
/// # async fn demo() -> agent_framework_core::error::Result<()> {
/// let client = OpenAIEmbeddingClient::from_env("text-embedding-3-small")?;
/// let batch = client
///     .get_embeddings(vec!["Hello, world!".into()], None)
///     .await?;
/// println!("{} dims", batch[0].dimensions());
/// # Ok(())
/// # }
/// ```
#[derive(Clone)]
pub struct OpenAIEmbeddingClient {
    inner: Arc<Inner>,
}

struct Inner {
    http: reqwest::Client,
    api_key: String,
    base_url: String,
    model: String,
    organization: Option<String>,
}

impl OpenAIEmbeddingClient {
    /// Create a client for the given API key and default embedding model.
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            inner: Arc::new(Inner {
                http: reqwest::Client::new(),
                api_key: api_key.into(),
                base_url: DEFAULT_BASE_URL.to_string(),
                model: model.into(),
                organization: None,
            }),
        }
    }

    /// Build a client from the environment: `OPENAI_API_KEY` (required) and
    /// optional `OPENAI_BASE_URL`. The model falls back through
    /// `OPENAI_EMBEDDING_MODEL` when the argument is empty — mirroring
    /// upstream's `model or OPENAI_EMBEDDING_MODEL` resolution.
    pub fn from_env(model: impl Into<String>) -> Result<Self> {
        let key = std::env::var("OPENAI_API_KEY")
            .map_err(|_| Error::Configuration("OPENAI_API_KEY is not set".into()))?;
        let mut model = model.into();
        if model.is_empty() {
            model = std::env::var("OPENAI_EMBEDDING_MODEL").map_err(|_| {
                Error::Configuration(
                    "no embedding model: pass one or set OPENAI_EMBEDDING_MODEL".into(),
                )
            })?;
        }
        let mut client = Self::new(key, model);
        if let Ok(base) = std::env::var("OPENAI_BASE_URL") {
            client = client.with_base_url(base);
        }
        Ok(client)
    }

    /// Override the base URL (for OpenAI-compatible servers).
    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
        arc_inner(&mut self.inner).base_url = base_url.into();
        self
    }

    /// Set the organization header.
    pub fn with_organization(mut self, org: impl Into<String>) -> Self {
        arc_inner(&mut self.inner).organization = Some(org.into());
        self
    }

    fn build_body(&self, values: &[String], options: Option<&EmbeddingGenerationOptions>) -> Value {
        let mut body = Map::new();
        let model = options
            .and_then(|o| o.model.clone())
            .unwrap_or_else(|| self.inner.model.clone());
        body.insert("model".into(), json!(model));
        body.insert("input".into(), json!(values));
        if let Some(options) = options {
            if let Some(dimensions) = options.dimensions {
                body.insert("dimensions".into(), json!(dimensions));
            }
            // Provider-specific extras understood by this endpoint
            // (`encoding_format`, `user`) are forwarded verbatim.
            for key in ["encoding_format", "user"] {
                if let Some(v) = options.additional_properties.get(key) {
                    body.insert(key.into(), v.clone());
                }
            }
        }
        Value::Object(body)
    }
}

/// `Arc::make_mut` over an `Inner` that is not `Clone`-derivable field-wise —
/// manual clone keeps the `reqwest::Client` (itself an `Arc` internally)
/// shared.
fn arc_inner(inner: &mut Arc<Inner>) -> &mut Inner {
    if Arc::strong_count(inner) != 1 {
        *inner = Arc::new(Inner {
            http: inner.http.clone(),
            api_key: inner.api_key.clone(),
            base_url: inner.base_url.clone(),
            model: inner.model.clone(),
            organization: inner.organization.clone(),
        });
    }
    Arc::get_mut(inner).expect("just ensured unique")
}

/// Decode one `embedding` field into a vector.
///
/// The field is a numeric array under the default `encoding_format: "float"`,
/// but a **base64 string** when the caller asks for `encoding_format:
/// "base64"` — a packed little-endian `f32` array, per the OpenAI embeddings
/// API and the Azure AI Inference surface that mirrors it. Every client here
/// forwards `encoding_format` verbatim, so refusing the encoded form would
/// reject a perfectly successful response for a documented option.
fn parse_embedding_vector(field: Option<&Value>) -> Result<Vec<f32>> {
    match field {
        Some(Value::Array(values)) => Ok(values
            .iter()
            .map(|v| v.as_f64().unwrap_or_default() as f32)
            .collect()),
        Some(Value::String(encoded)) => {
            let bytes = BASE64_STANDARD.decode(encoded.as_bytes()).map_err(|e| {
                Error::service(format!("embeddings item has undecodable base64: {e}"))
            })?;
            if bytes.len() % 4 != 0 {
                return Err(Error::service(format!(
                    "base64 embedding decodes to {} bytes, not a whole number of f32s",
                    bytes.len()
                )));
            }
            // `split_first_chunk::<4>` hands back a `&[u8; 4]` directly, so
            // `from_le_bytes` needs no length check or fallible conversion.
            // (`chunks_exact(4)` reads the same but draws clippy's
            // `chunks_exact_to_as_chunks`, whose suggested `as_chunks` only
            // stabilised at this crate's exact MSRV — `split_first_chunk` has
            // been stable since 1.77, well under it.)
            let mut vector = Vec::with_capacity(bytes.len() / 4);
            let mut rest = bytes.as_slice();
            while let Some((chunk, tail)) = rest.split_first_chunk::<4>() {
                vector.push(f32::from_le_bytes(*chunk));
                rest = tail;
            }
            Ok(vector)
        }
        _ => Err(Error::service("embeddings item missing 'embedding' vector")),
    }
}

/// Parse an OpenAI-shaped embeddings response
/// (`{"data": [{"embedding": [...], "index": n}], "model": .., "usage": ..}`)
/// into a [`GeneratedEmbeddings`], restoring input order via `index`.
///
/// `embedding` may be a numeric array (the default `encoding_format:
/// "float"`) or a base64 string of packed little-endian `f32`s
/// (`encoding_format: "base64"`); both decode to the same vector.
pub fn parse_embeddings_response(value: &Value) -> Result<GeneratedEmbeddings> {
    let model = value.get("model").and_then(Value::as_str);
    let data = value
        .get("data")
        .and_then(Value::as_array)
        .ok_or_else(|| Error::service("embeddings response missing 'data' array"))?;

    let mut indexed: Vec<(usize, Embedding)> = Vec::with_capacity(data.len());
    for (position, item) in data.iter().enumerate() {
        let vector = parse_embedding_vector(item.get("embedding"))?;
        let index = item
            .get("index")
            .and_then(Value::as_u64)
            .map(|i| i as usize)
            .unwrap_or(position);
        indexed.push((
            index,
            Embedding {
                vector,
                model: model.map(String::from),
            },
        ));
    }
    indexed.sort_by_key(|(i, _)| *i);

    let mut batch = GeneratedEmbeddings::new(indexed.into_iter().map(|(_, e)| e).collect());
    if let Some(usage) = value.get("usage") {
        let input = usage.get("prompt_tokens").and_then(Value::as_u64);
        let total = usage.get("total_tokens").and_then(Value::as_u64);
        if input.is_some() || total.is_some() {
            batch.usage = Some(UsageDetails {
                input_token_count: input,
                total_token_count: total,
                ..Default::default()
            });
        }
    }
    Ok(batch)
}

#[async_trait::async_trait]
impl EmbeddingClient for OpenAIEmbeddingClient {
    async fn get_embeddings(
        &self,
        values: Vec<String>,
        options: Option<EmbeddingGenerationOptions>,
    ) -> Result<GeneratedEmbeddings> {
        let body = self.build_body(&values, options.as_ref());
        let url = format!("{}/embeddings", self.inner.base_url.trim_end_matches('/'));
        let mut req = self
            .inner
            .http
            .post(&url)
            .bearer_auth(&self.inner.api_key)
            .json(&body);
        if let Some(org) = &self.inner.organization {
            req = req.header("OpenAI-Organization", org);
        }
        let resp = req
            .send()
            .await
            .map_err(|e| Error::service(format!("request failed: {e}")))?;
        if !resp.status().is_success() {
            let status = resp.status();
            let retry_after = parse_retry_after(resp.headers());
            let text = resp.text().await.unwrap_or_default();
            return Err(classify_service_error(
                status.as_u16(),
                &text,
                format!("OpenAI API error {status}: {text}"),
                retry_after,
            ));
        }
        let value: Value = resp
            .json()
            .await
            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
        parse_embeddings_response(&value)
    }

    fn model(&self) -> Option<&str> {
        Some(&self.inner.model)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn build_body_includes_model_input_and_dimensions() {
        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
        let options = EmbeddingGenerationOptions::new().with_dimensions(256);
        let body = client.build_body(&["a".into(), "b".into()], Some(&options));
        assert_eq!(body["model"], "text-embedding-3-small");
        assert_eq!(body["input"], json!(["a", "b"]));
        assert_eq!(body["dimensions"], 256);
    }

    #[test]
    fn build_body_option_model_overrides_default() {
        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
        let options = EmbeddingGenerationOptions::new().with_model("text-embedding-3-large");
        let body = client.build_body(&["a".into()], Some(&options));
        assert_eq!(body["model"], "text-embedding-3-large");
    }

    #[test]
    fn build_body_forwards_known_additional_properties_only() {
        let client = OpenAIEmbeddingClient::new("sk-test", "m");
        let mut options = EmbeddingGenerationOptions::new();
        options
            .additional_properties
            .insert("encoding_format".into(), json!("float"));
        options
            .additional_properties
            .insert("unrelated".into(), json!(true));
        let body = client.build_body(&["a".into()], Some(&options));
        assert_eq!(body["encoding_format"], "float");
        assert!(body.get("unrelated").is_none());
    }

    #[test]
    fn parse_response_restores_index_order_and_usage() {
        let value = json!({
            "model": "text-embedding-3-small",
            "data": [
                { "index": 1, "embedding": [0.3, 0.4] },
                { "index": 0, "embedding": [0.1, 0.2] },
            ],
            "usage": { "prompt_tokens": 5, "total_tokens": 5 }
        });
        let batch = parse_embeddings_response(&value).unwrap();
        assert_eq!(batch.len(), 2);
        assert_eq!(batch[0].vector, vec![0.1, 0.2]);
        assert_eq!(batch[1].vector, vec![0.3, 0.4]);
        assert_eq!(batch[0].model.as_deref(), Some("text-embedding-3-small"));
        let usage = batch.usage.as_ref().unwrap();
        assert_eq!(usage.input_token_count, Some(5));
        assert_eq!(usage.total_token_count, Some(5));
    }

    #[test]
    fn parse_response_missing_data_errors() {
        assert!(parse_embeddings_response(&json!({})).is_err());
    }

    /// Every client here forwards `encoding_format` verbatim, so a caller can
    /// ask for `"base64"` and get embeddings back as packed little-endian f32
    /// strings. Rejecting those would fail a successful response for a
    /// documented option, so the parser accepts both encodings.
    #[test]
    fn parse_response_decodes_base64_embeddings() {
        let vector: [f32; 3] = [1.0, -2.5, 0.125];
        let mut bytes = Vec::new();
        for f in vector {
            bytes.extend_from_slice(&f.to_le_bytes());
        }
        let encoded = BASE64_STANDARD.encode(&bytes);

        let value = json!({
            "data": [{ "index": 0, "embedding": encoded }],
            "model": "text-embedding-3-small",
        });
        let batch = parse_embeddings_response(&value).unwrap();
        assert_eq!(batch.embeddings[0].vector, vector.to_vec());
    }

    #[test]
    fn parse_response_mixes_base64_and_array_items_by_index() {
        let encoded = BASE64_STANDARD.encode(2.0f32.to_le_bytes());
        let value = json!({
            "data": [
                { "index": 1, "embedding": encoded },
                { "index": 0, "embedding": [1.0] },
            ],
        });
        let batch = parse_embeddings_response(&value).unwrap();
        assert_eq!(batch.embeddings[0].vector, vec![1.0]);
        assert_eq!(batch.embeddings[1].vector, vec![2.0]);
    }

    #[test]
    fn parse_response_rejects_undecodable_or_misaligned_base64() {
        // Not valid base64 at all.
        let bad = json!({ "data": [{ "index": 0, "embedding": "!!!not base64!!!" }] });
        assert!(parse_embeddings_response(&bad).is_err());

        // Valid base64, but not a whole number of f32s — better to say so
        // than to silently hand back a truncated vector.
        let misaligned = BASE64_STANDARD.encode([1u8, 2, 3, 4, 5]);
        let value = json!({ "data": [{ "index": 0, "embedding": misaligned }] });
        let err = parse_embeddings_response(&value).unwrap_err();
        assert!(err.to_string().contains("whole number of f32s"), "{err}");
    }
}