1use 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#[derive(Debug)]
24pub struct TranslationOutput {
25 pub body: Value,
26 pub diagnostics: Vec<TranslationDiagnostic>,
27}
28
29#[derive(Debug)]
31pub struct RequestIrOutput {
32 pub request: LlmRequest,
33 pub diagnostics: Vec<TranslationDiagnostic>,
34}
35
36#[derive(Debug)]
38pub struct ResponseIrOutput {
39 pub response: LlmResponse,
40 pub diagnostics: Vec<TranslationDiagnostic>,
41}
42
43#[derive(Default)]
45pub struct FormatRegistry {
46 codecs: BTreeMap<FormatId, Arc<dyn FormatCodec>>,
47}
48
49impl FormatRegistry {
50 pub fn new() -> Self {
52 Self::default()
53 }
54
55 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 pub fn register(&mut self, codec: impl FormatCodec + 'static) {
66 self.codecs.insert(codec.format(), Arc::new(codec));
67 }
68
69 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
79pub 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 pub fn new(registry: FormatRegistry) -> Self {
97 Self {
98 registry,
99 stream_registry: StreamCodecRegistry::with_builtins(),
100 }
101 }
102
103 pub fn with_registries(registry: FormatRegistry, stream_registry: StreamCodecRegistry) -> Self {
105 Self {
106 registry,
107 stream_registry,
108 }
109 }
110
111 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 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 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 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 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 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 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 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
258fn 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}