1use std::collections::BTreeMap;
7use std::time::Duration;
8
9use async_trait::async_trait;
10use reqwest::{Url, header::HeaderValue};
11use serde::Deserialize;
12use serde_json::{Value, json};
13use switchyard_protocol::{
14 BooleanEstimate, DecisionAnswer, DecisionKind, DecisionRequest, DecisionResponse,
15 DecisionValue, ModelId, Probability, ProviderConfidence, RoutedDecisionClient, ScoreValue,
16 Usage,
17};
18
19use crate::client::convert_reqwest_error;
20use crate::{LlmClientError, Result, metrics};
21
22pub struct OpenAiDecisionClient {
24 client: reqwest::Client,
25 endpoint: Url,
26 api_key: String,
27}
28
29impl OpenAiDecisionClient {
30 pub fn new(endpoint: Url, api_key: String, timeout: Duration) -> Result<Self> {
33 if HeaderValue::try_from(format!("Bearer {api_key}")).is_err() {
34 return Err(LlmClientError::Configuration {
35 message: "OpenAI Decisions API key cannot be encoded as an HTTP header".into(),
36 });
37 }
38 let client = reqwest::Client::builder()
39 .timeout(timeout)
40 .redirect(reqwest::redirect::Policy::none())
41 .build()
42 .map_err(convert_reqwest_error)?;
43 Ok(Self {
44 client,
45 endpoint,
46 api_key,
47 })
48 }
49}
50
51#[async_trait]
52impl RoutedDecisionClient for OpenAiDecisionClient {
53 async fn call(&self, request: DecisionRequest) -> Result<DecisionResponse> {
54 let response = self
55 .client
56 .post(self.endpoint.clone())
57 .bearer_auth(&self.api_key)
58 .json(&encode(&request)?)
59 .send()
60 .await
61 .inspect_err(|_| metrics::record_upstream_attempt(None))
62 .map_err(convert_reqwest_error)?;
63 let status = response.status();
64 let body = response
65 .bytes()
66 .await
67 .inspect_err(|_| metrics::record_upstream_attempt(None))
68 .map_err(convert_reqwest_error)?;
69 metrics::record_upstream_attempt(Some(status.as_u16()));
70 if !status.is_success() {
71 return Err(LlmClientError::UpstreamHttp {
72 status,
73 body: String::from_utf8_lossy(&body).replace(&self.api_key, "[REDACTED]"),
74 });
75 }
76 let response: WireResponse =
77 serde_json::from_slice(&body).map_err(|source| LlmClientError::InvalidResponse {
78 source: Box::new(source),
79 })?;
80 let mut questions = request.questions;
81 let mut answers = BTreeMap::new();
82 for answer in response.answers {
83 let invalid = || {
84 LlmClientError::ResponseTranslation(format!(
85 "OpenAI Decisions answer {:?} has an unexpected question name, kind, or value",
86 answer.name,
87 ))
88 };
89 let question = questions.remove(&answer.name).ok_or_else(invalid)?;
90 let value = match (&question.kind, answer.value) {
91 (DecisionKind::Boolean { .. }, WireValue::Predicate { probability }) => {
92 DecisionValue::Boolean(BooleanEstimate::ProbabilityTrue(probability))
93 }
94 (
95 DecisionKind::Choice { options },
96 WireValue::Choice {
97 choice,
98 probabilities,
99 },
100 ) => {
101 let count = probabilities.len();
102 let probabilities: BTreeMap<_, _> = probabilities
103 .into_iter()
104 .map(|p| (p.value, p.probability))
105 .collect();
106 if count != options.len()
107 || probabilities.len() != options.len()
108 || !options.iter().any(|option| option.id == choice)
109 || options
110 .iter()
111 .any(|option| !probabilities.contains_key(&option.id))
112 {
113 return Err(invalid());
114 }
115 DecisionValue::Choice {
116 selected: choice,
117 probabilities: Some(probabilities),
118 }
119 }
120 (
121 DecisionKind::Score { levels },
122 WireValue::Score {
123 score,
124 probabilities,
125 },
126 ) => {
127 let mut probabilities = probabilities;
128 probabilities.sort_by_key(|p| p.value);
129 if levels.is_empty()
130 || !(0.0..=(levels.len() - 1) as f64).contains(&score.0)
131 || probabilities.len() != levels.len()
132 || probabilities
133 .iter()
134 .enumerate()
135 .any(|(index, p)| p.value != index)
136 {
137 return Err(invalid());
138 }
139 DecisionValue::Score {
140 value: score,
141 probabilities: Some(
142 probabilities.into_iter().map(|p| p.probability).collect(),
143 ),
144 }
145 }
146 (_, WireValue::Refusal) => {
147 return Err(LlmClientError::ResponseTranslation(format!(
148 "OpenAI Decisions refused question {:?}",
149 answer.name,
150 )));
151 }
152 _ => return Err(invalid()),
153 };
154 answers.insert(
155 answer.name,
156 DecisionAnswer {
157 value,
158 provider_confidence: answer.confidence,
159 },
160 );
161 }
162 if !questions.is_empty() {
163 return Err(LlmClientError::ResponseTranslation(
164 "OpenAI Decisions answer names do not match the requested question keys".into(),
165 ));
166 }
167 let usage = response.usage;
168 let cached = usage.input_tokens_details.cached_tokens;
169 let created = usage.input_tokens_details.cache_write_tokens;
170 Ok(DecisionResponse {
171 id: None,
172 model: Some(response.model),
173 answers,
174 usage: Usage {
175 input_tokens: Some(
176 usage
177 .input_tokens
178 .saturating_sub(cached)
179 .saturating_sub(created),
180 ),
181 cache: Usage::cache_details(Some(cached), Some(created)),
182 output_tokens: Some(usage.output_tokens),
183 total_tokens: Some(usage.total_tokens),
184 reasoning_tokens: Some(usage.output_tokens_details.reasoning_tokens),
185 ..Usage::default()
186 },
187 })
188 }
189}
190
191fn encode(request: &DecisionRequest) -> Result<Value> {
192 let model = request
193 .model
194 .as_ref()
195 .ok_or_else(|| LlmClientError::InvalidRequest {
196 message: "OpenAI Decisions requires a selected model".into(),
197 })?;
198 let questions: Vec<_> = request
199 .questions
200 .iter()
201 .map(|(name, question)| {
202 let mut wire = json!({"name": name, "instructions": text(&question.instructions)});
203 match &question.kind {
204 DecisionKind::Boolean {
205 true_description,
206 false_description,
207 } => {
208 wire["type"] = json!("predicate");
209 if true_description.is_some() || false_description.is_some() {
210 wire["instructions"] = json!(
211 json!({
212 "instructions": question.instructions,
213 "true": true_description, "false": false_description,
214 })
215 .to_string()
216 );
217 }
218 }
219 DecisionKind::Choice { options } => {
220 wire["type"] = json!("choice");
221 wire["choices"] = Value::Array(
222 options
223 .iter()
224 .map(|option| {
225 let mut choice = json!({"value": option.id});
226 if let Some(description) = &option.description {
227 choice["description"] = json!(text(description));
228 }
229 choice
230 })
231 .collect(),
232 );
233 }
234 DecisionKind::Score { levels } => {
235 wire["type"] = json!("score");
236 wire["levels"] = Value::Array(
237 levels
238 .iter()
239 .map(|level| json!({"label": text(level)}))
240 .collect(),
241 );
242 }
243 }
244 wire
245 })
246 .collect();
247 Ok(json!({"model": model, "input": text(&request.input), "questions": questions}))
250}
251
252fn text(value: &Value) -> String {
254 value
255 .as_str()
256 .map(str::to_owned)
257 .unwrap_or_else(|| value.to_string())
258}
259
260#[derive(Deserialize)]
261struct WireResponse {
262 model: ModelId,
263 answers: Vec<WireAnswer>,
264 usage: WireUsage,
265}
266
267#[derive(Deserialize)]
268struct WireAnswer {
269 name: String,
270 #[serde(flatten)]
271 value: WireValue,
272 confidence: Option<ProviderConfidence>,
273}
274
275#[derive(Deserialize)]
276#[serde(tag = "type", rename_all = "snake_case")]
277enum WireValue {
278 Predicate {
279 probability: Probability,
280 },
281 Choice {
282 choice: String,
283 probabilities: Vec<WireProbability<String>>,
284 },
285 Score {
286 score: ScoreValue,
287 probabilities: Vec<WireProbability<usize>>,
288 },
289 Refusal,
290}
291
292#[derive(Deserialize)]
293struct WireProbability<T> {
294 value: T,
295 probability: Probability,
296}
297
298#[derive(Deserialize)]
299struct WireUsage {
300 input_tokens: u64,
301 output_tokens: u64,
302 total_tokens: u64,
303 input_tokens_details: InputTokenDetails,
304 output_tokens_details: OutputTokenDetails,
305}
306
307#[derive(Deserialize)]
308struct InputTokenDetails {
309 cached_tokens: u64,
310 cache_write_tokens: u64,
311}
312
313#[derive(Deserialize)]
314struct OutputTokenDetails {
315 reasoning_tokens: u64,
316}