1use crate::config::VlmConfig;
7use anyhow::{Context, Result, bail};
8use async_openai::types::chat::{
9 ChatCompletionMessageToolCall, ChatCompletionMessageToolCalls,
10 ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestMessage,
11 ChatCompletionRequestMessageContentPartImage, ChatCompletionRequestMessageContentPartText,
12 ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestToolMessageArgs,
13 ChatCompletionRequestUserMessageArgs, ChatCompletionRequestUserMessageContent,
14 ChatCompletionRequestUserMessageContentPart, ChatCompletionStreamOptions, ChatCompletionTool,
15 ChatCompletionTools, CreateChatCompletionRequestArgs, FunctionCall, FunctionObject,
16 FunctionObjectArgs, ImageDetail, ImageUrl, ResponseFormat,
17};
18use futures_util::stream::{Stream, StreamExt};
19use serde::{Deserialize, Serialize};
20use serde_json::Value;
21use std::collections::BTreeMap;
22use std::pin::Pin;
23use std::time::Duration;
24
25const MAX_OPEN_RETRIES: usize = 3;
26
27fn rejects_optional_request_fields(status: reqwest::StatusCode, body: &str) -> bool {
31 if status != reqwest::StatusCode::BAD_REQUEST
32 && status != reqwest::StatusCode::UNPROCESSABLE_ENTITY
33 {
34 return false;
35 }
36 let body = body.to_ascii_lowercase();
37 ["stream_options", "include_usage", "prompt_cache_key"]
38 .iter()
39 .any(|field| body.contains(field))
40}
41
42fn open_retry_delay(
43 status: reqwest::StatusCode,
44 retry_after: Option<&str>,
45 retry_index: usize,
46) -> Option<Duration> {
47 if retry_index >= MAX_OPEN_RETRIES
48 || !(status == reqwest::StatusCode::TOO_MANY_REQUESTS || status.is_server_error())
49 {
50 return None;
51 }
52 let server_seconds = retry_after
53 .and_then(|value| value.trim().parse::<u64>().ok())
54 .map(|seconds| seconds.clamp(1, 10));
55 let seconds = server_seconds.unwrap_or_else(|| 1_u64 << retry_index.min(3));
56 Some(Duration::from_secs(seconds))
57}
58
59#[derive(Serialize, Deserialize, Clone)]
86pub struct Message {
87 pub role: String,
90
91 #[serde(skip_serializing_if = "Option::is_none")]
94 pub name: Option<String>,
95
96 #[serde(skip_serializing_if = "Option::is_none")]
99 pub content: Option<String>,
100
101 #[serde(skip_serializing_if = "Option::is_none")]
105 pub tool_calls: Option<Vec<ToolCall>>,
106
107 #[serde(skip_serializing_if = "Option::is_none")]
110 pub tool_call_id: Option<String>,
111
112 #[serde(skip_serializing_if = "Option::is_none")]
116 pub image_base64: Option<String>,
117}
118
119#[derive(Serialize, Deserialize, Clone)]
120pub struct ToolCall {
121 pub id: String,
122 #[serde(rename = "type")]
123 pub kind: String,
124 pub function: FnCall,
125}
126
127#[derive(Serialize, Deserialize, Clone)]
128pub struct FnCall {
129 pub name: String,
130 pub arguments: String,
131}
132
133#[derive(Serialize, Clone)]
134pub struct ToolDef {
135 #[serde(rename = "type")]
136 kind: String,
137 function: FnDef,
138}
139
140#[derive(Serialize, Clone)]
141struct FnDef {
142 name: String,
143 description: String,
144 parameters: Value,
145}
146
147impl Message {
148 pub fn system(content: &str) -> Self {
149 Self {
150 role: "system".into(),
151 name: None,
152 content: Some(content.into()),
153 tool_calls: None,
154 tool_call_id: None,
155 image_base64: None,
156 }
157 }
158 pub fn user(content: &str) -> Self {
159 Self {
160 role: "user".into(),
161 name: None,
162 content: Some(content.into()),
163 tool_calls: None,
164 tool_call_id: None,
165 image_base64: None,
166 }
167 }
168 pub fn user_with_image(content: &str, image_base64: String) -> Self {
169 Self {
170 role: "user".into(),
171 name: None,
172 content: Some(content.into()),
173 tool_calls: None,
174 tool_call_id: None,
175 image_base64: Some(image_base64),
176 }
177 }
178 pub fn assistant(content: &str) -> Self {
179 Self {
180 role: "assistant".into(),
181 name: None,
182 content: Some(content.into()),
183 tool_calls: None,
184 tool_call_id: None,
185 image_base64: None,
186 }
187 }
188 pub fn tool_result(id: &str, content: &str) -> Self {
189 Self {
190 role: "tool".into(),
191 name: None,
192 content: Some(content.into()),
193 tool_calls: None,
194 tool_call_id: Some(id.into()),
195 image_base64: None,
196 }
197 }
198}
199
200pub enum VlmStreamItem {
203 TextDelta(String),
204 ToolCall(ToolCall),
205 Usage(VlmUsage),
208 Finish,
212}
213
214#[derive(Clone, Debug, PartialEq, Eq)]
215pub struct VlmUsage {
216 pub prompt_tokens: u64,
217 pub completion_tokens: u64,
218 pub cached_tokens: Option<u64>,
219}
220
221#[derive(Clone)]
225pub struct VlmClient {
226 inner: reqwest::Client,
227 api_base: String,
228 api_key: String,
229 model: String,
230}
231
232impl VlmClient {
233 pub fn new(cfg: &VlmConfig) -> Self {
234 Self {
235 inner: reqwest::Client::new(),
236 api_base: cfg.upstream.trim_end_matches('/').to_string(),
237 api_key: cfg.api_key.clone(),
238 model: cfg.model.clone(),
239 }
240 }
241
242 pub async fn chat_stream(
248 &self,
249 messages: &[Message],
250 tools: &[ToolDef],
251 prompt_cache_key: Option<&str>,
252 ) -> Result<Pin<Box<dyn Stream<Item = Result<VlmStreamItem>> + Send>>> {
253 let oai_messages = build_openai_messages(messages)?;
254 let oai_tools = build_openai_tools(tools)?;
255
256 let mut req_builder = CreateChatCompletionRequestArgs::default();
257 req_builder
258 .model(&self.model)
259 .messages(oai_messages)
260 .stream(true)
261 .stream_options(ChatCompletionStreamOptions {
262 include_usage: Some(true),
263 include_obfuscation: None,
264 })
265 .response_format(ResponseFormat::JsonObject);
266 if !oai_tools.is_empty() {
267 req_builder.tools(oai_tools);
268 }
269 if let Some(prompt_cache_key) = prompt_cache_key {
270 req_builder.prompt_cache_key(prompt_cache_key);
271 }
272 let request = req_builder
273 .build()
274 .context("build chat completion request")?;
275 let mut request_body = serde_json::to_value(request)
276 .context("serialize chat completion request for transport")?;
277
278 let url = format!("{}/chat/completions", self.api_base);
279 let mut retry_index = 0;
280 let mut compatibility_fallback_attempted = false;
281 let response = loop {
282 let response = self
283 .inner
284 .post(&url)
285 .bearer_auth(&self.api_key)
286 .header(reqwest::header::ACCEPT, "text/event-stream")
287 .header(reqwest::header::CONTENT_TYPE, "application/json")
288 .json(&request_body)
289 .send()
290 .await
291 .context("open VLM chat stream")?;
292 let status = response.status();
293 if status.is_success() {
294 break response;
295 }
296 let retry_after = response
297 .headers()
298 .get(reqwest::header::RETRY_AFTER)
299 .and_then(|value| value.to_str().ok())
300 .map(str::to_string);
301 let text = response.text().await.unwrap_or_default();
302 if rejects_optional_request_fields(status, &text) && !compatibility_fallback_attempted {
303 let removed = request_body.as_object_mut().is_some_and(|body| {
304 let stream_options = body.remove("stream_options").is_some();
305 let prompt_cache_key = body.remove("prompt_cache_key").is_some();
306 stream_options || prompt_cache_key
307 });
308 if !removed {
309 bail!("open VLM chat stream: HTTP {status}: {text}");
310 }
311 robonix_scribe::warn!(
312 "[pilot/vlm] upstream rejected optional cache/usage fields with HTTP {status}; retrying without them"
313 );
314 compatibility_fallback_attempted = true;
315 continue;
316 }
317 if let Some(delay) = open_retry_delay(status, retry_after.as_deref(), retry_index) {
318 robonix_scribe::warn!(
319 "[pilot/vlm] open stream HTTP {status}; retry {}/{} in {:.1}s",
320 retry_index + 1,
321 MAX_OPEN_RETRIES,
322 delay.as_secs_f64()
323 );
324 tokio::time::sleep(delay).await;
325 retry_index += 1;
326 continue;
327 }
328 bail!("open VLM chat stream: HTTP {status}: {text}");
329 };
330 let mut upstream = response.bytes_stream();
331
332 let (tx, rx) = tokio::sync::mpsc::channel::<Result<VlmStreamItem>>(64);
337 tokio::spawn(async move {
338 let mut tc_acc: BTreeMap<u32, AccumulatedToolCall> = BTreeMap::new();
339 let mut finish = "stop".to_string();
340 let mut buf = String::new();
341 let mut done = false;
342 while !done {
343 let chunk = tokio::select! {
344 _ = tx.closed() => return,
345 chunk = upstream.next() => chunk,
346 };
347 let Some(chunk) = chunk else {
348 break;
349 };
350 match chunk {
351 Ok(bytes) => {
352 buf.push_str(&String::from_utf8_lossy(&bytes));
353 while let Some(pos) = buf.find('\n') {
354 let line: String = buf.drain(..=pos).collect();
355 match process_stream_line(
356 line.trim_end(),
357 &mut tc_acc,
358 &mut finish,
359 &tx,
360 )
361 .await
362 {
363 Ok(true) => {
364 done = true;
365 break;
366 }
367 Ok(false) => {}
368 Err(e) => {
369 let _ = tx.send(Err(e)).await;
370 return;
371 }
372 }
373 }
374 }
375 Err(e) => {
376 let _ = tx
377 .send(Err(anyhow::anyhow!("VLM stream chunk error: {e}")))
378 .await;
379 return;
380 }
381 }
382 }
383 if !buf.trim().is_empty()
384 && let Err(e) =
385 process_stream_line(buf.trim_end(), &mut tc_acc, &mut finish, &tx).await
386 {
387 let _ = tx.send(Err(e)).await;
388 return;
389 }
390
391 for (_, tc) in tc_acc {
392 if tc.id.is_empty() && tc.name.is_empty() {
393 continue;
394 }
395 let item = VlmStreamItem::ToolCall(ToolCall {
396 id: tc.id,
397 kind: "function".to_string(),
398 function: FnCall {
399 name: tc.name,
400 arguments: tc.arguments,
401 },
402 });
403 if tx.send(Ok(item)).await.is_err() {
404 return;
405 }
406 }
407 let _ = finish; let _ = tx.send(Ok(VlmStreamItem::Finish)).await;
409 });
410
411 Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)))
412 }
413}
414
415#[cfg(test)]
416mod tests {
417 use super::{
418 AccumulatedToolCall, MAX_OPEN_RETRIES, VlmStreamItem, VlmUsage, open_retry_delay,
419 parse_usage, process_stream_line, rejects_optional_request_fields,
420 };
421 use serde_json::json;
422 use std::collections::BTreeMap;
423 use std::time::Duration;
424
425 #[test]
426 fn transient_open_errors_use_bounded_backoff() {
427 assert_eq!(
428 open_retry_delay(reqwest::StatusCode::TOO_MANY_REQUESTS, None, 0),
429 Some(Duration::from_secs(1))
430 );
431 assert_eq!(
432 open_retry_delay(reqwest::StatusCode::SERVICE_UNAVAILABLE, Some("7"), 1),
433 Some(Duration::from_secs(7))
434 );
435 assert_eq!(
436 open_retry_delay(reqwest::StatusCode::BAD_REQUEST, None, 0),
437 None
438 );
439 assert_eq!(
440 open_retry_delay(
441 reqwest::StatusCode::TOO_MANY_REQUESTS,
442 None,
443 MAX_OPEN_RETRIES
444 ),
445 None
446 );
447 }
448
449 #[test]
450 fn usage_includes_provider_prompt_cache_tokens() {
451 let usage = parse_usage(&json!({
452 "choices": [],
453 "usage": {
454 "prompt_tokens": 1200,
455 "completion_tokens": 80,
456 "total_tokens": 1280,
457 "prompt_tokens_details": {"cached_tokens": 900}
458 }
459 }));
460 assert_eq!(
461 usage,
462 Some(VlmUsage {
463 prompt_tokens: 1200,
464 completion_tokens: 80,
465 cached_tokens: Some(900),
466 })
467 );
468 }
469
470 #[test]
471 fn optional_field_fallback_does_not_mask_unrelated_client_errors() {
472 assert!(rejects_optional_request_fields(
473 reqwest::StatusCode::BAD_REQUEST,
474 "unknown field prompt_cache_key"
475 ));
476 assert!(rejects_optional_request_fields(
477 reqwest::StatusCode::UNPROCESSABLE_ENTITY,
478 "stream_options is not permitted"
479 ));
480 assert!(!rejects_optional_request_fields(
481 reqwest::StatusCode::BAD_REQUEST,
482 "invalid model name"
483 ));
484 assert!(!rejects_optional_request_fields(
485 reqwest::StatusCode::UNAUTHORIZED,
486 "prompt_cache_key"
487 ));
488 }
489
490 #[tokio::test]
491 async fn choice_less_usage_chunk_reaches_the_consumer() {
492 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
493 let mut calls = BTreeMap::<u32, AccumulatedToolCall>::new();
494 let mut finish = String::new();
495 process_stream_line(
496 r#"data: {"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":80,"prompt_tokens_details":{"cached_tokens":900}}}"#,
497 &mut calls,
498 &mut finish,
499 &tx,
500 )
501 .await
502 .unwrap();
503 assert!(matches!(
504 rx.recv().await,
505 Some(Ok(VlmStreamItem::Usage(VlmUsage {
506 prompt_tokens: 1200,
507 completion_tokens: 80,
508 cached_tokens: Some(900),
509 })))
510 ));
511 }
512}
513
514#[derive(Default)]
515struct AccumulatedToolCall {
516 id: String,
517 name: String,
518 arguments: String,
519}
520
521async fn process_stream_line(
522 line: &str,
523 tc_acc: &mut BTreeMap<u32, AccumulatedToolCall>,
524 finish: &mut String,
525 tx: &tokio::sync::mpsc::Sender<Result<VlmStreamItem>>,
526) -> Result<bool> {
527 let line = line.trim();
528 if line.is_empty() || line.starts_with(':') {
529 return Ok(false);
530 }
531 let Some(data) = line.strip_prefix("data:") else {
532 return Ok(false);
533 };
534 let data = data.trim();
535 if data == "[DONE]" {
536 return Ok(true);
537 }
538
539 let v: Value = serde_json::from_str(data)
540 .with_context(|| format!("deserialize VLM stream chunk: {data}"))?;
541 if let Some(usage) = parse_usage(&v)
542 && tx.send(Ok(VlmStreamItem::Usage(usage))).await.is_err()
543 {
544 return Ok(true);
545 }
546 let Some(choice) = v
547 .get("choices")
548 .and_then(Value::as_array)
549 .and_then(|choices| choices.first())
550 else {
551 return Ok(false);
552 };
553
554 if let Some(content) = choice
555 .get("delta")
556 .and_then(|delta| delta.get("content"))
557 .and_then(Value::as_str)
558 && !content.is_empty()
559 && tx
560 .send(Ok(VlmStreamItem::TextDelta(content.to_string())))
561 .await
562 .is_err()
563 {
564 return Ok(true);
565 }
566 if let Some(tc_chunks) = choice
567 .get("delta")
568 .and_then(|delta| delta.get("tool_calls"))
569 .and_then(Value::as_array)
570 {
571 for tc in tc_chunks {
572 let index = tc.get("index").and_then(Value::as_u64).unwrap_or(0) as u32;
573 let entry = tc_acc.entry(index).or_default();
574 if let Some(id) = tc.get("id").and_then(Value::as_str) {
575 entry.id = id.to_string();
576 }
577 if let Some(func) = tc.get("function") {
578 if let Some(name) = func.get("name").and_then(Value::as_str) {
579 entry.name.push_str(name);
580 }
581 if let Some(args) = func.get("arguments").and_then(Value::as_str) {
582 entry.arguments.push_str(args);
583 }
584 }
585 }
586 }
587 if let Some(fr) = choice.get("finish_reason").and_then(Value::as_str) {
588 *finish = fr.to_string();
589 }
590 Ok(false)
591}
592
593fn parse_usage(value: &Value) -> Option<VlmUsage> {
596 let usage = value.get("usage")?.as_object()?;
597 let prompt_tokens = usage.get("prompt_tokens")?.as_u64()?;
598 let completion_tokens = usage.get("completion_tokens")?.as_u64()?;
599 let cached_tokens = usage
600 .get("prompt_tokens_details")
601 .and_then(Value::as_object)
602 .and_then(|details| details.get("cached_tokens"))
603 .and_then(Value::as_u64);
604 Some(VlmUsage {
605 prompt_tokens,
606 completion_tokens,
607 cached_tokens,
608 })
609}
610
611fn build_openai_messages(messages: &[Message]) -> Result<Vec<ChatCompletionRequestMessage>> {
612 let mut out = Vec::with_capacity(messages.len());
613 for m in messages {
614 let msg = match m.role.as_str() {
615 "system" => ChatCompletionRequestSystemMessageArgs::default()
616 .content(m.content.clone().unwrap_or_default())
617 .build()?
618 .into(),
619 "user" => {
620 if let Some(image) = &m.image_base64 {
621 let text = m.content.clone().unwrap_or_default();
622 let mut parts: Vec<ChatCompletionRequestUserMessageContentPart> = Vec::new();
623 if !text.is_empty() {
624 parts.push(ChatCompletionRequestUserMessageContentPart::Text(
625 ChatCompletionRequestMessageContentPartText { text },
626 ));
627 }
628 let url = format!("data:image/jpeg;base64,{image}");
629 parts.push(ChatCompletionRequestUserMessageContentPart::ImageUrl(
630 ChatCompletionRequestMessageContentPartImage {
631 image_url: ImageUrl {
632 url,
633 detail: Some(ImageDetail::Auto),
634 },
635 },
636 ));
637 ChatCompletionRequestUserMessageArgs::default()
638 .content(ChatCompletionRequestUserMessageContent::Array(parts))
639 .build()?
640 .into()
641 } else {
642 ChatCompletionRequestUserMessageArgs::default()
643 .content(m.content.clone().unwrap_or_default())
644 .build()?
645 .into()
646 }
647 }
648 "assistant" => {
649 let mut b = ChatCompletionRequestAssistantMessageArgs::default();
650 if let Some(c) = &m.content
651 && !c.is_empty()
652 {
653 b.content(c.clone());
654 }
655 if let Some(tcs) = &m.tool_calls {
656 let oai_tcs: Vec<ChatCompletionMessageToolCalls> = tcs
657 .iter()
658 .map(|tc| {
659 ChatCompletionMessageToolCalls::Function(
660 ChatCompletionMessageToolCall {
661 id: tc.id.clone(),
662 function: FunctionCall {
663 name: tc.function.name.clone(),
664 arguments: tc.function.arguments.clone(),
665 },
666 },
667 )
668 })
669 .collect();
670 b.tool_calls(oai_tcs);
671 }
672 b.build()?.into()
673 }
674 "tool" => {
675 let id = m.tool_call_id.clone().unwrap_or_default();
676 ChatCompletionRequestToolMessageArgs::default()
677 .tool_call_id(id)
678 .content(m.content.clone().unwrap_or_default())
679 .build()?
680 .into()
681 }
682 other => anyhow::bail!("unknown message role '{other}'"),
683 };
684 out.push(msg);
685 }
686 Ok(out)
687}
688
689fn build_openai_tools(tools: &[ToolDef]) -> Result<Vec<ChatCompletionTools>> {
690 tools
691 .iter()
692 .map(|t| -> Result<ChatCompletionTools> {
693 let func: FunctionObject = FunctionObjectArgs::default()
694 .name(&t.function.name)
695 .description(&t.function.description)
696 .parameters(t.function.parameters.clone())
697 .build()?;
698 Ok(ChatCompletionTools::Function(ChatCompletionTool {
699 function: func,
700 }))
701 })
702 .collect()
703}