dynamo-llm 1.4.0

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::NvCreateEmbeddingResponse;
use crate::protocols::{
    Annotated,
    codec::{Message, SseCodecError},
    convert_sse_stream,
    openai::stream_aggregator::{StreamAggregable, aggregate_stream},
};

use dynamo_runtime::{engine::DataStream, error::DynamoError};
use futures::Stream;

impl StreamAggregable for NvCreateEmbeddingResponse {
    fn empty() -> Self {
        Self::empty()
    }

    fn merge(&mut self, next: Self) {
        self.inner.data.extend(next.inner.data);
        self.inner.usage.prompt_tokens += next.inner.usage.prompt_tokens;
        self.inner.usage.total_tokens += next.inner.usage.total_tokens;
    }
}

impl NvCreateEmbeddingResponse {
    /// Converts an SSE stream into a [`NvCreateEmbeddingResponse`].
    pub async fn from_sse_stream(
        stream: DataStream<Result<Message, SseCodecError>>,
    ) -> Result<NvCreateEmbeddingResponse, DynamoError> {
        let stream = convert_sse_stream::<NvCreateEmbeddingResponse>(stream);
        NvCreateEmbeddingResponse::from_annotated_stream(stream).await
    }

    /// Aggregates an annotated stream of embedding responses into a final response.
    pub async fn from_annotated_stream(
        stream: impl Stream<Item = Annotated<NvCreateEmbeddingResponse>>,
    ) -> Result<NvCreateEmbeddingResponse, DynamoError> {
        aggregate_stream(stream).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dynamo_runtime::error::{BackendError, ErrorType};
    use futures::stream;

    fn create_test_embedding_response(
        embeddings: Vec<dynamo_protocols::types::Embedding>,
        prompt_tokens: u32,
        total_tokens: u32,
    ) -> Annotated<NvCreateEmbeddingResponse> {
        let response = NvCreateEmbeddingResponse {
            inner: dynamo_protocols::types::CreateEmbeddingResponse {
                object: "list".to_string(),
                model: "test-model".to_string(),
                data: embeddings,
                usage: dynamo_protocols::types::EmbeddingUsage {
                    prompt_tokens,
                    total_tokens,
                },
            },
        };

        Annotated::from_data(response)
    }

    #[tokio::test]
    async fn test_empty_stream() {
        let stream = stream::empty();
        let result = NvCreateEmbeddingResponse::from_annotated_stream(Box::pin(stream)).await;

        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.inner.data.len(), 0);
        assert_eq!(response.inner.object, "list");
        assert_eq!(response.inner.model, "embedding");
    }

    #[tokio::test]
    async fn test_single_embedding() {
        let embedding = dynamo_protocols::types::Embedding {
            index: 0,
            object: "embedding".to_string(),
            embedding: dynamo_protocols::types::EmbeddingVector::Float(vec![0.1, 0.2, 0.3]),
        };

        let annotated = create_test_embedding_response(vec![embedding.clone()], 10, 10);
        let stream = stream::iter(vec![annotated]);

        let result = NvCreateEmbeddingResponse::from_annotated_stream(Box::pin(stream)).await;

        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.inner.data.len(), 1);
        assert_eq!(response.inner.data[0].index, 0);
        assert_eq!(
            response.inner.data[0].embedding,
            dynamo_protocols::types::EmbeddingVector::Float(vec![0.1, 0.2, 0.3])
        );
        assert_eq!(response.inner.usage.prompt_tokens, 10);
        assert_eq!(response.inner.usage.total_tokens, 10);
    }

    #[tokio::test]
    async fn test_multiple_embeddings() {
        let embedding1 = dynamo_protocols::types::Embedding {
            index: 0,
            object: "embedding".to_string(),
            embedding: dynamo_protocols::types::EmbeddingVector::Float(vec![0.1, 0.2, 0.3]),
        };

        let embedding2 = dynamo_protocols::types::Embedding {
            index: 1,
            object: "embedding".to_string(),
            embedding: dynamo_protocols::types::EmbeddingVector::Float(vec![0.4, 0.5, 0.6]),
        };

        let annotated1 = create_test_embedding_response(vec![embedding1.clone()], 5, 5);
        let annotated2 = create_test_embedding_response(vec![embedding2.clone()], 7, 7);
        let stream = stream::iter(vec![annotated1, annotated2]);

        let result = NvCreateEmbeddingResponse::from_annotated_stream(Box::pin(stream)).await;

        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.inner.data.len(), 2);
        assert_eq!(response.inner.data[0].index, 0);
        assert_eq!(response.inner.data[1].index, 1);
        assert_eq!(response.inner.usage.prompt_tokens, 12);
        assert_eq!(response.inner.usage.total_tokens, 12);
    }

    #[tokio::test]
    async fn test_error_in_stream() {
        let error_annotated =
            Annotated::<NvCreateEmbeddingResponse>::from_error("Test error".to_string());
        let stream = stream::iter(vec![error_annotated]);

        let result = NvCreateEmbeddingResponse::from_annotated_stream(Box::pin(stream)).await;

        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Test error"));
    }

    #[tokio::test]
    async fn test_typed_error_after_data_is_preserved() {
        let embedding = dynamo_protocols::types::Embedding {
            index: 0,
            object: "embedding".to_string(),
            embedding: dynamo_protocols::types::EmbeddingVector::Float(vec![0.1]),
        };
        let response = create_test_embedding_response(vec![embedding], 1, 1);
        let error = Annotated::<NvCreateEmbeddingResponse> {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: None,
            error: Some(
                DynamoError::builder()
                    .error_type(ErrorType::Backend(BackendError::InvalidArgument))
                    .message("invalid embedding request")
                    .build(),
            ),
        };

        let result =
            NvCreateEmbeddingResponse::from_annotated_stream(stream::iter(vec![response, error]))
                .await;

        let error = result.expect_err("typed error must win over partial data");
        assert_eq!(
            error.error_type(),
            ErrorType::Backend(BackendError::InvalidArgument)
        );
        assert_eq!(error.message(), "invalid embedding request");
    }

    #[tokio::test]
    async fn test_base64_embeddings_aggregate() {
        // Verifies that `data[].embedding` can be a base64 string, not
        // just a float array.
        let embedding = dynamo_protocols::types::Embedding {
            index: 0,
            object: "embedding".to_string(),
            embedding: dynamo_protocols::types::EmbeddingVector::Base64(
                "AAAAAAAAgD8AAABA".to_string(), // [0.0, 1.0, 2.0] as little-endian f32 bytes
            ),
        };
        let annotated = create_test_embedding_response(vec![embedding], 3, 3);

        // Round-trip through serde to prove the wire format also matches:
        // the Python handler emits JSON, the aggregator deserializes it.
        let json = serde_json::to_string(&annotated.data.as_ref().unwrap()).unwrap();
        let parsed: NvCreateEmbeddingResponse = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.inner.data.len(), 1);
        match &parsed.inner.data[0].embedding {
            dynamo_protocols::types::EmbeddingVector::Base64(s) => {
                assert_eq!(s, "AAAAAAAAgD8AAABA");
            }
            other => panic!("expected base64 variant, got {:?}", other),
        }

        let stream = stream::iter(vec![annotated]);
        let result = NvCreateEmbeddingResponse::from_annotated_stream(Box::pin(stream)).await;
        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.inner.data.len(), 1);
        match &response.inner.data[0].embedding {
            dynamo_protocols::types::EmbeddingVector::Base64(s) => {
                assert_eq!(s, "AAAAAAAAgD8AAABA");
            }
            other => panic!("expected base64 variant, got {:?}", other),
        }
    }
}