liter-llm 1.18.3

use base64::Engine as _;
use serde::{Deserialize, Serialize};

use super::common::{ImageUrl, Usage};
use crate::cost;

/// The format in which the embedding vectors are returned.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum EmbeddingFormat {
    /// 32-bit floating-point numbers (default).
    Float,
    /// Base64-encoded string representation of the floats.
    Base64,
}

/// Embedding request.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EmbeddingRequest {
    /// Model ID (e.g., `"text-embedding-3-small"`).
    pub model: String,
    /// Text, texts, or multimodal content to embed.
    pub input: EmbeddingInput,
    /// Output format: float (native) or base64.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub encoding_format: Option<EmbeddingFormat>,
    /// Requested embedding dimensions (if supported by the model).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub dimensions: Option<u32>,
    /// User identifier for request tracking.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

/// Text, texts, or multimodal content to embed.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum EmbeddingInput {
    /// Single text string.
    Single(String),
    /// Multiple text strings (batch embedding).
    Multiple(Vec<String>),
    /// Text and image parts for a single multimodal embedding.
    Multimodal(Vec<EmbeddingContentPart>),
}

#[cfg_attr(alef, alef(skip))]
impl Default for EmbeddingInput {
    fn default() -> Self {
        Self::Single(String::new())
    }
}

/// A content part in a multimodal embedding input.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum EmbeddingContentPart {
    /// Plain text.
    #[serde(rename = "text")]
    Text { text: String },
    /// Image identified by a data URL or HTTP/HTTPS URL.
    #[serde(rename = "image_url")]
    ImageUrl { image_url: ImageUrl },
    /// Image encoded as a complete data URL.
    #[serde(rename = "image_base64")]
    ImageBase64 { image_base64: String },
}

impl EmbeddingContentPart {
    /// Create a text embedding part.
    pub fn text(text: impl Into<String>) -> Self {
        Self::Text { text: text.into() }
    }

    /// Create an image embedding part from a data URL or HTTP/HTTPS URL.
    pub fn image_url(url: impl Into<String>) -> Self {
        Self::ImageUrl {
            image_url: ImageUrl {
                url: url.into(),
                detail: None,
            },
        }
    }

    /// Create an image embedding part from a complete data URL.
    pub fn image_base64(data_url: impl Into<String>) -> Self {
        Self::ImageBase64 {
            image_base64: data_url.into(),
        }
    }

    /// Create an image embedding part from raw bytes.
    pub fn image_bytes(bytes: &[u8], mime_type: Option<&str>) -> Self {
        Self::image_base64(crate::image::encode_data_url(bytes, mime_type))
    }
}

/// Embedding response.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EmbeddingResponse {
    /// Always `"list"` from OpenAI-compatible APIs.  Stored as a plain
    /// `String` so non-standard provider values do not break deserialization.
    pub object: String,
    /// List of embeddings.
    pub data: Vec<EmbeddingObject>,
    /// Model used to generate embeddings.
    pub model: String,
    /// Token usage (input tokens only; embeddings have zero output tokens).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,
}

impl EmbeddingResponse {
    /// Estimate the cost of this embedding request based on embedded pricing data.
    ///
    /// Returns `None` if:
    /// - the `model` field is not present in the embedded pricing registry, or
    /// - the `usage` field is absent from the response.
    ///
    /// Embedding models only charge for input tokens; output cost is zero.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let cost = response.estimated_cost();
    /// if let Some(usd) = cost {
    ///     println!("Embedding cost: ${usd:.8}");
    /// }
    /// ```
    #[cfg_attr(alef, alef(skip))]
    #[must_use]
    pub fn estimated_cost(&self) -> Option<f64> {
        let usage = self.usage.as_ref()?;
        cost::completion_cost(&self.model, usage.prompt_tokens, usage.completion_tokens)
    }
}

/// A single embedding vector.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EmbeddingObject {
    /// Always `"embedding"` from OpenAI-compatible APIs.  Stored as a plain
    /// `String` so non-standard provider values do not break deserialization.
    pub object: String,
    /// The embedding vector.
    ///
    /// Providers may return this as a JSON float array or, when
    /// `encoding_format: "base64"` was requested, as a base64 string of
    /// little-endian `f32` bytes. Base64 responses are decoded on read; this
    /// field always serializes back out as a JSON float array.
    #[serde(deserialize_with = "deserialize_embedding")]
    pub embedding: Vec<f32>,
    /// Index in the batch (corresponds to input order).
    pub index: u32,
}

