Skip to main content

nemo_relay/codec/
traits.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! LLM codec traits for bidirectional request translation.
5
6use crate::api::llm::LlmRequest;
7use crate::api::runtime::LlmCodecIdentity;
8use crate::error::Result;
9use crate::json::Json;
10
11use super::request::AnnotatedLlmRequest;
12use super::response::AnnotatedLlmResponse;
13
14// ---------------------------------------------------------------------------
15// LlmCodec trait
16// ---------------------------------------------------------------------------
17
18/// A bidirectional translator between opaque [`LlmRequest`] content and
19/// structured [`AnnotatedLlmRequest`].
20///
21/// Codecs are implemented by framework integrations and provider adapters
22/// because each SDK has its own request format. A codec is supplied per call by
23/// the caller; the built-in provider codecs can also be selected from a raw
24/// payload via [`crate::codec::resolve`].
25///
26/// # Design
27///
28/// - **Synchronous**: `decode`/`encode` are pure data transforms (JSON
29///   restructuring), not I/O operations. This matches existing guardrails
30///   and request intercepts.
31/// - **`Send + Sync`**: Required because [`NemoRelayContextState`](crate::api::runtime::NemoRelayContextState)
32///   is behind `Arc<RwLock<>>` and accessed from async contexts.
33/// - **Trait object**: Codecs are registered at runtime by callers or bindings,
34///   so the Rust core cannot know concrete types at compile time.
35///   Store as `Arc<dyn LlmCodec>`.
36pub trait LlmCodec: Send + Sync {
37    /// Return this codec's identity for LLM sanitizer context.
38    ///
39    /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless
40    /// they have a stable runtime registration ID. Callers must not infer a
41    /// provider surface from an opaque implementation or request shape.
42    fn codec_identity(&self) -> LlmCodecIdentity {
43        LlmCodecIdentity::Opaque
44    }
45
46    /// Parse opaque request content into structured form.
47    fn decode(&self, request: &LlmRequest) -> Result<AnnotatedLlmRequest>;
48
49    /// Merge structured changes back into the opaque request.
50    ///
51    /// The `original` parameter is the pre-intercept [`LlmRequest`], used to
52    /// preserve fields that the Codec does not structurally model. Implementations
53    /// MUST use merge-not-replace semantics: overlay structured changes onto
54    /// the original content, do not construct a fresh content object.
55    fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result<LlmRequest>;
56}
57
58// ---------------------------------------------------------------------------
59// LlmResponseCodec trait
60// ---------------------------------------------------------------------------
61
62/// Decode-only codec for LLM API responses.
63///
64/// Unlike [`LlmCodec`] (which is bidirectional for requests), response codecs
65/// are introspection-only: they parse a raw response into structured form but
66/// never need to encode back. This matches the pipeline design where responses
67/// are observed, not modified.
68///
69/// # Design
70///
71/// - **Synchronous**: `decode_response` is a pure data transform (JSON parsing),
72///   not an I/O operation.
73/// - **`Send + Sync`**: Required for storage in `Arc` behind `RwLock`.
74/// - **Trait object**: Codecs are registered at runtime, stored as
75///   `Arc<dyn LlmResponseCodec>`.
76/// - **Fallible**: Returns `Result`; managed call sites may omit annotations on
77///   decode failure, while manual lifecycle bindings may surface the error.
78///
79/// # Two-Phase Decode
80///
81/// Implementations should use a two-phase decode pattern:
82/// 1. Deserialize raw JSON into API-specific intermediate structs
83/// 2. Map intermediate structs into the normalized `AnnotatedLlmResponse`
84pub trait LlmResponseCodec: Send + Sync {
85    /// Return this codec's identity for LLM sanitizer context.
86    ///
87    /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless
88    /// they have a stable runtime registration ID.
89    fn codec_identity(&self) -> LlmCodecIdentity {
90        LlmCodecIdentity::Opaque
91    }
92
93    /// Parse a raw JSON response into normalized structured form.
94    ///
95    /// Implementations should return `Err` only for genuinely unparseable input.
96    fn decode_response(&self, response: &Json) -> Result<AnnotatedLlmResponse>;
97}
98
99/// Additive encoder for materializing a normalized response in a provider wire format.
100///
101/// Existing response codecs remain decode-only. Components that perform cross-protocol dispatch
102/// can supply this companion trait without changing the managed observability pipeline.
103pub trait LlmResponseEncoder: Send + Sync {
104    /// Encode a normalized response into the target provider's buffered JSON representation.
105    fn encode_response(&self, response: &AnnotatedLlmResponse) -> Result<Json>;
106}