1use reqwest::Method;
4use serde::Serialize;
5use serde_json::{Map, Value};
6
7use crate::client::encode::request_options;
8use crate::client::{Client, meta::parse_json};
9use crate::error::{Error, Result};
10use crate::http::{HttpResponse, headers, metadata};
11use crate::retry::{RequestOptions, RetryPolicy, RetryState};
12use crate::types::{
13 ChatCompletion, ChatMessage, CostEstimate, GenerateResult, Grammar, ImageInput,
14 RequestMetadata, ResponseInputMessage, ResponseResult,
15};
16
17pub(crate) fn escape_model(model: &str) -> String {
19 model.replace('/', "__")
20}
21
22fn params_object<T: Serialize>(params: &T) -> Result<Map<String, Value>> {
24 match serde_json::to_value(params)
25 .map_err(|err| Error::invalid(format!("could not encode request parameters: {err}")))?
26 {
27 Value::Object(map) => Ok(map),
28 other => Err(Error::invalid(format!(
29 "request parameters must serialize to an object, got {other}"
30 ))),
31 }
32}
33
34pub(crate) fn images_for_json(images: &[ImageInput]) -> Result<Vec<Value>> {
36 use base64::Engine as _;
37
38 images
39 .iter()
40 .map(|image| {
41 let (data, format) = image.resolve()?;
42 Ok(serde_json::json!({
43 "data": base64::engine::general_purpose::STANDARD.encode(&data),
44 "format": format,
45 }))
46 })
47 .collect()
48}
49
50impl Client {
51 pub(crate) async fn send_json<T: serde::de::DeserializeOwned>(
53 &self,
54 path: &str,
55 body: &Value,
56 policy: RetryPolicy,
57 model: Option<&str>,
58 options: &RequestOptions,
59 owner: &str,
60 ) -> Result<(T, Option<RequestMetadata>)> {
61 let response = self
62 .send_json_raw(path, body, policy, model, options)
63 .await?;
64 let decoded: T = parse_json(&response.0, owner)?;
65 let envelope: Value = parse_json(&response.0, owner)?;
66 Ok((
67 decoded,
68 metadata::parse(&response.0.headers, Some(&envelope), response.1),
69 ))
70 }
71
72 async fn send_json_raw(
73 &self,
74 path: &str,
75 body: &Value,
76 policy: RetryPolicy,
77 model: Option<&str>,
78 options: &RequestOptions,
79 ) -> Result<(HttpResponse, u32)> {
80 let routing = self.routing(options.gpu.as_deref());
81 let encoded = serde_json::to_vec(body)
82 .map_err(|err| Error::invalid(format!("could not encode the request body: {err}")))?;
83
84 let request = self
85 .request(Method::POST, path)?
86 .json_headers()
87 .maybe_header(headers::MACHINE_PROFILE, routing.profile.as_deref())
88 .maybe_header(headers::POOL, routing.pool.as_deref())
89 .body(encoded);
90
91 let mut state = RetryState::new(policy, options, model);
92 let response = self.send(request, &mut state).await?;
93 let retries = state.retries();
94 Ok((response, retries))
95 }
96
97 pub async fn estimate(&self, endpoint: &str, request: Value) -> Result<CostEstimate> {
102 if !endpoint.starts_with('/') {
103 return Err(Error::invalid(format!(
104 "estimate endpoint must be an absolute path, got {endpoint:?}"
105 )));
106 }
107 if !request.is_object() {
108 return Err(Error::invalid("estimate request must be a JSON object"));
109 }
110
111 let envelope = serde_json::json!({"endpoint": endpoint, "request": request});
112 let encoded = serde_json::to_vec(&envelope).map_err(|err| {
113 Error::invalid(format!("could not encode the estimate envelope: {err}"))
114 })?;
115 let prepared = self
116 .request(Method::POST, "/v1/estimate")?
117 .json_headers()
118 .body(encoded);
119
120 let options = self.metadata_options();
121 let mut state = RetryState::new(RetryPolicy::ESTIMATE, &options, None);
122 let response = self.send(prepared, &mut state).await?;
123 parse_json(&response, "estimate")
124 }
125
126 pub fn generate(
128 &self,
129 model: impl Into<String>,
130 prompt: impl Into<String>,
131 max_new_tokens: u32,
132 ) -> GenerateRequest {
133 GenerateRequest {
134 client: self.clone(),
135 model: model.into(),
136 prompt: prompt.into(),
137 images: Vec::new(),
138 params: GenerateParams {
139 max_new_tokens,
140 ..GenerateParams::default()
141 },
142 extra_body: None,
143 options: self.request_options(),
144 }
145 }
146
147 pub fn chat(
149 &self,
150 model: impl Into<String>,
151 messages: impl IntoIterator<Item = ChatMessage>,
152 ) -> ChatRequest {
153 ChatRequest {
154 client: self.clone(),
155 model: model.into(),
156 messages: messages.into_iter().collect(),
157 params: ChatParams::default(),
158 extra_body: None,
159 options: self.request_options(),
160 }
161 }
162
163 pub fn responses(&self, model: impl Into<String>, input: ResponseInput) -> ResponsesRequest {
165 ResponsesRequest {
166 client: self.clone(),
167 model: model.into(),
168 input,
169 params: ResponseParams::default(),
170 options: self.request_options(),
171 }
172 }
173}
174
175macro_rules! setters {
177 ($($(#[$meta:meta])* $name:ident: $ty:ty),* $(,)?) => {
178 $(
179 $(#[$meta])*
180 pub fn $name(mut self, value: $ty) -> Self {
181 self.params.$name = Some(value);
182 self
183 }
184 )*
185 };
186}
187
188macro_rules! into_setters {
190 ($($(#[$meta:meta])* $name:ident: $ty:ty),* $(,)?) => {
191 $(
192 $(#[$meta])*
193 pub fn $name(mut self, value: impl Into<$ty>) -> Self {
194 self.params.$name = Some(value.into());
195 self
196 }
197 )*
198 };
199}
200
201macro_rules! stop_setter {
203 () => {
204 pub fn stop(mut self, stop: impl IntoIterator<Item = impl Into<String>>) -> Self {
206 self.params.stop = Some(stop.into_iter().map(Into::into).collect());
207 self
208 }
209 };
210}
211
212#[derive(Debug, Clone, Default, Serialize)]
214pub(crate) struct GenerateParams {
215 pub(crate) max_new_tokens: u32,
216 #[serde(skip_serializing_if = "Option::is_none")]
217 pub(crate) temperature: Option<f64>,
218 #[serde(skip_serializing_if = "Option::is_none")]
219 pub(crate) top_p: Option<f64>,
220 #[serde(skip_serializing_if = "Option::is_none")]
221 pub(crate) stop: Option<Vec<String>>,
222 #[serde(skip_serializing_if = "Option::is_none")]
223 pub(crate) frequency_penalty: Option<f64>,
224 #[serde(skip_serializing_if = "Option::is_none")]
225 pub(crate) presence_penalty: Option<f64>,
226 #[serde(skip_serializing_if = "Option::is_none")]
227 pub(crate) grammar: Option<Grammar>,
228 #[serde(skip_serializing_if = "Option::is_none")]
229 pub(crate) seed: Option<i64>,
230 #[serde(skip_serializing_if = "Option::is_none")]
231 pub(crate) logit_bias: Option<Value>,
232 #[serde(skip_serializing_if = "Option::is_none")]
233 pub(crate) routing_key: Option<String>,
234 #[serde(skip_serializing_if = "Option::is_none")]
235 pub(crate) prompt_cache_key: Option<String>,
236 #[serde(skip_serializing_if = "Option::is_none")]
237 pub(crate) safety_identifier: Option<String>,
238 #[serde(skip_serializing_if = "Option::is_none")]
239 pub(crate) lora_adapter: Option<String>,
240 #[serde(skip_serializing_if = "Option::is_none")]
241 pub(crate) options: Option<Value>,
242 #[serde(skip_serializing_if = "Option::is_none")]
243 pub(crate) logprobs: Option<bool>,
244 #[serde(skip_serializing_if = "Option::is_none")]
245 pub(crate) top_logprobs: Option<u32>,
246}
247
248pub struct GenerateRequest {
250 pub(crate) client: Client,
251 pub(crate) model: String,
252 pub(crate) prompt: String,
253 pub(crate) images: Vec<ImageInput>,
254 pub(crate) params: GenerateParams,
255 pub(crate) extra_body: Option<Value>,
256 pub(crate) options: RequestOptions,
257}
258
259impl GenerateRequest {
260 request_options!();
261
262 setters! {
263 temperature: f64,
265 top_p: f64,
267 frequency_penalty: f64,
269 presence_penalty: f64,
271 grammar: Grammar,
273 seed: i64,
275 logit_bias: Value,
277 options: Value,
279 }
280
281 into_setters! {
282 routing_key: String,
284 prompt_cache_key: String,
286 safety_identifier: String,
288 lora_adapter: String,
290 }
291
292 stop_setter!();
293
294 pub fn images(mut self, images: impl IntoIterator<Item = ImageInput>) -> Self {
296 self.images = images.into_iter().collect();
297 self
298 }
299
300 pub fn extra_body(mut self, extra: Value) -> Self {
302 self.extra_body = Some(extra);
303 self
304 }
305
306 pub(crate) fn body(&self, stream: bool) -> Result<Value> {
307 let mut body = params_object(&self.params)?;
308 body.insert("prompt".to_string(), Value::String(self.prompt.clone()));
309 if !self.images.is_empty() {
310 body.insert(
311 "images".to_string(),
312 Value::Array(images_for_json(&self.images)?),
313 );
314 }
315 if stream {
316 body.insert("stream".to_string(), Value::Bool(true));
317 }
318 if self.params.logprobs != Some(true) {
320 body.remove("top_logprobs");
321 }
322 merge_extra_body(&mut body, self.extra_body.as_ref());
323 body.insert("prompt".to_string(), Value::String(self.prompt.clone()));
325 body.insert(
326 "max_new_tokens".to_string(),
327 Value::from(self.params.max_new_tokens),
328 );
329 if stream {
330 body.insert("stream".to_string(), Value::Bool(true));
331 }
332 Ok(Value::Object(body))
333 }
334
335 pub(crate) fn path(&self) -> String {
336 format!("/v1/generate/{}", escape_model(&self.model))
337 }
338
339 pub async fn send(self) -> Result<GenerateResult> {
341 let body = self.body(false)?;
342 let (mut result, request): (GenerateResult, _) = self
343 .client
344 .send_json(
345 &self.path(),
346 &body,
347 RetryPolicy::GENERATE,
348 Some(&self.model),
349 &self.options,
350 "generate",
351 )
352 .await?;
353 result.request = request;
354 Ok(result)
355 }
356
357 pub async fn estimate(self) -> Result<CostEstimate> {
359 let body = self.body(false)?;
360 self.client.estimate(&self.path(), body).await
361 }
362}
363
364#[derive(Debug, Clone, Default, Serialize)]
366pub(crate) struct ChatParams {
367 #[serde(skip_serializing_if = "Option::is_none")]
368 pub(crate) max_completion_tokens: Option<u32>,
369 #[serde(skip_serializing_if = "Option::is_none")]
370 pub(crate) max_tokens: Option<u32>,
371 #[serde(skip_serializing_if = "Option::is_none")]
372 pub(crate) temperature: Option<f64>,
373 #[serde(skip_serializing_if = "Option::is_none")]
374 pub(crate) top_p: Option<f64>,
375 #[serde(skip_serializing_if = "Option::is_none")]
376 pub(crate) top_k: Option<u32>,
377 #[serde(skip_serializing_if = "Option::is_none")]
378 pub(crate) repetition_penalty: Option<f64>,
379 #[serde(skip_serializing_if = "Option::is_none")]
380 pub(crate) stop: Option<Vec<String>>,
381 #[serde(skip_serializing_if = "Option::is_none")]
382 pub(crate) tools: Option<Vec<Value>>,
383 #[serde(skip_serializing_if = "Option::is_none")]
384 pub(crate) tool_choice: Option<Value>,
385 #[serde(skip_serializing_if = "Option::is_none")]
386 pub(crate) parallel_tool_calls: Option<bool>,
387 #[serde(skip_serializing_if = "Option::is_none")]
388 pub(crate) response_format: Option<Value>,
389 #[serde(skip_serializing_if = "Option::is_none")]
390 pub(crate) frequency_penalty: Option<f64>,
391 #[serde(skip_serializing_if = "Option::is_none")]
392 pub(crate) presence_penalty: Option<f64>,
393 #[serde(skip_serializing_if = "Option::is_none")]
394 pub(crate) n: Option<u32>,
395 #[serde(skip_serializing_if = "Option::is_none")]
396 pub(crate) best_of: Option<u32>,
397 #[serde(skip_serializing_if = "Option::is_none")]
398 pub(crate) logprobs: Option<bool>,
399 #[serde(skip_serializing_if = "Option::is_none")]
400 pub(crate) top_logprobs: Option<u32>,
401 #[serde(skip_serializing_if = "Option::is_none")]
402 pub(crate) logit_bias: Option<Value>,
403 #[serde(skip_serializing_if = "Option::is_none")]
404 pub(crate) seed: Option<i64>,
405 #[serde(skip_serializing_if = "Option::is_none")]
406 pub(crate) user: Option<String>,
407 #[serde(skip_serializing_if = "Option::is_none")]
408 pub(crate) safety_identifier: Option<String>,
409 #[serde(skip_serializing_if = "Option::is_none")]
410 pub(crate) lora_adapter: Option<String>,
411 #[serde(skip_serializing_if = "Option::is_none")]
412 pub(crate) stream_options: Option<Value>,
413}
414
415pub struct ChatRequest {
417 pub(crate) client: Client,
418 pub(crate) model: String,
419 pub(crate) messages: Vec<ChatMessage>,
420 pub(crate) params: ChatParams,
421 pub(crate) extra_body: Option<Value>,
422 pub(crate) options: RequestOptions,
423}
424
425impl ChatRequest {
426 request_options!();
427
428 setters! {
429 max_completion_tokens: u32,
431 max_tokens: u32,
433 temperature: f64,
435 top_p: f64,
437 top_k: u32,
439 repetition_penalty: f64,
441 tools: Vec<Value>,
443 tool_choice: Value,
445 parallel_tool_calls: bool,
447 response_format: Value,
449 frequency_penalty: f64,
451 presence_penalty: f64,
453 n: u32,
455 best_of: u32,
457 logprobs: bool,
459 top_logprobs: u32,
461 logit_bias: Value,
463 seed: i64,
465 stream_options: Value,
467 }
468
469 into_setters! {
470 user: String,
472 safety_identifier: String,
474 lora_adapter: String,
476 }
477
478 stop_setter!();
479
480 pub fn extra_body(mut self, extra: Value) -> Self {
484 self.extra_body = Some(extra);
485 self
486 }
487
488 pub(crate) fn body(&self, stream: bool) -> Result<Value> {
489 let mut params = self.params.clone();
490 if stream {
491 params.best_of = None;
494 }
495 let mut body = params_object(¶ms)?;
496 body.insert("model".to_string(), Value::String(self.model.clone()));
497 body.insert(
498 "messages".to_string(),
499 serde_json::to_value(&self.messages)
500 .map_err(|err| Error::invalid(format!("could not encode chat messages: {err}")))?,
501 );
502 if stream {
503 body.insert("stream".to_string(), Value::Bool(true));
504 }
505 merge_extra_body(&mut body, self.extra_body.as_ref());
506 Ok(Value::Object(body))
507 }
508
509 pub async fn send(self) -> Result<ChatCompletion> {
511 let body = self.body(false)?;
512 let (mut result, request): (ChatCompletion, _) = self
513 .client
514 .send_json(
515 "/v1/chat/completions",
516 &body,
517 RetryPolicy::STREAM,
518 Some(&self.model),
519 &self.options,
520 "chat completion",
521 )
522 .await?;
523 result.request = request;
524 Ok(result)
525 }
526
527 pub async fn estimate(self) -> Result<CostEstimate> {
529 let body = self.body(false)?;
530 self.client.estimate("/v1/chat/completions", body).await
531 }
532}
533
534#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
536#[serde(untagged)]
537pub enum ResponseInput {
538 Text(String),
540 Messages(Vec<ResponseInputMessage>),
542}
543
544impl From<&str> for ResponseInput {
545 fn from(value: &str) -> Self {
546 Self::Text(value.to_string())
547 }
548}
549
550impl From<String> for ResponseInput {
551 fn from(value: String) -> Self {
552 Self::Text(value)
553 }
554}
555
556impl From<Vec<ResponseInputMessage>> for ResponseInput {
557 fn from(value: Vec<ResponseInputMessage>) -> Self {
558 Self::Messages(value)
559 }
560}
561
562#[derive(Debug, Clone, Default, Serialize)]
564pub(crate) struct ResponseParams {
565 #[serde(skip_serializing_if = "Option::is_none")]
566 pub(crate) max_output_tokens: Option<u32>,
567 #[serde(skip_serializing_if = "Option::is_none")]
568 pub(crate) temperature: Option<f64>,
569 #[serde(skip_serializing_if = "Option::is_none")]
570 pub(crate) top_p: Option<f64>,
571 #[serde(skip_serializing_if = "Option::is_none")]
572 pub(crate) seed: Option<i64>,
573}
574
575pub struct ResponsesRequest {
577 client: Client,
578 model: String,
579 input: ResponseInput,
580 params: ResponseParams,
581 options: RequestOptions,
582}
583
584impl ResponsesRequest {
585 request_options!();
586
587 setters! {
588 max_output_tokens: u32,
590 temperature: f64,
592 top_p: f64,
594 seed: i64,
596 }
597
598 fn body(&self) -> Result<Value> {
599 let mut body = params_object(&self.params)?;
600 body.insert("model".to_string(), Value::String(self.model.clone()));
601 body.insert(
602 "input".to_string(),
603 serde_json::to_value(&self.input).map_err(|err| {
604 Error::invalid(format!("could not encode the response input: {err}"))
605 })?,
606 );
607 Ok(Value::Object(body))
608 }
609
610 pub async fn send(self) -> Result<ResponseResult> {
612 let body = self.body()?;
613 let (mut result, request): (ResponseResult, _) = self
614 .client
615 .send_json(
616 "/v1/responses",
617 &body,
618 RetryPolicy::STREAM,
619 Some(&self.model),
620 &self.options,
621 "response",
622 )
623 .await?;
624 result.request = request;
625 Ok(result)
626 }
627
628 pub async fn estimate(self) -> Result<CostEstimate> {
630 let body = self.body()?;
631 self.client.estimate("/v1/responses", body).await
632 }
633}
634
635fn merge_extra_body(body: &mut Map<String, Value>, extra: Option<&Value>) {
636 if let Some(Value::Object(extra)) = extra {
637 for (key, value) in extra {
638 body.insert(key.clone(), value.clone());
639 }
640 }
641}
642
643#[cfg(test)]
644mod tests {
645 use super::*;
646 use serde_json::json;
647
648 fn client() -> Client {
649 Client::new("https://sie.invalid").unwrap()
650 }
651
652 #[test]
653 fn generate_paths_escape_slashes_in_the_model_id() {
654 assert_eq!(escape_model("BAAI/bge-m3"), "BAAI__bge-m3");
655 assert_eq!(escape_model("qwen3"), "qwen3");
656 let request = client().generate("org/model", "hi", 16);
657 assert_eq!(request.path(), "/v1/generate/org__model");
658 }
659
660 #[test]
661 fn generate_body_carries_only_what_was_set() {
662 let body = client().generate("m", "Once upon", 32).body(false).unwrap();
663 assert_eq!(body, json!({"prompt": "Once upon", "max_new_tokens": 32}));
664 }
665
666 #[test]
667 fn generate_body_includes_every_optional_that_was_set() {
668 let body = client()
669 .generate("m", "p", 8)
670 .temperature(0.7)
671 .top_p(0.9)
672 .stop(vec!["\n".to_string()])
673 .grammar(Grammar::regex("[0-9]+"))
674 .seed(42)
675 .routing_key("tenant-1")
676 .lora_adapter("sql")
677 .body(false)
678 .unwrap();
679 assert_eq!(body["temperature"], json!(0.7));
680 assert_eq!(body["top_p"], json!(0.9));
681 assert_eq!(body["stop"], json!(["\n"]));
682 assert_eq!(body["grammar"], json!({"regex": "[0-9]+"}));
683 assert_eq!(body["seed"], json!(42));
684 assert_eq!(body["routing_key"], json!("tenant-1"));
685 assert_eq!(body["lora_adapter"], json!("sql"));
686 }
687
688 #[test]
689 fn extra_body_cannot_rewrite_the_fields_that_define_the_request() {
690 let body = client()
691 .generate("m", "real prompt", 16)
692 .extra_body(json!({"prompt": "hijacked", "max_new_tokens": 9999, "custom": true}))
693 .body(true)
694 .unwrap();
695 assert_eq!(body["prompt"], json!("real prompt"));
696 assert_eq!(body["max_new_tokens"], json!(16));
697 assert_eq!(body["stream"], json!(true));
698 assert_eq!(body["custom"], json!(true));
699 }
700
701 #[test]
702 fn top_logprobs_is_dropped_unless_logprobs_is_on() {
703 let mut request = client().generate("m", "p", 4);
704 request.params.top_logprobs = Some(5);
705 assert!(request.body(true).unwrap().get("top_logprobs").is_none());
706
707 let mut request = client().generate("m", "p", 4);
708 request.params.logprobs = Some(true);
709 request.params.top_logprobs = Some(5);
710 assert_eq!(request.body(true).unwrap()["top_logprobs"], json!(5));
711 }
712
713 #[test]
714 fn chat_body_always_names_the_model_and_messages() {
715 let body = client()
716 .chat("qwen3", [ChatMessage::user("hi")])
717 .temperature(0.2)
718 .body(false)
719 .unwrap();
720 assert_eq!(
721 body,
722 json!({"model": "qwen3", "messages": [{"role": "user", "content": "hi"}], "temperature": 0.2})
723 );
724 assert!(body.get("stream").is_none());
725 }
726
727 #[test]
728 fn chat_extra_body_is_merged_last() {
729 let body = client()
730 .chat("m", [ChatMessage::user("hi")])
731 .temperature(0.2)
732 .extra_body(json!({"temperature": 0.9, "future_field": 1}))
733 .body(false)
734 .unwrap();
735 assert_eq!(body["temperature"], json!(0.9));
736 assert_eq!(body["future_field"], json!(1));
737 }
738
739 #[test]
740 fn best_of_is_dropped_when_streaming() {
741 let request = client().chat("m", [ChatMessage::user("hi")]).best_of(4);
742 assert_eq!(request.body(false).unwrap()["best_of"], json!(4));
743 let streamed = request.body(true).unwrap();
744 assert!(streamed.get("best_of").is_none());
745 assert_eq!(streamed["stream"], json!(true));
746 }
747
748 #[test]
749 fn responses_body_takes_text_or_messages() {
750 let text = client()
751 .responses("m", ResponseInput::from("Summarize this"))
752 .max_output_tokens(64)
753 .body()
754 .unwrap();
755 assert_eq!(
756 text,
757 json!({"model": "m", "input": "Summarize this", "max_output_tokens": 64})
758 );
759
760 let messages = client()
761 .responses("m", vec![ResponseInputMessage::user("hi")].into())
762 .body()
763 .unwrap();
764 assert_eq!(
765 messages["input"],
766 json!([{"role": "user", "content": "hi"}])
767 );
768 }
769
770 #[tokio::test]
771 async fn estimate_rejects_a_relative_endpoint_or_a_non_object_request() {
772 let client = client();
773 assert!(
774 client
775 .estimate("v1/chat/completions", json!({}))
776 .await
777 .is_err()
778 );
779 assert!(
780 client
781 .estimate("/v1/chat/completions", json!([]))
782 .await
783 .is_err()
784 );
785 }
786}