nemo_relay/codec/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! LLM codec types, traits, and built-in implementations.
5//!
6//! This module provides the type system and traits for bidirectional
7//! request codecs ([`traits::LlmCodec`] / [`request::AnnotatedLlmRequest`]),
8//! the decode-only response codec
9//! ([`traits::LlmResponseCodec`] / [`response::AnnotatedLlmResponse`]), and
10//! the streaming response codec
11//! ([`streaming::StreamingCodec`]) used with the managed
12//! streaming LLM execution pipeline.
13//!
14//! [`resolve`] is the detect-then-decode entry point for selecting a built-in
15//! provider codec from a raw payload when no codec annotation is present.
16
17pub mod anthropic;
18pub mod model_pricing;
19pub mod openai_chat;
20pub mod openai_responses;
21pub mod optimization;
22pub mod request;
23pub mod resolve;
24pub mod response;
25pub mod streaming;
26pub mod traits;
27
28use nemo_relay_types::Json;
29
30use crate::error::{FlowError, Result};
31
32fn optional_bool(
33    obj: &serde_json::Map<String, Json>,
34    key: &str,
35    surface: &str,
36) -> Result<Option<bool>> {
37    match obj.get(key) {
38        Some(Json::Null) | None => Ok(None),
39        Some(Json::Bool(value)) => Ok(Some(*value)),
40        Some(_) => Err(FlowError::InvalidArgument(format!(
41            "{surface} {key} must be a boolean or null"
42        ))),
43    }
44}
45
46fn optional_u64(
47    obj: &serde_json::Map<String, Json>,
48    key: &str,
49    surface: &str,
50) -> Result<Option<u64>> {
51    match obj.get(key) {
52        Some(Json::Null) | None => Ok(None),
53        Some(value) if value.as_u64().is_some() => Ok(value.as_u64()),
54        Some(_) => Err(FlowError::InvalidArgument(format!(
55            "{surface} {key} must be a non-negative integer or null"
56        ))),
57    }
58}
59
60fn optional_i64(
61    obj: &serde_json::Map<String, Json>,
62    key: &str,
63    surface: &str,
64) -> Result<Option<i64>> {
65    match obj.get(key) {
66        Some(Json::Null) | None => Ok(None),
67        Some(value) if value.as_i64().is_some() => Ok(value.as_i64()),
68        Some(_) => Err(FlowError::InvalidArgument(format!(
69            "{surface} {key} must be an integer or null"
70        ))),
71    }
72}
73
74fn optional_f64(
75    obj: &serde_json::Map<String, Json>,
76    key: &str,
77    surface: &str,
78) -> Result<Option<f64>> {
79    match obj.get(key) {
80        Some(Json::Null) | None => Ok(None),
81        Some(value) if value.as_f64().is_some() => Ok(value.as_f64()),
82        Some(_) => Err(FlowError::InvalidArgument(format!(
83            "{surface} {key} must be a number or null"
84        ))),
85    }
86}
87
88fn optional_string(
89    obj: &serde_json::Map<String, Json>,
90    key: &str,
91    surface: &str,
92) -> Result<Option<String>> {
93    match obj.get(key) {
94        Some(Json::Null) | None => Ok(None),
95        Some(Json::String(value)) => Ok(Some(value.clone())),
96        Some(_) => Err(FlowError::InvalidArgument(format!(
97            "{surface} {key} must be a string or null"
98        ))),
99    }
100}
101
102fn optional_object(
103    obj: &serde_json::Map<String, Json>,
104    key: &str,
105    surface: &str,
106) -> Result<Option<Json>> {
107    match obj.get(key) {
108        Some(Json::Null) | None => Ok(None),
109        Some(value @ Json::Object(_)) => Ok(Some(value.clone())),
110        Some(_) => Err(FlowError::InvalidArgument(format!(
111            "{surface} {key} must be an object or null"
112        ))),
113    }
114}
115
116fn optional_array(
117    obj: &serde_json::Map<String, Json>,
118    key: &str,
119    surface: &str,
120) -> Result<Option<Json>> {
121    match obj.get(key) {
122        Some(Json::Null) | None => Ok(None),
123        Some(value @ Json::Array(_)) => Ok(Some(value.clone())),
124        Some(_) => Err(FlowError::InvalidArgument(format!(
125            "{surface} {key} must be an array or null"
126        ))),
127    }
128}
129
130fn encode_changed_items<T, F>(
131    edited: &[T],
132    baseline: &[T],
133    original: Option<&[Json]>,
134    encode: F,
135) -> Result<Vec<Json>>
136where
137    T: PartialEq,
138    F: FnMut(&T) -> Result<Json>,
139{
140    encode_changed_items_with_patch(
141        edited,
142        baseline,
143        original,
144        encode,
145        |original, _, _, baseline_value, edited_value| {
146            patch_changed_json(original, baseline_value, edited_value)
147        },
148    )
149}
150
151fn encode_changed_items_with_patch<T, F, P>(
152    edited: &[T],
153    baseline: &[T],
154    original: Option<&[Json]>,
155    mut encode: F,
156    mut patch: P,
157) -> Result<Vec<Json>>
158where
159    T: PartialEq,
160    F: FnMut(&T) -> Result<Json>,
161    P: FnMut(&Json, &T, &T, &Json, &Json) -> Result<Json>,
162{
163    let alignment = if original.is_some() {
164        align_changed_items(edited, baseline)?
165    } else {
166        vec![None; edited.len()]
167    };
168    edited
169        .iter()
170        .enumerate()
171        .map(|(index, item)| {
172            if let Some(baseline_index) = alignment[index] {
173                let baseline_item = &baseline[baseline_index];
174                if let Some(original_item) = original.and_then(|items| items.get(baseline_index)) {
175                    if baseline_item == item {
176                        return Ok(original_item.clone());
177                    }
178                    let baseline_value = encode(baseline_item)?;
179                    let edited_value = encode(item)?;
180                    return patch(
181                        original_item,
182                        baseline_item,
183                        item,
184                        &baseline_value,
185                        &edited_value,
186                    );
187                }
188            }
189            encode(item)
190        })
191        .collect()
192}
193
194fn align_changed_items<T: PartialEq>(edited: &[T], baseline: &[T]) -> Result<Vec<Option<usize>>> {
195    let mut alignment = vec![None; edited.len()];
196    let mut used = vec![false; baseline.len()];
197
198    for index in 0..edited.len().min(baseline.len()) {
199        if edited[index] == baseline[index] {
200            alignment[index] = Some(index);
201            used[index] = true;
202        }
203    }
204
205    let mut structurally_changed = edited.len() != baseline.len();
206    for (edited_index, item) in edited.iter().enumerate() {
207        if alignment[edited_index].is_some() {
208            continue;
209        }
210        if let Some(baseline_index) = baseline
211            .iter()
212            .enumerate()
213            .position(|(baseline_index, candidate)| !used[baseline_index] && candidate == item)
214        {
215            alignment[edited_index] = Some(baseline_index);
216            used[baseline_index] = true;
217            structurally_changed |= baseline_index != edited_index;
218        }
219    }
220
221    let unmatched_edited = alignment
222        .iter()
223        .enumerate()
224        .filter_map(|(index, item)| item.is_none().then_some(index))
225        .collect::<Vec<_>>();
226    let unmatched_baseline = used
227        .iter()
228        .enumerate()
229        .filter_map(|(index, item)| (!item).then_some(index))
230        .collect::<Vec<_>>();
231
232    if unmatched_edited.len() > 1 && unmatched_baseline.len() > 1 {
233        return Err(FlowError::InvalidArgument(
234            "cannot safely preserve provider fields for multiple edited array items without stable identities"
235                .into(),
236        ));
237    }
238
239    if !structurally_changed && edited.len() == baseline.len() {
240        for index in unmatched_edited {
241            alignment[index] = Some(index);
242        }
243    }
244
245    Ok(alignment)
246}
247
248fn patch_changed_array(original: &[Json], baseline: &[Json], edited: &[Json]) -> Result<Vec<Json>> {
249    let alignment = align_changed_items(edited, baseline)?;
250    edited
251        .iter()
252        .enumerate()
253        .map(|(edited_index, edited_value)| {
254            if let Some(baseline_index) = alignment[edited_index]
255                && let (Some(original_value), Some(baseline_value)) =
256                    (original.get(baseline_index), baseline.get(baseline_index))
257            {
258                return patch_changed_json(original_value, baseline_value, edited_value);
259            }
260            Ok(edited_value.clone())
261        })
262        .collect()
263}
264
265fn patch_changed_json(original: &Json, baseline: &Json, edited: &Json) -> Result<Json> {
266    if baseline == edited {
267        return Ok(original.clone());
268    }
269
270    match (original, baseline, edited) {
271        (Json::Object(original), Json::Object(baseline), Json::Object(edited)) => {
272            if ["type", "role"]
273                .into_iter()
274                .any(|key| baseline.get(key) != edited.get(key))
275            {
276                return Ok(Json::Object(edited.clone()));
277            }
278            let mut patched = original.clone();
279            for key in baseline.keys().filter(|key| !edited.contains_key(*key)) {
280                patched.remove(key);
281            }
282            for (key, edited_value) in edited {
283                if baseline.get(key) == Some(edited_value) {
284                    continue;
285                }
286                let value = match (original.get(key), baseline.get(key)) {
287                    (Some(original_value), Some(baseline_value)) => {
288                        patch_changed_json(original_value, baseline_value, edited_value)?
289                    }
290                    _ => edited_value.clone(),
291                };
292                patched.insert(key.clone(), value);
293            }
294            Ok(Json::Object(patched))
295        }
296        (Json::Array(original), Json::Array(baseline), Json::Array(edited)) => Ok(Json::Array(
297            patch_changed_array(original, baseline, edited)?,
298        )),
299        _ => Ok(edited.clone()),
300    }
301}
302
303#[cfg(test)]
304#[path = "../../tests/unit/codec/parity_tests.rs"]
305mod parity_tests;
nemo_relay/codec/mod.rs

nemo_relay/codec/
mod.rs
