ai_sdk_core/embed/

single.rs

use super::EmbeddingConfig;
use crate::error::EmbedError;
use crate::retry::RetryPolicy;
use crate::Result;
use ai_sdk_provider::{EmbedOptions, EmbeddingModel, EmbeddingUsage};
use std::marker::PhantomData;
use std::sync::Arc;

/// Builder for single value embedding
pub struct EmbedBuilder<M, V, Val> {
    model: M,
    value: V,
    config: EmbeddingConfig,
    _marker: PhantomData<Val>,
}

// 1. Initial State
impl<Val> EmbedBuilder<(), (), Val>
where
    Val: Send + Sync + Clone + 'static,
{
    /// Create a new embed builder
    pub fn new() -> Self {
        Self {
            model: (),
            value: (),
            config: EmbeddingConfig::default(),
            _marker: PhantomData,
        }
    }
}

impl<Val> Default for EmbedBuilder<(), (), Val>
where
    Val: Send + Sync + Clone + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

// 2. Configuration Setters (Available in ANY state)
impl<M, V, Val> EmbedBuilder<M, V, Val> {
    /// Set custom retry policy
    pub fn retry_policy(mut self, retry_policy: RetryPolicy) -> Self {
        self.config.retry_policy = retry_policy;
        self
    }
}

// 3. State Transition: Set Model
impl<V, Val> EmbedBuilder<(), V, Val>
where
    Val: Send + Sync + Clone + 'static,
{
    /// Set the embedding model
    pub fn model<Mod: EmbeddingModel<Val> + 'static>(
        self,
        model: Mod,
    ) -> EmbedBuilder<Arc<dyn EmbeddingModel<Val>>, V, Val> {
        EmbedBuilder {
            model: Arc::new(model),
            value: self.value,
            config: self.config,
            _marker: PhantomData,
        }
    }
}

// 4. State Transition: Set Value
impl<M, Val> EmbedBuilder<M, (), Val> {
    /// Set the value to embed
    pub fn value(self, value: Val) -> EmbedBuilder<M, Val, Val> {
        EmbedBuilder {
            model: self.model,
            value,
            config: self.config,
            _marker: PhantomData,
        }
    }
}

// 5. Execution Logic (Requires Model and Value)
impl<Val> EmbedBuilder<Arc<dyn EmbeddingModel<Val>>, Val, Val>
where
    Val: Send + Sync + Clone + 'static,
{
    /// Execute the embedding
    pub async fn execute(self) -> Result<EmbedResult<Val>> {
        let model = self.model;
        let value = self.value;
        let config = self.config;

        // Call model with retry
        let response = config
            .retry_policy
            .retry(|| {
                let options = EmbedOptions {
                    values: vec![value.clone()],
                    provider_options: None,
                    headers: None,
                };
                let model = model.clone();
                async move { model.do_embed(options).await }
            })
            .await
            .map_err(EmbedError::ProviderError)?;

        // Extract first embedding
        let embedding = response
            .embeddings
            .into_iter()
            .next()
            .ok_or(EmbedError::EmptyResponse)?;

        Ok(EmbedResult {
            value,
            embedding,
            usage: response.usage,
        })
    }
}

/// Result of embedding a single value
#[derive(Debug, Clone)]
pub struct EmbedResult<Val> {
    /// The original value that was embedded
    pub value: Val,
    /// The embedding vector
    pub embedding: Vec<f32>,
    /// Token usage information
    pub usage: Option<EmbeddingUsage>,
}

impl<Val> EmbedResult<Val> {
    /// Get the embedding vector
    pub fn embedding(&self) -> &[f32] {
        &self.embedding
    }

    /// Get the original value
    pub fn value(&self) -> &Val {
        &self.value
    }

    /// Get token usage
    pub fn usage(&self) -> Option<&EmbeddingUsage> {
        self.usage.as_ref()
    }
}

/// Entry point function for embedding a single value
pub fn embed<Val>() -> EmbedBuilder<(), (), Val>
where
    Val: Send + Sync + Clone + 'static,
{
    EmbedBuilder::new()
}

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

    #[test]
    fn test_embed_builder_defaults() {
        let builder = embed::<String>();
        let _ = builder; // prevent unused warning
    }

    // Note: Missing model/value tests are removed because they are compile-time errors now.

    #[tokio::test]
    async fn test_embed_success() {
        use ai_sdk_provider::{EmbedResponse, Result};
        use async_trait::async_trait;

        struct DummyModel;
        #[async_trait]
        impl EmbeddingModel<String> for DummyModel {
            fn provider(&self) -> &str {
                "test"
            }
            fn model_id(&self) -> &str {
                "dummy"
            }
            async fn max_embeddings_per_call(&self) -> Option<usize> {
                Some(100)
            }
            async fn supports_parallel_calls(&self) -> bool {
                true
            }
            async fn do_embed(&self, options: EmbedOptions<String>) -> Result<EmbedResponse> {
                assert_eq!(options.values.len(), 1);
                assert_eq!(options.values[0], "test value");
                Ok(EmbedResponse {
                    embeddings: vec![vec![0.1, 0.2, 0.3]],
                    usage: Some(EmbeddingUsage { tokens: 10 }),
                    provider_metadata: None,
                    response: None,
                    // ..Default::default() // If needed
                })
            }
        }

        let result = embed()
            .model(DummyModel)
            .value("test value".to_string())
            .execute()
            .await
            .unwrap();

        assert_eq!(result.value(), "test value");
        assert_eq!(result.embedding().len(), 3);
        assert_eq!(result.embedding(), &[0.1, 0.2, 0.3]);
        assert_eq!(result.usage().unwrap().tokens, 10);
    }
}