dynamo_llm/protocols/openai/
embeddings.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use serde::{Deserialize, Serialize};
17use validator::Validate;
18
19mod aggregator;
20mod nvext;
21
22pub use nvext::{NvExt, NvExtProvider};
23// pub use delta::DeltaGenerator;
24pub use aggregator::DeltaAggregator;
25
26use dynamo_runtime::protocols::annotated::AnnotationsProvider;
27
28#[derive(Serialize, Deserialize, Validate, Debug, Clone)]
29pub struct NvCreateEmbeddingRequest {
30    #[serde(flatten)]
31    pub inner: async_openai::types::CreateEmbeddingRequest,
32
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub nvext: Option<NvExt>,
35}
36
37/// A response structure for unary chat completion responses, embedding OpenAI's
38/// `CreateChatCompletionResponse`.
39///
40/// # Fields
41/// - `inner`: The base OpenAI unary chat completion response, embedded
42///   using `serde(flatten)`.
43#[derive(Serialize, Deserialize, Validate, Debug, Clone)]
44pub struct NvCreateEmbeddingResponse {
45    #[serde(flatten)]
46    pub inner: async_openai::types::CreateEmbeddingResponse,
47}
48
49impl NvCreateEmbeddingResponse {
50    pub fn empty() -> Self {
51        Self {
52            inner: async_openai::types::CreateEmbeddingResponse {
53                object: "list".to_string(),
54                model: "embedding".to_string(),
55                data: vec![],
56                usage: async_openai::types::EmbeddingUsage {
57                    prompt_tokens: 0,
58                    total_tokens: 0,
59                },
60            },
61        }
62    }
63}
64
65/// Implements `NvExtProvider` for `NvCreateEmbeddingRequest`,
66/// providing access to NVIDIA-specific extensions.
67impl NvExtProvider for NvCreateEmbeddingRequest {
68    /// Returns a reference to the optional `NvExt` extension, if available.
69    fn nvext(&self) -> Option<&NvExt> {
70        self.nvext.as_ref()
71    }
72}
73
74/// Implements `AnnotationsProvider` for `NvCreateEmbeddingRequest`,
75/// enabling retrieval and management of request annotations.
76impl AnnotationsProvider for NvCreateEmbeddingRequest {
77    /// Retrieves the list of annotations from `NvExt`, if present.
78    fn annotations(&self) -> Option<Vec<String>> {
79        self.nvext
80            .as_ref()
81            .and_then(|nvext| nvext.annotations.clone())
82    }
83
84    /// Checks whether a specific annotation exists in the request.
85    ///
86    /// # Arguments
87    /// * `annotation` - A string slice representing the annotation to check.
88    ///
89    /// # Returns
90    /// `true` if the annotation exists, `false` otherwise.
91    fn has_annotation(&self, annotation: &str) -> bool {
92        self.nvext
93            .as_ref()
94            .and_then(|nvext| nvext.annotations.as_ref())
95            .map(|annotations| annotations.contains(&annotation.to_string()))
96            .unwrap_or(false)
97    }
98}