/// Deserialization helper for [`EmbeddingObject::embedding`].
///
/// Accepts either a JSON array of floats or a base64-encoded string of
/// little-endian `f32` bytes (the OpenAI-compatible `encoding_format:
/// "base64"` response shape).
///
/// Uses a [`Visitor`](serde::de::Visitor) rather than an untagged enum: an
/// untagged enum buffers the whole value into an intermediate before
/// re-deserializing, which measured ~3x slower on the base64 path. `visit_seq`
/// handles the float array and `visit_str` decodes the base64 string, both with
/// zero buffering.
fn deserialize_embedding<'de, D>(deserializer: D) -> Result<Vec<f32>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    struct EmbeddingVisitor;

    impl<'de> serde::de::Visitor<'de> for EmbeddingVisitor {
        type Value = Vec<f32>;

        fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
            formatter.write_str("a float array or a base64-encoded string of little-endian f32 bytes")
        }

        fn visit_str<E>(self, value: &str) -> Result<Vec<f32>, E>
        where
            E: serde::de::Error,
        {
            let bytes = base64::engine::general_purpose::STANDARD
                .decode(value)
                .map_err(|e| E::custom(format!("invalid base64 embedding: {e}")))?;
            if bytes.len() % 4 != 0 {
                return Err(E::custom(format!(
                    "base64 embedding length {} is not a multiple of 4",
                    bytes.len()
                )));
            }
            Ok(bytes
                .chunks_exact(4)
                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                .collect())
        }

        fn visit_seq<A>(self, mut seq: A) -> Result<Vec<f32>, A::Error>
        where
            A: serde::de::SeqAccess<'de>,
        {
            let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(0));
            while let Some(value) = seq.next_element::<f32>()? {
                out.push(value);
            }
            Ok(out)
        }
    }

    deserializer.deserialize_any(EmbeddingVisitor)
}

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

    #[test]
    fn multimodal_input_round_trips_text_and_image_url() {
        let input = EmbeddingInput::Multimodal(vec![
            EmbeddingContentPart::text("a red bicycle"),
            EmbeddingContentPart::image_url("https://example.com/bicycle.png"),
        ]);

        let json = serde_json::to_string(&input).expect("serialization should not fail");
        assert_eq!(
            json,
            r#"[{"type":"text","text":"a red bicycle"},{"type":"image_url","image_url":{"url":"https://example.com/bicycle.png"}}]"#
        );
        let parsed: EmbeddingInput = serde_json::from_str(&json).expect("deserialization should not fail");
        assert_eq!(parsed, input);
    }

    #[test]
    fn image_bytes_encode_as_data_url() {
        let part = EmbeddingContentPart::image_bytes(b"image", Some("image/png"));
        assert_eq!(
            part,
            EmbeddingContentPart::ImageBase64 {
                image_base64: "data:image/png;base64,aW1hZ2U=".into(),
            }
        );
    }

    fn embedding_body(embedding_json: &str) -> String {
        format!(r#"{{"object":"embedding","index":0,"embedding":{embedding_json}}}"#)
    }

    #[test]
    fn base64_embedding_round_trips_bit_exact() {
        let src: [f32; 5] = [1.0, -2.5, 12.375, 0.0, f32::MIN_POSITIVE];
        let mut bytes = Vec::with_capacity(src.len() * 4);
        for v in src {
            bytes.extend_from_slice(&v.to_le_bytes());
        }
        let encoded = base64::engine::general_purpose::STANDARD.encode(&bytes);
        let body = embedding_body(&format!("{encoded:?}"));

        let obj: EmbeddingObject = serde_json::from_str(&body).expect("base64 embedding should deserialize");
        assert_eq!(obj.embedding, src, "decoded floats must match source bit-exactly");
    }

    #[test]
    fn openai_base64_anchor_decodes_to_one() {
        let body = embedding_body(r#""AACAPw==""#);
        let obj: EmbeddingObject = serde_json::from_str(&body).expect("anchor base64 embedding should deserialize");
        assert_eq!(obj.embedding, vec![1.0_f32]);
    }

    #[test]
    fn float_array_body_still_parses() {
        let body = embedding_body("[0.1,0.2,0.3]");
        let obj: EmbeddingObject = serde_json::from_str(&body).expect("float array embedding should deserialize");
        assert_eq!(obj.embedding, vec![0.1_f32, 0.2, 0.3]);
    }

    #[test]
    fn odd_length_base64_errors_with_multiple_of_four_message() {
        let encoded = base64::engine::general_purpose::STANDARD.encode(b"abcdef");
        let body = embedding_body(&format!("{encoded:?}"));

        let err = serde_json::from_str::<EmbeddingObject>(&body).expect_err("6-byte base64 payload must error");
        assert!(
            err.to_string().contains("not a multiple of 4"),
            "expected multiple-of-4 message, got: {err}"
        );
    }

    #[test]
    fn invalid_base64_errors() {
        let body = embedding_body(r#""not valid base64!!!""#);
        let err = serde_json::from_str::<EmbeddingObject>(&body).expect_err("non-base64 string must error");
        assert!(
            err.to_string().contains("invalid base64 embedding"),
            "expected invalid base64 message, got: {err}"
        );
    }

    #[test]
    fn serialization_stays_a_json_array() {
        let body = embedding_body(r#""AACAPw==""#);
        let obj: EmbeddingObject = serde_json::from_str(&body).expect("anchor base64 embedding should deserialize");

        let serialized = serde_json::to_value(&obj).expect("serialize back to JSON");
        assert!(
            serialized["embedding"].is_array(),
            "expected embedding field to serialize as an array, got: {serialized}"
        );
    }
}