Skip to main content

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::codecs::anthropic::AnthropicMessagesCodec;
12use crate::codecs::openai_chat::OpenAiChatCodec;
13use crate::codecs::responses::OpenAiResponsesCodec;
14use crate::codecs::stream::{StreamCodecRegistry, StreamTranslationState};
15use crate::codecs::FormatCodec;
16use crate::diagnostic::TranslationDiagnostic;
17use crate::error::{Result, TranslationError};
18use crate::format::FormatId;
19use crate::llm::{LlmRequest, LlmResponse};
20use crate::policy::TranslationPolicy;
21
22/// Encoded translation result with any diagnostics emitted along the way.
23#[derive(Debug)]
24pub struct TranslationOutput {
25    pub body: Value,
26    pub diagnostics: Vec<TranslationDiagnostic>,
27}
28
29/// Decoded request IR plus diagnostics.
30#[derive(Debug)]
31pub struct RequestIrOutput {
32    pub request: LlmRequest,
33    pub diagnostics: Vec<TranslationDiagnostic>,
34}
35
36/// Decoded response IR plus diagnostics.
37#[derive(Debug)]
38pub struct ResponseIrOutput {
39    pub response: LlmResponse,
40    pub diagnostics: Vec<TranslationDiagnostic>,
41}
42
43/// Registry mapping wire formats to buffered codecs.
44#[derive(Default)]
45pub struct FormatRegistry {
46    codecs: BTreeMap<FormatId, Arc<dyn FormatCodec>>,
47}
48
49impl FormatRegistry {
50    /// Creates an empty format registry.
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// Creates a registry populated with the built-in provider codecs.
56    pub fn with_builtins() -> Self {
57        let mut registry = Self::new();
58        registry.register(OpenAiChatCodec);
59        registry.register(AnthropicMessagesCodec);
60        registry.register(OpenAiResponsesCodec);
61        registry
62    }
63
64    /// Registers or replaces a codec for its declared format.
65    pub fn register(&mut self, codec: impl FormatCodec + 'static) {
66        self.codecs.insert(codec.format(), Arc::new(codec));
67    }
68
69    /// Looks up a codec by format identifier.
70    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn FormatCodec>> {
71        let format = format.into();
72        self.codecs
73            .get(&format)
74            .cloned()
75            .ok_or_else(|| TranslationError::Other(format!("no codec registered for {format}")))
76    }
77}
78
79/// Stateless request/response translator that routes through the neutral IR.
80pub struct TranslationEngine {
81    registry: FormatRegistry,
82    stream_registry: StreamCodecRegistry,
83}
84
85impl Default for TranslationEngine {
86    fn default() -> Self {
87        Self {
88            registry: FormatRegistry::with_builtins(),
89            stream_registry: StreamCodecRegistry::with_builtins(),
90        }
91    }
92}
93
94impl TranslationEngine {
95    /// Creates an engine from an explicit buffered codec registry.
96    pub fn new(registry: FormatRegistry) -> Self {
97        Self {
98            registry,
99            stream_registry: StreamCodecRegistry::with_builtins(),
100        }
101    }
102
103    /// Creates an engine from explicit buffered and streaming codec registries.
104    pub fn with_registries(registry: FormatRegistry, stream_registry: StreamCodecRegistry) -> Self {
105        Self {
106            registry,
107            stream_registry,
108        }
109    }
110
111    /// Decodes a request body into the neutral request IR.
112    pub fn decode_request(
113        &self,
114        source: impl Into<FormatId>,
115        body: &Value,
116        policy: &TranslationPolicy,
117    ) -> Result<RequestIrOutput> {
118        let source = source.into();
119        let decoded = self.registry.codec(source)?.decode_request(body, policy)?;
120        Ok(RequestIrOutput {
121            request: decoded.request,
122            diagnostics: decoded.diagnostics,
123        })
124    }
125
126    /// Encodes a neutral request IR into a target wire format.
127    pub fn encode_request(
128        &self,
129        target: impl Into<FormatId>,
130        request: &LlmRequest,
131        policy: &TranslationPolicy,
132    ) -> Result<TranslationOutput> {
133        let target = target.into();
134        let encoded = self
135            .registry
136            .codec(target)?
137            .encode_request(request, policy)?;
138        Ok(TranslationOutput {
139            body: encoded.body,
140            diagnostics: encoded.diagnostics,
141        })
142    }
143
144    /// Translates a request body from source format to target format.
145    pub fn translate_request(
146        &self,
147        source: impl Into<FormatId>,
148        target: impl Into<FormatId>,
149        body: &Value,
150        policy: &TranslationPolicy,
151    ) -> Result<TranslationOutput> {
152        let source = source.into();
153        let target = target.into();
154        let decoded = self
155            .registry
156            .codec(source.clone())?
157            .decode_request(body, policy)?;
158        let encoded = self
159            .registry
160            .codec(target.clone())?
161            .encode_request(&decoded.request, policy)?;
162        Ok(TranslationOutput {
163            body: encoded.body,
164            diagnostics: with_formats(decoded.diagnostics, encoded.diagnostics, source, target),
165        })
166    }
167
168    /// Decodes a response body into the neutral response IR.
169    pub fn decode_response(
170        &self,
171        source: impl Into<FormatId>,
172        body: &Value,
173        policy: &TranslationPolicy,
174    ) -> Result<ResponseIrOutput> {
175        let source = source.into();
176        let decoded = self.registry.codec(source)?.decode_response(body, policy)?;
177        Ok(ResponseIrOutput {
178            response: decoded.response,
179            diagnostics: decoded.diagnostics,
180        })
181    }
182
183    /// Encodes a neutral response IR into a target wire format.
184    pub fn encode_response(
185        &self,
186        target: impl Into<FormatId>,
187        response: &LlmResponse,
188        policy: &TranslationPolicy,
189    ) -> Result<TranslationOutput> {
190        let target = target.into();
191        let encoded = self
192            .registry
193            .codec(target)?
194            .encode_response(response, policy)?;
195        Ok(TranslationOutput {
196            body: encoded.body,
197            diagnostics: encoded.diagnostics,
198        })
199    }
200
201    /// Translates a response body from source format to target format.
202    pub fn translate_response(
203        &self,
204        source: impl Into<FormatId>,
205        target: impl Into<FormatId>,
206        body: &Value,
207        policy: &TranslationPolicy,
208    ) -> Result<TranslationOutput> {
209        let source = source.into();
210        let target = target.into();
211        let decoded = self
212            .registry
213            .codec(source.clone())?
214            .decode_response(body, policy)?;
215        let encoded = self
216            .registry
217            .codec(target.clone())?
218            .encode_response(&decoded.response, policy)?;
219        Ok(TranslationOutput {
220            body: encoded.body,
221            diagnostics: with_formats(decoded.diagnostics, encoded.diagnostics, source, target),
222        })
223    }
224
225    /// Translates one streaming source event into zero or more target events.
226    pub fn translate_event(
227        &self,
228        state: &mut StreamTranslationState,
229        source: impl Into<FormatId>,
230        target: impl Into<FormatId>,
231        event: &Value,
232    ) -> Result<Vec<Value>> {
233        let source = source.into();
234        let target = target.into();
235        let source_codec = self.stream_registry.codec(source.clone())?;
236        let target_codec = self.stream_registry.codec(target.clone())?;
237        let canonical = source_codec.decode_event(state, event);
238        state.source = Some(source);
239        state.target = Some(target);
240        Ok(canonical
241            .into_iter()
242            .flat_map(|event| target_codec.encode_event(state, event))
243            .collect())
244    }
245
246    /// Finishes target-provider stream emission after the source stream closes.
247    pub fn finish_stream(
248        &self,
249        state: &mut StreamTranslationState,
250        target: impl Into<FormatId>,
251    ) -> Result<Vec<Value>> {
252        let target = target.into();
253        let target_codec = self.stream_registry.codec(target)?;
254        Ok(target_codec.finish(state))
255    }
256}
257
258// Attaches source and target formats to every diagnostic emitted across both passes.
259fn with_formats(
260    decoded: Vec<TranslationDiagnostic>,
261    encoded: Vec<TranslationDiagnostic>,
262    source: FormatId,
263    target: FormatId,
264) -> Vec<TranslationDiagnostic> {
265    decoded
266        .into_iter()
267        .chain(encoded)
268        .map(|diagnostic| diagnostic.with_formats(source.clone(), target.clone()))
269        .collect()
270}