switchyard_translation/
engine.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Registry-backed translation engine for buffered requests and responses.
5
6use std::collections::BTreeMap;
7use std::sync::Arc;
8
9use serde_json::Value;
10
11use crate::LlmResponseStreamEvent;
12use crate::codecs::FormatCodec;
13use crate::codecs::anthropic::AnthropicMessagesCodec;
14use crate::codecs::openai_chat::OpenAiChatCodec;
15use crate::codecs::responses::OpenAiResponsesCodec;
16use crate::codecs::stream::{
17    StreamCodecRegistry, StreamTranslationState, encode_response_stream_event,
18};
19use crate::diagnostic::TranslationDiagnostic;
20use crate::error::{Result, TranslationError};
21use crate::format::FormatId;
22use crate::llm::{AggLlmResponse, LlmRequest, ProviderExtensions};
23use crate::policy::TranslationPolicy;
24
25/// Encoded translation result with any diagnostics emitted along the way.
26#[derive(Debug)]
27pub struct TranslationOutput {
28    pub body: Value,
29    pub diagnostics: Vec<TranslationDiagnostic>,
30}
31
32/// Decoded request IR plus diagnostics.
33#[derive(Debug)]
34pub struct RequestIrOutput {
35    pub request: LlmRequest,
36    pub diagnostics: Vec<TranslationDiagnostic>,
37}
38
39/// Decoded response IR plus diagnostics.
40#[derive(Debug)]
41pub struct ResponseIrOutput {
42    pub response: AggLlmResponse,
43    pub diagnostics: Vec<TranslationDiagnostic>,
44}
45
46/// Registry mapping wire formats to buffered codecs.
47#[derive(Default)]
48pub struct FormatRegistry {
49    codecs: BTreeMap<FormatId, Arc<dyn FormatCodec>>,
50}
51
52impl FormatRegistry {
53    /// Creates an empty format registry.
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    /// Creates a registry populated with the built-in provider codecs.
59    pub fn with_builtins() -> Self {
60        let mut registry = Self::new();
61        registry.register(OpenAiChatCodec);
62        registry.register(AnthropicMessagesCodec);
63        registry.register(OpenAiResponsesCodec);
64        registry
65    }
66
67    /// Registers or replaces a codec for its declared format.
68    pub fn register(&mut self, codec: impl FormatCodec + 'static) {
69        self.codecs.insert(codec.format(), Arc::new(codec));
70    }
71
72    /// Looks up a codec by format identifier.
73    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn FormatCodec>> {
74        let format = format.into();
75        self.codecs
76            .get(&format)
77            .cloned()
78            .ok_or_else(|| TranslationError::Other(format!("no codec registered for {format}")))
79    }
80}
81
82/// Stateless request/response translator that routes through the neutral IR.
83pub struct TranslationEngine {
84    registry: FormatRegistry,
85    stream_registry: StreamCodecRegistry,
86}
87
88impl Default for TranslationEngine {
89    fn default() -> Self {
90        Self {
91            registry: FormatRegistry::with_builtins(),
92            stream_registry: StreamCodecRegistry::with_builtins(),
93        }
94    }
95}
96
97impl TranslationEngine {
98    /// Creates an engine from an explicit buffered codec registry.
99    pub fn new(registry: FormatRegistry) -> Self {
100        Self {
101            registry,
102            stream_registry: StreamCodecRegistry::with_builtins(),
103        }
104    }
105
106    /// Creates an engine from explicit buffered and streaming codec registries.
107    pub fn with_registries(registry: FormatRegistry, stream_registry: StreamCodecRegistry) -> Self {
108        Self {
109            registry,
110            stream_registry,
111        }
112    }
113
114    /// Decodes a request body into the neutral request IR.
115    pub fn decode_request(
116        &self,
117        source: impl Into<FormatId>,
118        body: &Value,
119        policy: &TranslationPolicy,
120    ) -> Result<RequestIrOutput> {
121        let source = source.into();
122        let decoded = self.registry.codec(source)?.decode_request(body, policy)?;
123        Ok(RequestIrOutput {
124            request: decoded.request,
125            diagnostics: decoded.diagnostics,
126        })
127    }
128
129    /// Encodes a neutral request IR into a target wire format.
130    ///
131    /// Preserved target bodies take precedence over normalized fields, including
132    /// `output.is_schema_enforced`. Clear the target's preservation entry or use
133    /// [`crate::policy::PreservationPolicy::Disabled`] to encode IR edits.
134    pub fn encode_request(
135        &self,
136        target: impl Into<FormatId>,
137        request: &LlmRequest,
138        policy: &TranslationPolicy,
139    ) -> Result<TranslationOutput> {
140        let target = target.into();
141        let encoded = self
142            .registry
143            .codec(target)?
144            .encode_request(request, policy)?;
145        Ok(TranslationOutput {
146            body: encoded.body,
147            diagnostics: encoded.diagnostics,
148        })
149    }
150
151    /// Translates a request body from source format to target format.
152    pub fn translate_request(
153        &self,
154        source: impl Into<FormatId>,
155        target: impl Into<FormatId>,
156        body: &Value,
157        policy: &TranslationPolicy,
158    ) -> Result<TranslationOutput> {
159        let source = source.into();
160        let target = target.into();
161        let decoded = self
162            .registry
163            .codec(source.clone())?
164            .decode_request(body, policy)?;
165        let encoded = self
166            .registry
167            .codec(target.clone())?
168            .encode_request(&decoded.request, policy)?;
169        Ok(TranslationOutput {
170            body: encoded.body,
171            diagnostics: with_formats(decoded.diagnostics, encoded.diagnostics, source, target),
172        })
173    }
174
175    /// Decodes a response body into the neutral response IR.
176    pub fn decode_response(
177        &self,
178        source: impl Into<FormatId>,
179        body: &Value,
180        policy: &TranslationPolicy,
181    ) -> Result<ResponseIrOutput> {
182        let source = source.into();
183        let decoded = self.registry.codec(source)?.decode_response(body, policy)?;
184        Ok(ResponseIrOutput {
185            response: decoded.response,
186            diagnostics: decoded.diagnostics,
187        })
188    }
189
190    /// Encodes a neutral response IR into a target wire format.
191    pub fn encode_response(
192        &self,
193        target: impl Into<FormatId>,
194        response: &AggLlmResponse,
195        policy: &TranslationPolicy,
196    ) -> Result<TranslationOutput> {
197        let target = target.into();
198        let encoded = self
199            .registry
200            .codec(target)?
201            .encode_response(response, policy)?;
202        Ok(TranslationOutput {
203            body: encoded.body,
204            diagnostics: encoded.diagnostics,
205        })
206    }
207
208    /// Encodes a neutral response IR while applying metadata retained from its request.
209    ///
210    /// Request extensions are used to restore caller-specific response details, such as
211    /// Codex tool namespaces, without exposing those details to the upstream provider.
212    pub fn encode_response_with_extensions(
213        &self,
214        target: impl Into<FormatId>,
215        response: &AggLlmResponse,
216        request_extensions: &ProviderExtensions,
217        policy: &TranslationPolicy,
218    ) -> Result<TranslationOutput> {
219        let mut output = self.encode_response(target, response, policy)?;
220        crate::codex_namespaces::restore_qualified_tool_names(
221            &mut output.body,
222            &crate::codex_namespaces::qualified_tool_origins(request_extensions),
223        );
224        crate::codex_custom_tools::restore_custom_tool_calls(
225            &mut output.body,
226            &crate::codex_custom_tools::custom_tool_names(request_extensions),
227        );
228        Ok(output)
229    }
230
231    /// Translates a response body from source format to target format.
232    pub fn translate_response(
233        &self,
234        source: impl Into<FormatId>,
235        target: impl Into<FormatId>,
236        body: &Value,
237        policy: &TranslationPolicy,
238    ) -> Result<TranslationOutput> {
239        let source = source.into();
240        let target = target.into();
241        let decoded = self
242            .registry
243            .codec(source.clone())?
244            .decode_response(body, policy)?;
245        let encoded = self
246            .registry
247            .codec(target.clone())?
248            .encode_response(&decoded.response, policy)?;
249        Ok(TranslationOutput {
250            body: encoded.body,
251            diagnostics: with_formats(decoded.diagnostics, encoded.diagnostics, source, target),
252        })
253    }
254
255    /// Translates one streaming source event into zero or more target events.
256    pub fn translate_event(
257        &self,
258        state: &mut StreamTranslationState,
259        source: impl Into<FormatId>,
260        target: impl Into<FormatId>,
261        event: &Value,
262    ) -> Result<Vec<Value>> {
263        let source = source.into();
264        let target = target.into();
265        let source_codec = self.stream_registry.codec(source.clone())?;
266        let target_codec = self.stream_registry.codec(target.clone())?;
267        if let Err(error) =
268            crate::codecs::responses::validate_stream_output(&source, &target, event)
269        {
270            return Ok(target_codec.encode_event(
271                state,
272                crate::LlmResponseChunk::DecodeError {
273                    message: error.to_string(),
274                },
275            ));
276        }
277        let canonical = source_codec.decode_event(state, event);
278        state.source = Some(source);
279        state.target = Some(target);
280        Ok(canonical
281            .into_iter()
282            .flat_map(|event| target_codec.encode_event(state, event))
283            .collect())
284    }
285
286    /// Decodes one provider event while retaining its parsed source JSON value.
287    ///
288    /// Takes ownership of `event` so preservation does not deep-copy provider JSON
289    /// on the per-event streaming path.
290    pub fn decode_stream_event(
291        &self,
292        state: &mut StreamTranslationState,
293        source: impl Into<FormatId>,
294        event: Value,
295    ) -> Result<LlmResponseStreamEvent> {
296        let source = source.into();
297        let source_codec = self.stream_registry.codec(source.clone())?;
298        state.source = Some(source.clone());
299        let normalized = source_codec.decode_event(state, &event);
300        Ok(LlmResponseStreamEvent::preserved(source, event, normalized))
301    }
302
303    /// Encodes one neutral or preserved stream event for a target provider.
304    ///
305    /// A preserved event replays its retained JSON value unchanged when its
306    /// source and target formats match. Cross-format encoding intentionally
307    /// uses only its normalized events.
308    pub fn encode_stream_event(
309        &self,
310        state: &mut StreamTranslationState,
311        target: impl Into<FormatId>,
312        event: LlmResponseStreamEvent,
313    ) -> Result<Vec<Value>> {
314        let target = target.into();
315        let target_codec = self.stream_registry.codec(target.clone())?;
316        state.target = Some(target.clone());
317        Ok(encode_response_stream_event(
318            state,
319            target_codec.as_ref(),
320            &target,
321            event,
322        ))
323    }
324
325    /// Finishes target-provider stream emission after the source stream closes.
326    pub fn finish_stream(
327        &self,
328        state: &mut StreamTranslationState,
329        target: impl Into<FormatId>,
330    ) -> Result<Vec<Value>> {
331        let target = target.into();
332        let target_codec = self.stream_registry.codec(target)?;
333        Ok(target_codec.finish(state))
334    }
335}
336
337// Attaches source and target formats to every diagnostic emitted across both passes.
338fn with_formats(
339    decoded: Vec<TranslationDiagnostic>,
340    encoded: Vec<TranslationDiagnostic>,
341    source: FormatId,
342    target: FormatId,
343) -> Vec<TranslationDiagnostic> {
344    decoded
345        .into_iter()
346        .chain(encoded)
347        .map(|diagnostic| diagnostic.with_formats(source.clone(), target.clone()))
348        .collect()
349}
switchyard_translation/engine.rs

switchyard_translation/
engine.rs
