1use std::sync::Arc;
4
5use tap::Pipe;
6
7use crate::error::{GatewayError, LegFailure};
8use crate::providers::genai_provider::Provider;
9use crate::providers::Catalog;
10use crate::resilience::{run_with_classifier, ResilienceError};
11use crate::routing::request::ChatRequest;
12use crate::routing::stream::{Accumulator, FinishReason, StreamItem, ToolCallOut};
13use crate::routing::table::ChainLeg;
14use futures::stream::{BoxStream, Stream, StreamExt};
15
16#[derive(Debug, Clone)]
18pub struct Completion {
19 pub provider: String,
20 pub model: String,
21 pub content: String,
22 pub tool_calls: Vec<ToolCallOut>,
23 pub finish_reason: FinishReason,
24 pub input_tokens: u64,
25 pub output_tokens: u64,
26}
27
28fn to_genai_request(req: &ChatRequest) -> genai::chat::ChatRequest {
32 use genai::chat::{ChatMessage, Tool, ToolCall, ToolResponse};
33
34 let mut chat = genai::chat::ChatRequest::default();
35
36 for m in &req.messages {
37 match m.role.as_str() {
38 "assistant" if m.tool_calls.is_some() => {
39 let calls: Vec<ToolCall> = m
40 .tool_calls
41 .as_ref()
42 .unwrap()
43 .iter()
44 .filter_map(openai_tool_call_to_genai)
45 .collect();
46 chat = chat.append_message(ChatMessage::from(calls));
47 }
48 "tool" => {
49 let call_id = m.tool_call_id.clone().unwrap_or_default();
50 let content = m
51 .content
52 .as_str()
53 .map(str::to_string)
54 .unwrap_or_else(|| m.content.to_string());
55 chat = chat.append_message(ChatMessage::from(ToolResponse { call_id, content }));
56 }
57 role => {
58 let msg = match role {
59 "system" => ChatMessage::system(genai::chat::MessageContent::from_parts(
60 crate::routing::content_parts::content_to_genai_parts(&m.content),
61 )),
62 "assistant" => ChatMessage::assistant(genai::chat::MessageContent::from_parts(
63 crate::routing::content_parts::content_to_genai_parts(&m.content),
64 )),
65 _ => ChatMessage::user(genai::chat::MessageContent::from_parts(
66 crate::routing::content_parts::content_to_genai_parts(&m.content),
67 )),
68 };
69 chat = chat.append_message(msg);
70 }
71 }
72 }
73
74 if let Some(tools) = &req.tools {
75 let mapped: Vec<Tool> = tools.iter().filter_map(openai_tool_to_genai).collect();
76 if !mapped.is_empty() {
77 chat = chat.with_tools(mapped);
78 }
79 }
80
81 chat
82}
83
84fn openai_tool_to_genai(v: &serde_json::Value) -> Option<genai::chat::Tool> {
86 let f = v.get("function")?;
87 let name = f.get("name")?.as_str()?.to_string();
88 let mut tool = genai::chat::Tool::new(name);
89 if let Some(desc) = f.get("description").and_then(|d| d.as_str()) {
90 tool = tool.with_description(desc);
91 }
92 if let Some(params) = f.get("parameters") {
93 tool = tool.with_schema(params.clone());
94 }
95 Some(tool)
96}
97
98fn openai_tool_call_to_genai(v: &serde_json::Value) -> Option<genai::chat::ToolCall> {
101 let call_id = v.get("id")?.as_str()?.to_string();
102 let f = v.get("function")?;
103 let fn_name = f.get("name")?.as_str()?.to_string();
104 let raw_args = f.get("arguments").and_then(|a| a.as_str()).unwrap_or("{}");
105 let fn_arguments =
106 serde_json::from_str(raw_args).unwrap_or_else(|_| serde_json::json!(raw_args));
107 Some(genai::chat::ToolCall {
108 call_id,
109 fn_name,
110 fn_arguments,
111 thought_signatures: None,
112 })
113}
114
115fn to_genai_options(req: &ChatRequest) -> genai::chat::ChatOptions {
116 genai::chat::ChatOptions::default()
117 .pipe(|o| match req.temperature {
118 Some(t) => o.with_temperature(t as f64),
119 None => o,
120 })
121 .pipe(|o| match &req.response_format {
122 Some(rf) => match rf.kind.as_str() {
123 "json_object" => o.with_response_format(genai::chat::ChatResponseFormat::JsonMode),
124 "json_schema" => {
125 if let Some(spec) = rf.json_schema.clone() {
126 o.with_response_format(genai::chat::ChatResponseFormat::JsonSpec(
127 genai::chat::JsonSpec::new("synapse", spec),
128 ))
129 } else {
130 o
131 }
132 }
133 _ => o,
134 },
135 None => o,
136 })
137}
138
139fn is_genai_retryable(e: &genai::Error) -> bool {
157 fn webc_retryable(we: &genai::webc::Error) -> bool {
159 match we {
160 genai::webc::Error::ResponseFailedStatus { status, .. } => status.is_server_error(),
161 genai::webc::Error::Reqwest(re) => re.is_timeout() || re.is_connect(),
162 _ => false,
163 }
164 }
165
166 match e {
167 genai::Error::WebModelCall { webc_error, .. } => webc_retryable(webc_error),
168 genai::Error::WebAdapterCall { webc_error, .. } => webc_retryable(webc_error),
169 genai::Error::HttpError { status, .. } => status.is_server_error(),
170 _ => false,
171 }
172}
173
174#[derive(Debug)]
176pub enum LegError {
177 Start(String),
179 MidStream(String),
181}
182
183impl std::fmt::Display for LegError {
184 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
185 match self {
186 LegError::Start(s) => write!(f, "start: {s}"),
187 LegError::MidStream(s) => write!(f, "mid-stream: {s}"),
188 }
189 }
190}
191
192fn map_stop_reason(sr: Option<&genai::chat::StopReason>) -> FinishReason {
194 match sr {
195 Some(genai::chat::StopReason::ToolCall(_)) => FinishReason::ToolCalls,
196 Some(genai::chat::StopReason::MaxTokens(_)) => FinishReason::Length,
197 _ => FinishReason::Stop,
198 }
199}
200
201#[derive(Default)]
203struct ToolIndexer {
204 ids: Vec<String>,
205}
206impl ToolIndexer {
207 fn index_of(&mut self, call_id: &str) -> (u32, bool) {
209 if let Some(pos) = self.ids.iter().position(|c| c == call_id) {
210 (pos as u32, false)
211 } else {
212 self.ids.push(call_id.to_string());
213 ((self.ids.len() - 1) as u32, true)
214 }
215 }
216}
217
218pub async fn stream_one_leg_standard(
221 provider: &Arc<Provider>,
222 model: &str,
223 req: &ChatRequest,
224) -> Result<impl Stream<Item = Result<StreamItem, LegError>>, LegError> {
225 let chat_req = to_genai_request(req);
226 let opts = to_genai_options(req)
227 .with_capture_usage(true)
228 .with_capture_tool_calls(true);
229
230 let resp = provider
231 .client
232 .exec_chat_stream(model.to_string(), chat_req, Some(&opts))
233 .await
234 .map_err(|e| LegError::Start(e.to_string()))?;
235
236 struct St<S> {
240 inner: S,
241 indexer: ToolIndexer,
242 input_tokens: u64,
243 output_tokens: u64,
244 }
245 let state = St {
246 inner: Box::pin(resp.stream),
247 indexer: ToolIndexer::default(),
248 input_tokens: 0,
249 output_tokens: 0,
250 };
251
252 let normalized = futures::stream::unfold(state, |mut st| async move {
253 loop {
254 match st.inner.next().await {
255 None => return None,
256 Some(Err(e)) => return Some((Err(LegError::MidStream(e.to_string())), st)),
257 Some(Ok(ev)) => match ev {
258 genai::chat::ChatStreamEvent::Chunk(c) => {
259 if !c.content.is_empty() {
260 return Some((Ok(StreamItem::Delta(c.content)), st));
261 }
262 }
263 genai::chat::ChatStreamEvent::ToolCallChunk(tc) => {
264 let call = tc.tool_call;
265 let (index, first) = st.indexer.index_of(&call.call_id);
266 let args = match &call.fn_arguments {
267 serde_json::Value::String(s) => s.clone(),
268 other => other.to_string(),
269 };
270 return Some((
271 Ok(StreamItem::ToolCallDelta {
272 index,
273 id: if first { Some(call.call_id) } else { None },
274 name: if first { Some(call.fn_name) } else { None },
275 args_fragment: args,
276 }),
277 st,
278 ));
279 }
280 genai::chat::ChatStreamEvent::End(end) => {
281 if let Some(u) = &end.captured_usage {
282 st.input_tokens = u.prompt_tokens.unwrap_or(0).max(0) as u64;
283 st.output_tokens = u.completion_tokens.unwrap_or(0).max(0) as u64;
284 }
285 let done = StreamItem::Done {
286 input_tokens: st.input_tokens,
287 output_tokens: st.output_tokens,
288 finish_reason: map_stop_reason(end.captured_stop_reason.as_ref()),
289 };
290 return Some((Ok(done), st));
291 }
292 _ => {} },
294 }
295 }
296 });
297
298 Ok(normalized)
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304
305 fn req(body: serde_json::Value) -> ChatRequest {
306 serde_json::from_value(body).unwrap()
307 }
308
309 #[test]
310 fn maps_tools_into_genai_request() {
311 let r = req(serde_json::json!({
312 "model": "m",
313 "messages": [{"role": "user", "content": "hi"}],
314 "tools": [{"type": "function", "function": {"name": "get_weather",
315 "description": "Lookup", "parameters": {"type": "object"}}}]
316 }));
317 let g = to_genai_request(&r);
318 let tools = g.tools.expect("tools mapped");
319 assert_eq!(tools.len(), 1);
320 assert_eq!(tools[0].name.to_string(), "get_weather");
321 assert_eq!(tools[0].description.as_deref(), Some("Lookup"));
322 }
323
324 #[test]
325 fn maps_assistant_tool_calls_and_tool_results() {
326 let r = req(serde_json::json!({
327 "model": "m",
328 "messages": [
329 {"role": "user", "content": "weather?"},
330 {"role": "assistant", "content": null, "tool_calls": [
331 {"id": "call_0", "type": "function", "function": {"name": "f", "arguments": "{\"c\":\"SF\"}"}}]},
332 {"role": "tool", "tool_call_id": "call_0", "content": "21C"}
333 ]
334 }));
335 let g = to_genai_request(&r);
336 assert_eq!(g.messages.len(), 3);
337 }
338
339 use crate::providers::genai_provider::{build_openai_compat_provider, OpenAiCompatConfig};
340 use crate::routing::stream::{FinishReason, StreamItem};
341 use futures::StreamExt;
342 use std::time::Duration;
343
344 #[tokio::test]
345 async fn standard_lane_streams_text_then_done() {
346 use wiremock::matchers::{method, path};
347 use wiremock::{Mock, MockServer, ResponseTemplate};
348
349 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"He\"}}]}\n\n\
351 data: {\"choices\":[{\"delta\":{\"content\":\"llo\"}}]}\n\n\
352 data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5}}\n\n\
353 data: [DONE]\n\n";
354 let mock = MockServer::start().await;
355 Mock::given(method("POST"))
356 .and(path("/v1/chat/completions"))
357 .respond_with(
358 ResponseTemplate::new(200)
359 .insert_header("content-type", "text/event-stream")
360 .set_body_string(sse),
361 )
362 .mount(&mock)
363 .await;
364
365 let provider = Arc::new(
366 build_openai_compat_provider(
367 "oai",
368 OpenAiCompatConfig {
369 base_url: format!("{}/v1", mock.uri()),
370 api_key: "k".into(),
371 request_timeout: Duration::from_secs(5),
372 endpoint_override: None,
373 },
374 )
375 .unwrap(),
376 );
377
378 let req = req(
379 serde_json::json!({"model":"m","messages":[{"role":"user","content":"hi"}],"stream":true}),
380 );
381 let mut stream = std::pin::pin!(stream_one_leg_standard(&provider, "m", &req)
382 .await
383 .expect("stream starts"));
384 let mut items = Vec::new();
385 while let Some(it) = stream.next().await {
386 items.push(it.expect("no mid-stream error"));
387 }
388 assert!(items
389 .iter()
390 .any(|i| matches!(i, StreamItem::Delta(t) if t == "He")));
391 assert!(matches!(
392 items.last().unwrap(),
393 StreamItem::Done {
394 input_tokens: 3,
395 output_tokens: 2,
396 finish_reason: FinishReason::Stop
397 }
398 ));
399 }
400
401 #[tokio::test]
402 async fn execute_buffered_falls_back_on_midstream_failure() {
403 use crate::providers::Catalog;
404 use crate::routing::table::ChainLeg;
405 use wiremock::matchers::{method, path};
406 use wiremock::{Mock, MockServer, ResponseTemplate};
407
408 let bad = MockServer::start().await;
410 Mock::given(method("POST"))
411 .and(path("/v1/chat/completions"))
412 .respond_with(
413 ResponseTemplate::new(200)
414 .insert_header("content-type", "text/event-stream")
415 .set_body_string("data: {\"choices\":[{\"delta\":{\"content\":\"par\"}}]}\n\n"),
416 ) .mount(&bad)
418 .await;
419 let good = MockServer::start().await;
421 Mock::given(method("POST"))
422 .and(path("/v1/chat/completions"))
423 .respond_with(
424 ResponseTemplate::new(200)
425 .insert_header("content-type", "text/event-stream")
426 .set_body_string(
427 "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n\
428 data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n\
429 data: [DONE]\n\n",
430 ),
431 )
432 .mount(&good)
433 .await;
434
435 let catalog = Catalog::for_test(vec![
436 ("p1", format!("{}/v1", bad.uri())),
437 ("p2", format!("{}/v1", good.uri())),
438 ]);
439 let legs = vec![
440 ChainLeg {
441 provider: "p1".into(),
442 model: "m".into(),
443 },
444 ChainLeg {
445 provider: "p2".into(),
446 model: "m".into(),
447 },
448 ];
449 let r =
450 req(serde_json::json!({"model":"route","messages":[{"role":"user","content":"hi"}]}));
451 let c = execute_buffered(&catalog, "route", &legs, &r)
452 .await
453 .unwrap();
454 assert_eq!(c.content, "ok");
455 assert_eq!(c.provider, "p2");
456 }
457
458 #[tokio::test]
459 async fn execute_streaming_commits_first_leg_with_items() {
460 use crate::providers::Catalog;
461 use crate::routing::stream::StreamItem;
462 use crate::routing::table::ChainLeg;
463 use futures::StreamExt;
464 use wiremock::matchers::{method, path};
465 use wiremock::{Mock, MockServer, ResponseTemplate};
466
467 let bad = MockServer::start().await;
469 Mock::given(method("POST"))
470 .and(path("/v1/chat/completions"))
471 .respond_with(ResponseTemplate::new(500))
472 .mount(&bad)
473 .await;
474 let good = MockServer::start().await;
475 Mock::given(method("POST")).and(path("/v1/chat/completions"))
476 .respond_with(ResponseTemplate::new(200)
477 .insert_header("content-type", "text/event-stream")
478 .set_body_string("data: {\"choices\":[{\"delta\":{\"content\":\"go\"}}]}\n\n\
479 data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n\
480 data: [DONE]\n\n"))
481 .mount(&good).await;
482
483 let catalog = Catalog::for_test(vec![
484 ("p1", format!("{}/v1", bad.uri())),
485 ("p2", format!("{}/v1", good.uri())),
486 ]);
487 let legs = vec![
488 ChainLeg {
489 provider: "p1".into(),
490 model: "m".into(),
491 },
492 ChainLeg {
493 provider: "p2".into(),
494 model: "m".into(),
495 },
496 ];
497 let r = req(
498 serde_json::json!({"model":"route","messages":[{"role":"user","content":"hi"}],"stream":true}),
499 );
500 let committed = execute_streaming(&catalog, "route", &legs, &r)
501 .await
502 .unwrap();
503 assert_eq!(committed.provider, "p2");
504 let mut stream = committed.stream;
505 let mut items = Vec::new();
506 while let Some(i) = stream.next().await {
507 items.push(i.unwrap());
508 }
509 assert!(items
510 .iter()
511 .any(|i| matches!(i, StreamItem::Delta(t) if t == "go")));
512 }
513
514 #[tokio::test]
515 async fn first_chunk_timeout_falls_back() {
516 use crate::providers::Catalog;
517 use crate::routing::table::ChainLeg;
518 use std::time::Duration;
519 use wiremock::matchers::{method, path};
520 use wiremock::{Mock, MockServer, ResponseTemplate};
521
522 let slow = MockServer::start().await;
533 Mock::given(method("POST")).and(path("/v1/chat/completions"))
534 .respond_with(ResponseTemplate::new(200)
535 .insert_header("content-type", "text/event-stream")
536 .set_delay(Duration::from_millis(400))
537 .set_body_string("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{}}\n\ndata: [DONE]\n\n"))
538 .mount(&slow).await;
539 let good = MockServer::start().await;
540 Mock::given(method("POST")).and(path("/v1/chat/completions"))
541 .respond_with(ResponseTemplate::new(200)
542 .insert_header("content-type", "text/event-stream")
543 .set_body_string("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n\
544 data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n\
545 data: [DONE]\n\n")).mount(&good).await;
546
547 let catalog = Catalog::for_test(vec![
548 ("p1", format!("{}/v1", slow.uri())),
549 ("p2", format!("{}/v1", good.uri())),
550 ]);
551 let legs = vec![
552 ChainLeg {
553 provider: "p1".into(),
554 model: "m".into(),
555 },
556 ChainLeg {
557 provider: "p2".into(),
558 model: "m".into(),
559 },
560 ];
561 let r =
562 req(serde_json::json!({"model":"route","messages":[{"role":"user","content":"hi"}]}));
563 let timeouts = StreamTimeouts {
564 first_chunk: Duration::from_millis(150),
565 idle: Duration::from_secs(5),
566 };
567 let c = execute_buffered_with_timeouts(&catalog, "route", &legs, &r, timeouts)
568 .await
569 .unwrap();
570 assert_eq!(c.provider, "p2");
571 }
572}
573
574async fn run_one_leg(
575 provider: &Arc<Provider>,
576 leg: &ChainLeg,
577 req: &ChatRequest,
578) -> Result<Completion, ResilienceError<genai::Error>> {
579 let client = provider.client.clone();
580 let model = leg.model.clone();
581 let chat_req = to_genai_request(req);
582 let opts = to_genai_options(req);
583
584 let resp = run_with_classifier(
585 move || {
586 let (client, model, chat_req, opts) = (
587 client.clone(),
588 model.clone(),
589 chat_req.clone(),
590 opts.clone(),
591 );
592 async move { client.exec_chat(model, chat_req, Some(&opts)).await }
593 },
594 provider.profile,
595 &provider.breaker,
596 provider.label,
597 is_genai_retryable,
598 )
599 .await?;
600
601 let content = resp.first_text().unwrap_or_default().to_string();
602 let usage = &resp.usage;
603 Ok(Completion {
604 provider: leg.provider.clone(),
605 model: leg.model.clone(),
606 content,
607 tool_calls: Vec::new(),
608 finish_reason: FinishReason::Stop,
609 input_tokens: usage.prompt_tokens.unwrap_or(0).max(0) as u64,
610 output_tokens: usage.completion_tokens.unwrap_or(0).max(0) as u64,
611 })
612}
613
614pub async fn execute_chain(
617 catalog: &Catalog,
618 route_name: &str,
619 legs: &[ChainLeg],
620 req: &ChatRequest,
621) -> Result<Completion, GatewayError> {
622 let mut failures: Vec<LegFailure> = Vec::new();
623 let mut all_circuit_open = true;
624 for leg in legs {
625 let provider = catalog.get(&leg.provider).ok_or_else(|| {
626 GatewayError::BadRequest(format!(
627 "route '{route_name}' references unbuilt provider '{}'",
628 leg.provider
629 ))
630 })?;
631 match run_one_leg(provider, leg, req).await {
632 Ok(c) => return Ok(c),
633 Err(ResilienceError::CircuitOpen { name }) => failures.push(LegFailure {
634 provider: leg.provider.clone(),
635 model: leg.model.clone(),
636 message: format!("circuit open: {name}"),
637 }),
638 Err(ResilienceError::Exhausted(e)) => {
639 all_circuit_open = false;
640 let retryable = is_genai_retryable(&e);
641 failures.push(LegFailure {
642 provider: leg.provider.clone(),
643 model: leg.model.clone(),
644 message: e.to_string(),
645 });
646 if !retryable {
647 break; }
649 }
650 }
651 }
652 if all_circuit_open && !failures.is_empty() {
653 return Err(GatewayError::AllCircuitsOpen(route_name.to_string()));
654 }
655 Err(GatewayError::AllLegsFailed {
656 route: route_name.to_string(),
657 failures,
658 })
659}
660
661#[derive(Debug, Clone, Copy)]
663pub struct StreamTimeouts {
664 pub first_chunk: std::time::Duration,
666 pub idle: std::time::Duration,
668}
669
670impl Default for StreamTimeouts {
671 fn default() -> Self {
672 Self {
673 first_chunk: std::time::Duration::from_secs(120),
674 idle: std::time::Duration::from_secs(60),
675 }
676 }
677}
678
679async fn buffer_one_leg_timed(
681 catalog: &Catalog,
682 leg: &ChainLeg,
683 req: &ChatRequest,
684 t: StreamTimeouts,
685) -> Result<Completion, LegError> {
686 let provider = catalog
687 .get(&leg.provider)
688 .ok_or_else(|| LegError::Start(format!("unbuilt provider '{}'", leg.provider)))?;
689 let stream = stream_one_leg_standard(provider, &leg.model, req).await?;
690 let mut stream = std::pin::pin!(stream);
691 let mut acc = Accumulator::default();
692 let mut first = true;
693 loop {
694 let budget = if first { t.first_chunk } else { t.idle };
695 match tokio::time::timeout(budget, stream.next()).await {
696 Err(_) => {
697 return Err(if first {
698 LegError::Start("first-chunk timeout".into())
699 } else {
700 LegError::MidStream("idle timeout".into())
701 })
702 }
703 Ok(None) => break,
704 Ok(Some(item)) => {
705 acc.push(item?);
706 first = false;
707 }
708 }
709 }
710 if !acc.got_done {
711 return Err(LegError::MidStream("stream ended before completion".into()));
712 }
713 Ok(Completion {
714 provider: leg.provider.clone(),
715 model: leg.model.clone(),
716 content: acc.content,
717 tool_calls: acc.tool_calls,
718 finish_reason: acc.finish_reason,
719 input_tokens: acc.input_tokens,
720 output_tokens: acc.output_tokens,
721 })
722}
723
724pub async fn execute_buffered_with_timeouts(
729 catalog: &Catalog,
730 route_name: &str,
731 legs: &[ChainLeg],
732 req: &ChatRequest,
733 t: StreamTimeouts,
734) -> Result<Completion, GatewayError> {
735 let mut failures: Vec<LegFailure> = Vec::new();
736 for leg in legs {
737 match buffer_one_leg_timed(catalog, leg, req, t).await {
738 Ok(c) => return Ok(c),
739 Err(e) => failures.push(LegFailure {
740 provider: leg.provider.clone(),
741 model: leg.model.clone(),
742 message: e.to_string(),
743 }),
744 }
745 }
746 Err(GatewayError::AllLegsFailed {
747 route: route_name.to_string(),
748 failures,
749 })
750}
751
752pub async fn execute_buffered(
762 catalog: &Catalog,
763 route_name: &str,
764 legs: &[ChainLeg],
765 req: &ChatRequest,
766) -> Result<Completion, GatewayError> {
767 execute_buffered_with_timeouts(catalog, route_name, legs, req, StreamTimeouts::default()).await
768}
769
770pub struct CommittedStream {
773 pub provider: String,
774 pub model: String,
775 pub stream: BoxStream<'static, Result<StreamItem, LegError>>,
776}
777
778impl CommittedStream {
779 pub fn single(
782 provider: String,
783 model: String,
784 stream: impl Stream<Item = Result<StreamItem, LegError>> + Send + 'static,
785 ) -> Self {
786 Self {
787 provider,
788 model,
789 stream: stream.boxed(),
790 }
791 }
792}
793
794pub async fn execute_streaming_with_timeouts(
799 catalog: &Catalog,
800 route_name: &str,
801 legs: &[ChainLeg],
802 req: &ChatRequest,
803 t: StreamTimeouts,
804) -> Result<CommittedStream, GatewayError> {
805 let mut failures: Vec<LegFailure> = Vec::new();
806
807 for leg in legs {
808 let provider = match catalog.get(&leg.provider) {
809 Some(p) => p,
810 None => {
811 failures.push(LegFailure {
812 provider: leg.provider.clone(),
813 model: leg.model.clone(),
814 message: format!("unbuilt provider '{}'", leg.provider),
815 });
816 continue;
817 }
818 };
819 let started = match stream_one_leg_standard(provider, &leg.model, req).await {
820 Ok(s) => s,
821 Err(e) => {
822 failures.push(LegFailure {
823 provider: leg.provider.clone(),
824 model: leg.model.clone(),
825 message: e.to_string(),
826 });
827 continue;
828 }
829 };
830 let mut stream = Box::pin(started);
831 match tokio::time::timeout(t.first_chunk, stream.next()).await {
832 Err(_) => {
833 failures.push(LegFailure {
834 provider: leg.provider.clone(),
835 model: leg.model.clone(),
836 message: "first-chunk timeout".into(),
837 });
838 }
839 Ok(Some(Ok(first))) => {
840 let rest = futures::stream::once(async move { Ok(first) }).chain(stream);
841 return Ok(CommittedStream {
842 provider: leg.provider.clone(),
843 model: leg.model.clone(),
844 stream: rest.boxed(),
845 });
846 }
847 Ok(Some(Err(e))) => {
848 failures.push(LegFailure {
849 provider: leg.provider.clone(),
850 model: leg.model.clone(),
851 message: e.to_string(),
852 });
853 }
854 Ok(None) => {
855 failures.push(LegFailure {
856 provider: leg.provider.clone(),
857 model: leg.model.clone(),
858 message: "empty stream".into(),
859 });
860 }
861 }
862 }
863 Err(GatewayError::AllLegsFailed {
864 route: route_name.to_string(),
865 failures,
866 })
867}
868
869pub async fn execute_streaming(
876 catalog: &Catalog,
877 route_name: &str,
878 legs: &[ChainLeg],
879 req: &ChatRequest,
880) -> Result<CommittedStream, GatewayError> {
881 execute_streaming_with_timeouts(catalog, route_name, legs, req, StreamTimeouts::default()).await
882}