From 01baad10a11c69b11d85a94596fbf6ccd42484c2 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Thu, 13 Aug 2026 22:32:08 -0700 Subject: [PATCH] feat: add specialized media model routing --- .gitignore | 1 + Cargo.lock | 2 + Cargo.toml | 1 + crates/libsy-llm-client/Cargo.toml | 2 + crates/libsy-llm-client/src/lib.rs | 2 + crates/libsy-llm-client/src/media.rs | 306 +++++++++++++ crates/libsy/README.md | 1 + crates/libsy/src/algorithms.rs | 1 + crates/libsy/src/algorithms/model_as_tool.rs | 425 +++++++++++++++++++ crates/libsy/src/lib.rs | 1 + crates/switchyard-py/src/libsy_bindings.rs | 13 +- crates/switchyard-server/CONFIGURATION.md | 30 ++ crates/switchyard-server/src/config.rs | 315 ++++++++++++-- examples/model-as-tool-media.toml | 41 ++ switchyard/libsy/algorithms.py | 3 +- switchyard_rust/libsy.py | 3 + tests/test_libsy_minimal_bindings.py | 51 +++ 17 files changed, 1155 insertions(+), 43 deletions(-) create mode 100644 crates/libsy-llm-client/src/media.rs create mode 100644 crates/libsy/src/algorithms/model_as_tool.rs create mode 100644 examples/model-as-tool-media.toml diff --git a/.gitignore b/.gitignore index fdc44128c..d68681faf 100644 --- a/.gitignore +++ b/.gitignore @@ -147,6 +147,7 @@ third party/**/*.* scratch/ plans/ .experiments/ +.switchyard/ # Benchmark result directories bench_results*/ diff --git a/Cargo.lock b/Cargo.lock index 3d0ea85f6..29f04e6a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2291,6 +2291,7 @@ name = "switchyard-llm-client" version = "0.2.0" dependencies = [ "async-trait", + "base64", "futures", "futures-util", "http", @@ -2303,6 +2304,7 @@ dependencies = [ "switchyard-libsy", "switchyard-protocol", "switchyard-translation", + "tempfile", "thiserror 2.0.18", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 07d133bb3..0ea78f101 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ rust-version = "1.96.1" [workspace.dependencies] async-stream = "0.3" async-trait = "0.1" +base64 = "0.22" futures = "0.3" futures-util = "0.3" http = "1" diff --git a/crates/libsy-llm-client/Cargo.toml b/crates/libsy-llm-client/Cargo.toml index 9a9ca6e0a..a22102c40 100644 --- a/crates/libsy-llm-client/Cargo.toml +++ b/crates/libsy-llm-client/Cargo.toml @@ -22,6 +22,7 @@ switchyard-protocol.workspace = true switchyard-translation.workspace = true reqwest.workspace = true async-trait.workspace = true +base64.workspace = true futures.workspace = true futures-util.workspace = true opentelemetry = { version = "0.32", default-features = false, features = ["metrics"] } @@ -38,4 +39,5 @@ tracing-opentelemetry.workspace = true opentelemetry_sdk = { version = "0.32", features = ["metrics", "testing", "trace"] } thiserror.workspace = true tracing-subscriber.workspace = true +tempfile = "3" wiremock = "0.6" diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index e10a17dd6..673f5bff6 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -19,6 +19,7 @@ pub mod backend; pub mod client; pub mod error; +pub mod media; pub mod metrics; mod observability; mod observation; @@ -28,6 +29,7 @@ pub mod run; pub use backend::{Backend, DEFAULT_MAX_RETRIES, HttpBackendConfig}; pub use client::{ModelConfig, TranslatingLlmClient}; pub use error::{LlmClientError, Result}; +pub use media::{CosmosMediaClient, CosmosMediaConfig}; pub use observation::{LlmCallObservation, RunObservation, RunObserver}; pub use raw::RawResponse; pub use run::{ClientRouter, run}; diff --git a/crates/libsy-llm-client/src/media.rs b/crates/libsy-llm-client/src/media.rs new file mode 100644 index 000000000..224690639 --- /dev/null +++ b/crates/libsy-llm-client/src/media.rs @@ -0,0 +1,306 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimal vLLM-Omni Cosmos image client for model-as-a-tool demos. + +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use base64::Engine; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::{Value, json}; +use switchyard_protocol::{ + LlmClientError, LlmResponse, Request, Response, RoutedLlmClient, prompt_text, text_response, +}; + +use crate::Result; + +const IMAGE_REQUEST_TIMEOUT: Duration = Duration::from_secs(30 * 60); + +/// Connection and artifact settings for [`CosmosMediaClient`]. +#[derive(Clone, Debug)] +pub struct CosmosMediaConfig { + /// vLLM-Omni base URL, normally `http://127.0.0.1:8000/v1`. + pub base_url: String, + /// Optional bearer token. + pub api_key: Option, + /// Additional static request headers. + pub extra_headers: BTreeMap, + /// Directory where generated PNG artifacts are written. + pub output_dir: PathBuf, +} + +/// Adapts the Cosmos image endpoint to the routed model-client contract. +/// +/// Input is a one-message text request. Output is a text response listing the generated local +/// paths, keeping the endpoint-specific image protocol outside libsy algorithms. +pub struct CosmosMediaClient { + base_url: String, + api_key: Option, + extra_headers: HeaderMap, + output_dir: PathBuf, + client: reqwest::Client, + sequence: AtomicU64, +} + +impl CosmosMediaClient { + /// Builds a client without contacting the server or creating the artifact directory. + pub fn new(config: CosmosMediaConfig) -> Result { + let base_url = config.base_url.trim().trim_end_matches('/').to_string(); + if base_url.is_empty() { + return Err(LlmClientError::Configuration { + message: "Cosmos media base_url must not be empty".to_string(), + }); + } + if config.output_dir.as_os_str().is_empty() { + return Err(LlmClientError::Configuration { + message: "Cosmos media output_dir must not be empty".to_string(), + }); + } + let extra_headers = header_map(&config.extra_headers)?; + let client = reqwest::Client::builder() + .timeout(IMAGE_REQUEST_TIMEOUT) + .build() + .map_err(transport_error)?; + Ok(Self { + base_url, + api_key: config.api_key, + extra_headers, + output_dir: config.output_dir, + client, + sequence: AtomicU64::new(0), + }) + } + + async fn generate(&self, request: Request) -> Result { + let model = + request + .llm_request + .model + .clone() + .ok_or_else(|| LlmClientError::InvalidRequest { + message: "Cosmos media request has no model".to_string(), + })?; + let prompt = prompt_text(&request.llm_request); + if prompt.trim().is_empty() { + return Err(LlmClientError::InvalidRequest { + message: "Cosmos media request requires a non-empty user prompt".to_string(), + }); + } + tokio::fs::create_dir_all(&self.output_dir) + .await + .map_err(|error| artifact_error(&self.output_dir, error))?; + + let image = self.generate_image(&model, &prompt).await?; + if image.is_empty() { + return Err(LlmClientError::ResponseTranslation( + "Cosmos returned an empty image artifact".to_string(), + )); + } + + let stem = self.next_stem(); + let image_path = self.output_dir.join(format!("{stem}.png")); + write_new(&image_path, &image).await?; + + let completion = format!( + "Generated an image with {model}:\n- Image: `{}`", + image_path.display() + ); + Ok(Response { + llm_response: LlmResponse::Agg(text_response(Some(model), completion)), + metadata: request.metadata, + }) + } + + async fn generate_image(&self, model: &str, prompt: &str) -> Result> { + let response = self + .authorized( + self.client + .post(format!("{}/images/generations", self.base_url)), + ) + .json(&json!({ + "model": model, + "prompt": prompt, + "negative_prompt": "blurry, distorted, low quality", + "size": "1024x1024", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 50, + "guidance_scale": 7.0, + "seed": 42 + })) + .send() + .await + .map_err(transport_error)?; + let body = successful_body(response).await?; + let payload: Value = + serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { + source: Box::new(error), + })?; + let encoded = payload + .get("data") + .and_then(Value::as_array) + .and_then(|data| data.first()) + .and_then(|item| item.get("b64_json")) + .and_then(Value::as_str) + .ok_or_else(|| { + LlmClientError::ResponseTranslation( + "Cosmos image response has no data[0].b64_json".to_string(), + ) + })?; + base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|error| LlmClientError::InvalidResponse { + source: Box::new(error), + }) + } + + fn authorized(&self, mut builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + if let Some(api_key) = &self.api_key { + builder = builder.bearer_auth(api_key); + } + builder.headers(self.extra_headers.clone()) + } + + fn next_stem(&self) -> String { + let epoch_millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); + format!("cosmos-{epoch_millis}-{}-{sequence}", std::process::id()) + } +} + +#[async_trait] +impl RoutedLlmClient for CosmosMediaClient { + async fn call(&self, request: Request) -> Result { + self.generate(request).await + } +} + +fn header_map(headers: &BTreeMap) -> Result { + let mut result = HeaderMap::new(); + for (name, value) in headers { + if ["authorization", "content-length", "content-type", "host"] + .iter() + .any(|reserved| name.eq_ignore_ascii_case(reserved)) + { + return Err(LlmClientError::Configuration { + message: format!("Cosmos media extra_headers cannot set reserved header {name:?}"), + }); + } + let name = HeaderName::from_bytes(name.as_bytes()).map_err(|error| { + LlmClientError::Configuration { + message: format!("invalid Cosmos media header name {name:?}: {error}"), + } + })?; + let value = + HeaderValue::from_str(value).map_err(|error| LlmClientError::Configuration { + message: format!("invalid Cosmos media header value: {error}"), + })?; + result.append(name, value); + } + Ok(result) +} + +async fn successful_body(response: reqwest::Response) -> Result> { + let status = response.status(); + let body = response.bytes().await.map_err(transport_error)?.to_vec(); + if status.is_success() { + return Ok(body); + } + Err(LlmClientError::UpstreamHttp { + status: status.as_u16(), + body: String::from_utf8_lossy(&body).into_owned(), + }) +} + +fn transport_error(error: reqwest::Error) -> LlmClientError { + if error.is_timeout() { + LlmClientError::Timeout { + source: Box::new(error), + } + } else { + LlmClientError::Transport { + source: Box::new(error), + } + } +} + +fn artifact_error(path: &std::path::Path, error: std::io::Error) -> LlmClientError { + LlmClientError::General(format!( + "failed to write generated media artifact {}: {error}", + path.display() + )) +} + +async fn write_new(path: &std::path::Path, bytes: &[u8]) -> Result<()> { + use tokio::io::AsyncWriteExt; + + let mut file = tokio::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .await + .map_err(|error| artifact_error(path, error))?; + file.write_all(bytes) + .await + .map_err(|error| artifact_error(path, error)) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use switchyard_protocol::{ + LlmResponse, Request, RoutedLlmClient, completion_text, text_request, + }; + use tempfile::tempdir; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::{CosmosMediaClient, CosmosMediaConfig}; + + #[tokio::test] + async fn generates_image_artifact_and_returns_path() + -> std::result::Result<(), Box> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/images/generations")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "data": [{"b64_json": "cG5n"}] + }))) + .mount(&server) + .await; + let output = tempdir()?; + let client = CosmosMediaClient::new(CosmosMediaConfig { + base_url: format!("{}/v1/", server.uri()), + api_key: None, + extra_headers: BTreeMap::new(), + output_dir: output.path().join("media"), + })?; + + let response = client + .call(Request { + llm_request: text_request(Some("nvidia/Cosmos3-Nano".to_string()), "A tiny robot"), + raw_request: None, + metadata: None, + }) + .await?; + let LlmResponse::Agg(response) = response.llm_response else { + panic!("media response must be buffered"); + }; + let completion = completion_text(&response); + assert!(completion.contains(".png")); + let mut files = std::fs::read_dir(output.path().join("media"))? + .collect::, _>>()?; + files.sort_by_key(std::fs::DirEntry::file_name); + assert_eq!(files.len(), 1); + assert_eq!(std::fs::read(files[0].path())?, b"png"); + Ok(()) + } +} diff --git a/crates/libsy/README.md b/crates/libsy/README.md index 895c69f82..b865485d3 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -24,6 +24,7 @@ tokio = { version = "1", features = ["macros", "rt"] } | [`Passthrough`] | Always call one configured target. | | [`Random`] | Select among any number of targets using uniform or weighted routing. | | [`LlmTaskClassifier`] | Ask a judge model to choose an efficient or capable target. | +| [`ModelAsTool`] | Let a primary model dispatch a prompt to a specialized media model. | | [`StageRouter`] | Route coding-agent turns from tool and progress signals, with an optional judge fallback. | [`Noop`] is a test helper, not a production routing algorithm. diff --git a/crates/libsy/src/algorithms.rs b/crates/libsy/src/algorithms.rs index c95ede17c..5c5618da1 100644 --- a/crates/libsy/src/algorithms.rs +++ b/crates/libsy/src/algorithms.rs @@ -8,6 +8,7 @@ pub mod fall_through; pub mod llm_class; +pub mod model_as_tool; pub mod noop; pub mod passthrough; pub mod rand; diff --git a/crates/libsy/src/algorithms/model_as_tool.rs b/crates/libsy/src/algorithms/model_as_tool.rs new file mode 100644 index 000000000..d52ebdd7a --- /dev/null +++ b/crates/libsy/src/algorithms/model_as_tool.rs @@ -0,0 +1,425 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Route a model-requested media tool call to a specialized generation model. + +use std::sync::Arc; + +use serde_json::json; +use switchyard_protocol::{ + AggLlmResponse, ContentBlock, Decision, LlmResponse, Metadata, ModelId, Request, Response, + ToolCall, ToolChoice, ToolDefinition, text_request, +}; + +use crate::core::algorithm::{Algorithm, Driver}; +use crate::{LibsyError, Result}; + +/// Name of the synthetic tool exposed to the primary model. +pub const GENERATE_MEDIA_TOOL_NAME: &str = "generate_media"; + +/// Routes a model-selected media tool call to a specialized generation model. +/// +/// The primary model receives one additional `generate_media(prompt)` tool. Normal answers pass +/// through unchanged. A matching tool call becomes a second routed model call whose host client +/// is responsible for turning the prompt into media. +pub struct ModelAsTool { + primary_target: ModelId, + media_target: ModelId, +} + +impl ModelAsTool { + /// Creates a router backed by a reasoning model and a specialized media model. + pub fn new(primary_target: impl Into, media_target: impl Into) -> Self { + Self { + primary_target: primary_target.into(), + media_target: media_target.into(), + } + } +} + +#[async_trait::async_trait] +impl Algorithm for ModelAsTool { + fn name(&self) -> &str { + "model_as_tool" + } + + async fn route(self: Arc, driver: Driver, mut request: Request) -> Result { + reject_reserved_tool_collision(&request)?; + + let stream_requested = request.llm_request.stream; + let request_metadata = request.metadata.clone(); + // The tool choice must be inspected before the algorithm can decide what to return. + request.llm_request.stream = false; + request + .llm_request + .extensions + .fields + .remove("stream_options"); + request + .llm_request + .extensions + .fields + .insert("parallel_tool_calls".to_string(), json!(false)); + request.llm_request.tools.push(media_tool()); + request + .llm_request + .tool_choice + .get_or_insert(ToolChoice::Auto); + // Same-format replay would encode the preserved request instead of the injected tool. + request.llm_request.preservation.requests.clear(); + + tracing::info!( + target = %self.primary_target, + tool = GENERATE_MEDIA_TOOL_NAME, + "offering specialized model as a tool" + ); + let primary_decision = Decision::new(self.primary_target.clone(), true); + driver.decide(primary_decision.clone()).await?; + let primary_response = driver.call_model(request, primary_decision).await?; + let Response { + llm_response, + metadata, + } = primary_response; + let aggregate = llm_response + .into_agg() + .await + .map_err(|error| LibsyError::external("inspecting model-as-tool response", error))?; + + let Some(tool_call) = find_media_tool_call(&aggregate)? else { + return Ok(response_from_aggregate( + aggregate, + metadata, + stream_requested, + )); + }; + let prompt = media_prompt(tool_call)?; + tracing::info!( + target = %self.media_target, + tool = GENERATE_MEDIA_TOOL_NAME, + "dispatching selected specialized model tool" + ); + + let media_request = Request { + llm_request: text_request(None, prompt), + raw_request: None, + metadata: request_metadata, + }; + let media_decision = Decision::new(self.media_target.clone(), true); + driver.decide(media_decision.clone()).await?; + let media_response = driver.call_model(media_request, media_decision).await?; + if !stream_requested { + return Ok(media_response); + } + + let Response { + llm_response, + metadata, + } = media_response; + let aggregate = llm_response + .into_agg() + .await + .map_err(|error| LibsyError::external("streaming model-as-tool response", error))?; + Ok(response_from_aggregate(aggregate, metadata, true)) + } +} + +fn reject_reserved_tool_collision(request: &Request) -> Result<()> { + if request + .llm_request + .tools + .iter() + .any(|tool| tool.name == GENERATE_MEDIA_TOOL_NAME) + { + return Err(LibsyError::AlgorithmError { + message: format!("request already defines reserved tool {GENERATE_MEDIA_TOOL_NAME:?}"), + }); + } + Ok(()) +} + +fn media_tool() -> ToolDefinition { + ToolDefinition { + name: GENERATE_MEDIA_TOOL_NAME.to_string(), + description: Some( + "Generate an image with a specialized local visual model. Call this tool alone only when visual output materially improves the answer. Supply a self-contained prompt describing the scene, composition, and style." + .to_string(), + ), + parameters: json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "A complete image generation prompt." + } + }, + "required": ["prompt"], + "additionalProperties": false + }), + strict: Some(true), + } +} + +fn find_media_tool_call(response: &AggLlmResponse) -> Result> { + let tool_calls = response + .outputs + .iter() + .flat_map(|output| &output.content) + .filter_map(|block| match block { + ContentBlock::ToolCall(tool_call) => Some(tool_call), + _ => None, + }) + .collect::>(); + let media_call = tool_calls + .iter() + .copied() + .find(|tool_call| tool_call.name == GENERATE_MEDIA_TOOL_NAME); + if media_call.is_some() && tool_calls.len() != 1 { + return Err(LibsyError::AlgorithmError { + message: format!( + "{GENERATE_MEDIA_TOOL_NAME} must be called alone; primary model returned {} tool calls", + tool_calls.len() + ), + }); + } + Ok(media_call) +} + +fn media_prompt(tool_call: &ToolCall) -> Result { + tool_call + .arguments + .get("prompt") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|prompt| !prompt.is_empty()) + .map(str::to_string) + .ok_or_else(|| LibsyError::AlgorithmError { + message: format!( + "{GENERATE_MEDIA_TOOL_NAME} call {} requires a non-empty string prompt", + tool_call.id + ), + }) +} + +fn response_from_aggregate( + aggregate: AggLlmResponse, + metadata: Option, + stream: bool, +) -> Response { + Response { + llm_response: if stream { + LlmResponse::Stream(aggregate.into_stream()) + } else { + LlmResponse::Agg(aggregate) + }, + metadata, + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use serde_json::json; + use switchyard_protocol::{ + AggLlmResponse, ContentBlock, LlmRequest, LlmResponse, Message, Request, Response, + ResponseOutput, Role, StopReason, ToolCall, ToolChoice, completion_text, text_response, + }; + + use super::{GENERATE_MEDIA_TOOL_NAME, ModelAsTool}; + use crate::core::algorithm::Algorithm; + use crate::core::testing::test_drive; + + fn request(stream: bool) -> Request { + Request { + llm_request: LlmRequest { + model: Some("auto".to_string()), + messages: vec![Message::text(Role::User, "Make a cinematic launch image")], + stream, + ..LlmRequest::default() + }, + raw_request: None, + metadata: None, + } + } + + fn tool_call_response(prompt: serde_json::Value) -> Response { + Response { + llm_response: LlmResponse::Agg(AggLlmResponse { + outputs: vec![ResponseOutput { + role: Role::Assistant, + content: vec![ContentBlock::ToolCall(ToolCall { + id: "media-1".to_string(), + name: GENERATE_MEDIA_TOOL_NAME.to_string(), + arguments: json!({"prompt": prompt}), + })], + stop_reason: Some(StopReason::ToolUse), + }], + ..AggLlmResponse::default() + }), + metadata: None, + } + } + + #[tokio::test] + async fn normal_answer_passes_through_after_tool_injection() -> crate::Result<()> { + let recorded = Arc::new(Mutex::new(None)); + let captured = Arc::clone(&recorded); + let algorithm: Arc = Arc::new(ModelAsTool::new("primary", "cosmos")); + let mut streamed_request = request(true); + streamed_request + .llm_request + .extensions + .fields + .insert("stream_options".to_string(), json!({"include_usage": true})); + let (trace, response) = + test_drive(algorithm, streamed_request, move |_decision, request| { + let captured = Arc::clone(&captured); + async move { + *captured.lock().expect("recording lock") = Some(request); + Ok(Response { + llm_response: LlmResponse::Agg(text_response( + Some("primary".to_string()), + "plain answer", + )), + metadata: None, + }) + } + }) + .await?; + + let request = recorded + .lock() + .expect("recording lock") + .take() + .expect("primary request"); + assert_eq!(request.llm_request.tools.len(), 1); + assert_eq!(request.llm_request.tools[0].name, GENERATE_MEDIA_TOOL_NAME); + assert_eq!(request.llm_request.tool_choice, Some(ToolChoice::Auto)); + assert!(!request.llm_request.stream); + assert!( + !request + .llm_request + .extensions + .fields + .contains_key("stream_options") + ); + assert_eq!( + request + .llm_request + .extensions + .fields + .get("parallel_tool_calls"), + Some(&json!(false)) + ); + assert!(request.llm_request.preservation.requests.is_empty()); + assert_eq!(trace.len(), 1); + assert!(trace[0].is_answer_call()); + let aggregate = response + .llm_response + .into_agg() + .await + .map_err(|error| crate::LibsyError::external("aggregating test response", error))?; + assert_eq!(completion_text(&aggregate), "plain answer"); + Ok(()) + } + + #[tokio::test] + async fn selected_tool_dispatches_prompt_to_media_target() -> crate::Result<()> { + let calls = Arc::new(Mutex::new(Vec::new())); + let captured = Arc::clone(&calls); + let algorithm: Arc = Arc::new(ModelAsTool::new("primary", "cosmos")); + let (trace, response) = test_drive( + algorithm, + request(false), + move |decision: switchyard_protocol::Decision, request: Request| { + let captured = Arc::clone(&captured); + async move { + captured + .lock() + .expect("recording lock") + .push((decision.selected_model_id().to_string(), request.clone())); + match decision.selected_model_id().as_str() { + "primary" => Ok(tool_call_response(json!("A chrome robot in rain"))), + "cosmos" => Ok(Response { + llm_response: LlmResponse::Agg(text_response( + Some("cosmos".to_string()), + "Image: output.png", + )), + metadata: None, + }), + other => panic!("unexpected target {other}"), + } + } + }, + ) + .await?; + + assert_eq!( + trace + .iter() + .map(|decision| decision.selected_model_id()) + .collect::>(), + ["primary", "cosmos"] + ); + assert!(trace[0].is_answer_call()); + assert!(trace[1].is_answer_call()); + let media_prompt = { + let calls = calls.lock().expect("recording lock"); + switchyard_protocol::prompt_text(&calls[1].1.llm_request) + }; + assert_eq!(media_prompt, "A chrome robot in rain"); + let aggregate = response + .llm_response + .into_agg() + .await + .map_err(|error| crate::LibsyError::external("aggregating test response", error))?; + assert!(completion_text(&aggregate).contains("output.png")); + Ok(()) + } + + #[tokio::test] + async fn rejects_reserved_tool_collision() { + let mut request = request(false); + request.llm_request.tools.push(super::media_tool()); + let algorithm: Arc = Arc::new(ModelAsTool::new("primary", "cosmos")); + let result = test_drive(algorithm, request, |_decision, _request| async move { + unreachable!("collision must fail before a model call") + }) + .await; + + assert!(matches!( + result, + Err(crate::LibsyError::AlgorithmError { message }) if message.contains("reserved tool") + )); + } + + #[tokio::test] + async fn rejects_empty_prompt_and_parallel_media_call() { + for response in [tool_call_response(json!(" ")), { + let mut response = tool_call_response(json!("A robot")); + let LlmResponse::Agg(aggregate) = &mut response.llm_response else { + unreachable!() + }; + aggregate.outputs[0] + .content + .push(ContentBlock::ToolCall(ToolCall { + id: "shell-1".to_string(), + name: "shell".to_string(), + arguments: json!({"command": "pwd"}), + })); + response + }] { + let response = Mutex::new(Some(response)); + let algorithm: Arc = Arc::new(ModelAsTool::new("primary", "cosmos")); + let result = test_drive(algorithm, request(false), move |_decision, _request| { + let response = response.lock().expect("response lock").take(); + async move { Ok(response.expect("one primary call")) } + }) + .await; + assert!(matches!( + result, + Err(crate::LibsyError::AlgorithmError { .. }) + )); + } + } +} diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 3b1ede9ac..9ff89d64a 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -18,6 +18,7 @@ pub use algorithms::llm_class::{ CustomClassifierConfig, CustomClassifierPolicy, LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig, }; +pub use algorithms::model_as_tool::ModelAsTool; pub use algorithms::noop::Noop; pub use algorithms::passthrough::Passthrough; pub use algorithms::rand::{Random, RandomClassifier}; diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 3b3423bb9..6571ce470 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -12,8 +12,8 @@ use pyo3::exceptions::{PyBaseException, PyStopAsyncIteration, PyTypeError, PyVal use pyo3::prelude::*; use switchyard_libsy::{ Algorithm, CallModel, ClassifierContractConfig, HandoffNoteConfig, - LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, - PickerMode, Random, StageRouter, StageRouterConfig, Step as RustStep, StepStream, + LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, ModelAsTool, + Noop, PickerMode, Random, StageRouter, StageRouterConfig, Step as RustStep, StepStream, TaskClassifierConfig, }; use switchyard_protocol::{ @@ -387,6 +387,14 @@ fn random_algorithm( }) } +/// Construct a model-as-a-tool router for a primary model and specialized media model. +#[pyfunction(name = "model_as_tool")] +fn model_as_tool_algorithm(primary_target: String, media_target: String) -> PyAlgorithm { + PyAlgorithm { + inner: Arc::new(ModelAsTool::new(primary_target, media_target)), + } +} + /// Construct task-level LLM classifier routing. #[pyfunction(name = "llm_task_classifier")] #[pyo3(signature = ( @@ -500,6 +508,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { libsy_module.add_class::()?; libsy_module.add_function(wrap_pyfunction!(noop_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(random_algorithm, &libsy_module)?)?; + libsy_module.add_function(wrap_pyfunction!(model_as_tool_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!( llm_task_classifier_algorithm, &libsy_module diff --git a/crates/switchyard-server/CONFIGURATION.md b/crates/switchyard-server/CONFIGURATION.md index efd438109..8dda46ffa 100644 --- a/crates/switchyard-server/CONFIGURATION.md +++ b/crates/switchyard-server/CONFIGURATION.md @@ -28,6 +28,36 @@ outside normal assistant `content`. To support another wire format, add its `ClientFormat` variant and explicit construction match in `src/config.rs`. Add a client type only when a second implementation exists. +## Route a local Cosmos model as a tool + +The demo-only `cosmos_media` client adapts vLLM-Omni's image endpoint to a normal routed model call. +It writes one PNG to `output_dir`, then returns its path as assistant text. Retries must be disabled +because generation has file-producing side effects. + +```toml +[llm_clients.cosmos] +format = "cosmos_media" +base_url = "http://127.0.0.1:8000/v1" +max_retries = 0 +output_dir = ".switchyard/media" + +[targets.cosmos] +id = "nvidia/Cosmos3-Nano" +llm_client = "cosmos" + +[routes.media] +id = "switchyard/media" +type = "model_as_tool" +primary_target = "frontier" +media_target = "cosmos" +tool_calling = true +``` + +`model_as_tool` appends a reserved `generate_media` function with one required string argument, +`prompt`. A matching tool call becomes a second libsy model call to the media target. Otherwise, +the primary response passes through unchanged. Python hosts serve both calls through the same +`Algorithm.run_stream()` interface. + ## Add an algorithm 1. Implement and export the algorithm from `libsy`. diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index d2e59fa18..bc5ee7f1c 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -5,20 +5,20 @@ use std::collections::{BTreeMap, HashSet}; use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use libsy::{ Algorithm, ClassifierContractConfig, CustomClassifierConfig, CustomClassifierPolicy, EscalationJudgeConfig, HandoffNoteConfig, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, - Noop, Passthrough, PickerMode, Random, StageRouter, StageRouterConfig, TargetPrompts, - TaskClassifierConfig, + ModelAsTool, Noop, Passthrough, PickerMode, Random, StageRouter, StageRouterConfig, + TargetPrompts, TaskClassifierConfig, }; use serde::Deserialize; use serde_json::Value; use switchyard_llm_client::{ - Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig, - TranslatingLlmClient, + Backend, ClientRouter, CosmosMediaClient, CosmosMediaConfig, DEFAULT_MAX_RETRIES, + HttpBackendConfig, ModelConfig, TranslatingLlmClient, }; use switchyard_protocol::{ModelId, RoutedLlmClient}; @@ -98,6 +98,7 @@ impl ServerConfig { "route {route_name} context_window must be greater than zero" ))); } + self.validate_specialized_route(route_name, config)?; let algorithm = build_algorithm(route_name, config, &targets)?; let client = self.build_client_router(config, &clients)?; let count_tokens_target = self.build_count_tokens_target(config, &clients); @@ -112,40 +113,136 @@ impl ServerConfig { ServerState::new_with_capabilities(routes) } - fn build_clients(&self) -> ServerResult>> { - let mut models_by_client = self + fn validate_specialized_route( + &self, + route_name: &str, + route: &RouteConfig, + ) -> ServerResult<()> { + let RouteConfig::ModelAsTool { + primary_target, + media_target, + .. + } = route + else { + return Ok(()); + }; + if matches!( + self.target_client_format(route_name, primary_target)?, + ClientFormat::CosmosMedia + ) { + return Err(ServerError::new(format!( + "model_as_tool route {route_name} primary_target must use an LLM client" + ))); + } + if !matches!( + self.target_client_format(route_name, media_target)?, + ClientFormat::CosmosMedia + ) { + return Err(ServerError::new(format!( + "model_as_tool route {route_name} media_target must use a cosmos_media client" + ))); + } + Ok(()) + } + + fn target_client_format( + &self, + route_name: &str, + target_name: &str, + ) -> ServerResult { + let target = self.targets.get(target_name).ok_or_else(|| { + ServerError::new(format!( + "route {route_name} references unknown target {target_name}" + )) + })?; + self.llm_clients + .get(&target.llm_client) + .map(|client| client.format) + .ok_or_else(|| { + ServerError::new(format!( + "target {target_name} references unknown llm client {}", + target.llm_client + )) + }) + } + + fn build_clients(&self) -> ServerResult> { + let mut targets_by_client = self .llm_clients .keys() .map(|name| (name.clone(), Vec::new())) - .collect::>>(); + .collect::>>(); for name in self.llm_clients.keys() { validate_value("llm client name", name)?; } for (target_name, target) in &self.targets { - let client_config = self.llm_clients.get(&target.llm_client).ok_or_else(|| { + self.llm_clients.get(&target.llm_client).ok_or_else(|| { ServerError::new(format!( "target {target_name} references unknown llm client {}", target.llm_client )) })?; - let model_configs = models_by_client + targets_by_client .get_mut(&target.llm_client) - .ok_or_else(|| ServerError::new("validated llm client was not initialized"))?; - model_configs.push(ModelConfig::new( - target.id.clone(), - build_backend(&target.llm_client, client_config, &target.extra_body)?, - None, - )); + .ok_or_else(|| ServerError::new("validated llm client was not initialized"))? + .push((target_name, target)); } let mut clients = BTreeMap::new(); - for (name, model_configs) in models_by_client { - let client = Arc::new( - TranslatingLlmClient::new(&model_configs) - .map_err(|error| ServerError::new(error.to_string()))?, - ); - clients.insert(name, client); + for (name, config) in &self.llm_clients { + let configured_targets = targets_by_client.remove(name).unwrap_or_default(); + let client = match config.format { + ClientFormat::CosmosMedia => { + if config.max_retries != 0 { + return Err(ServerError::new(format!( + "Cosmos media llm client {name} requires max_retries = 0 to avoid duplicate generation" + ))); + } + if let Some((target_name, _)) = configured_targets + .iter() + .find(|(_, target)| !target.extra_body.is_empty()) + { + return Err(ServerError::new(format!( + "Cosmos media target {target_name} does not support extra_body" + ))); + } + BuiltClient::Cosmos(Arc::new( + CosmosMediaClient::new(CosmosMediaConfig { + base_url: config.base_url.clone(), + api_key: resolve_api_key(name, config)?, + extra_headers: config.extra_headers.clone(), + output_dir: config + .output_dir + .clone() + .unwrap_or_else(|| PathBuf::from(".switchyard/media")), + }) + .map_err(|error| ServerError::new(error.to_string()))?, + )) + } + _ => { + if config.output_dir.is_some() { + return Err(ServerError::new(format!( + "llm client {name} output_dir is only valid for cosmos_media" + ))); + } + let model_configs = configured_targets + .into_iter() + .map(|(_, target)| { + Ok(ModelConfig::new( + target.id.clone(), + build_backend(name, config, &target.extra_body)?, + None, + )) + }) + .collect::>>()?; + BuiltClient::Translating(Arc::new( + TranslatingLlmClient::new(&model_configs) + .map_err(|error| ServerError::new(error.to_string()))?, + )) + } + }; + clients.insert(name.clone(), client); } Ok(clients) } @@ -171,7 +268,7 @@ impl ServerConfig { fn build_client_router( &self, route: &RouteConfig, - clients: &BTreeMap>, + clients: &BTreeMap, ) -> ServerResult { let by_model = route .callable_target_names() @@ -183,8 +280,7 @@ impl ServerConfig { let client = clients.get(&target.llm_client).ok_or_else(|| { ServerError::new(format!("target {name} has no constructed llm client")) })?; - let client: Arc = client.clone(); - Ok((target.id.clone(), client)) + Ok((target.id.clone(), client.routed())) }) .collect::>()?; Ok(ClientRouter::new(by_model)) @@ -193,7 +289,7 @@ impl ServerConfig { fn build_count_tokens_target( &self, route_config: &RouteConfig, - clients: &BTreeMap>, + clients: &BTreeMap, ) -> Option { route_config .routing_target_names() @@ -202,6 +298,7 @@ impl ServerConfig { .filter_map(|(index, name)| { let target = self.targets.get(name)?; let client = clients.get(&target.llm_client)?; + let client = client.translating()?; client.supports_count_tokens(&target.id).then_some(( count_tokens_priority(name, &target.id), index, @@ -217,6 +314,27 @@ impl ServerConfig { } } +enum BuiltClient { + Translating(Arc), + Cosmos(Arc), +} + +impl BuiltClient { + fn routed(&self) -> Arc { + match self { + Self::Translating(client) => client.clone(), + Self::Cosmos(client) => client.clone(), + } + } + + fn translating(&self) -> Option<&Arc> { + match self { + Self::Translating(client) => Some(client), + Self::Cosmos(_) => None, + } + } +} + // Prefer known Claude families, then preserve the route's target order. fn count_tokens_priority(target_name: &str, model_id: &ModelId) -> usize { let target_name = target_name.to_ascii_lowercase(); @@ -237,6 +355,7 @@ struct LlmClientConfig { extra_headers: BTreeMap, #[serde(default = "default_max_retries")] max_retries: u32, + output_dir: Option, } #[derive(Debug, Deserialize)] @@ -256,6 +375,8 @@ enum ClientFormat { OpenAiResponses, #[serde(rename = "anthropic_messages")] AnthropicMessages, + #[serde(rename = "cosmos_media")] + CosmosMedia, } #[derive(Clone, Debug, Deserialize)] @@ -356,6 +477,17 @@ enum RouteConfig { reasoning: Option, target: String, }, + ModelAsTool { + id: ModelId, + #[serde(default)] + context_window: Option, + #[serde(default)] + tool_calling: Option, + #[serde(default)] + reasoning: Option, + primary_target: String, + media_target: String, + }, LlmClassifier { id: ModelId, #[serde(default)] @@ -468,6 +600,7 @@ impl RouteConfig { Noop { id, .. } | Random { id, .. } | LlmClassifier { id, .. } + | ModelAsTool { id, .. } | Passthrough { id, .. } | StageRouter { id, .. } => id, } @@ -479,6 +612,7 @@ impl RouteConfig { Self::Noop { .. } => Vec::new(), Self::Random { targets, .. } => targets.iter().map(String::as_str).collect(), Self::Passthrough { target, .. } => vec![target], + Self::ModelAsTool { primary_target, .. } => vec![primary_target], Self::LlmClassifier { mode, strong_target, @@ -525,6 +659,7 @@ impl RouteConfig { classifier: Some(classifier), .. } => names.push(&classifier.target), + Self::ModelAsTool { media_target, .. } => names.push(media_target), _ => {} } names @@ -551,6 +686,12 @@ impl RouteConfig { reasoning, .. } + | ModelAsTool { + context_window, + tool_calling, + reasoning, + .. + } | LlmClassifier { context_window, tool_calling, @@ -771,7 +912,26 @@ fn build_backend( "llm client {client_name} max_retries must be at most {MAX_CONFIGURED_RETRIES}" ))); } - let api_key = config + let api_key = resolve_api_key(client_name, config)?; + let http = HttpBackendConfig { + base_url: base_url.to_string(), + api_key, + extra_headers: config.extra_headers.clone(), + extra_body: extra_body.clone(), + max_retries: config.max_retries, + }; + match config.format { + ClientFormat::OpenAiChat => Ok(Backend::OpenAiChat(http)), + ClientFormat::OpenAiResponses => Ok(Backend::OpenAiResponses(http)), + ClientFormat::AnthropicMessages => Ok(Backend::Anthropic(http)), + ClientFormat::CosmosMedia => Err(ServerError::new(format!( + "Cosmos media llm client {client_name} cannot be built as a translating backend" + ))), + } +} + +fn resolve_api_key(client_name: &str, config: &LlmClientConfig) -> ServerResult> { + config .api_key_env .as_deref() .map(|variable| { @@ -792,19 +952,7 @@ fn build_backend( } Ok(api_key) }) - .transpose()?; - let http = HttpBackendConfig { - base_url: base_url.to_string(), - api_key, - extra_headers: config.extra_headers.clone(), - extra_body: extra_body.clone(), - max_retries: config.max_retries, - }; - Ok(match config.format { - ClientFormat::OpenAiChat => Backend::OpenAiChat(http), - ClientFormat::OpenAiResponses => Backend::OpenAiResponses(http), - ClientFormat::AnthropicMessages => Backend::Anthropic(http), - }) + .transpose() } const fn default_max_retries() -> u32 { @@ -834,6 +982,15 @@ fn build_algorithm( let target = resolve_target_model_id(route_name, target, targets)?; Ok(Arc::new(Passthrough::new(target))) } + RouteConfig::ModelAsTool { + primary_target, + media_target, + .. + } => { + let primary = resolve_target_model_id(route_name, primary_target, targets)?; + let media = resolve_target_model_id(route_name, media_target, targets)?; + Ok(Arc::new(ModelAsTool::new(primary, media))) + } RouteConfig::LlmClassifier { classifier_target, .. } => { @@ -1097,6 +1254,84 @@ target = "weak" Ok(()) } + #[test] + fn builds_model_as_tool_with_cosmos_media_client() -> ServerResult<()> { + let config = r#" +schema_version = 1 + +[llm_clients.primary] +format = "openai_chat" +base_url = "https://example.test/v1" + +[llm_clients.cosmos] +format = "cosmos_media" +base_url = "http://127.0.0.1:8000/v1" +max_retries = 0 +output_dir = ".switchyard/test-media" + +[targets.primary] +id = "frontier/model" +llm_client = "primary" + +[targets.cosmos] +id = "nvidia/Cosmos3-Nano" +llm_client = "cosmos" + +[routes.media] +id = "switchyard/media" +type = "model_as_tool" +primary_target = "primary" +media_target = "cosmos" +tool_calling = true +"#; + + let state = server_state_from_toml(config)?; + assert_eq!(state.models().collect::>(), ["switchyard/media"]); + Ok(()) + } + + #[test] + fn validates_model_as_tool_client_roles_and_retry_safety() { + let valid = r#" +schema_version = 1 + +[llm_clients.primary] +format = "openai_chat" +base_url = "https://example.test/v1" + +[llm_clients.cosmos] +format = "cosmos_media" +base_url = "http://127.0.0.1:8000/v1" +max_retries = 0 + +[targets.primary] +id = "frontier/model" +llm_client = "primary" + +[targets.cosmos] +id = "nvidia/Cosmos3-Nano" +llm_client = "cosmos" + +[routes.media] +id = "switchyard/media" +type = "model_as_tool" +primary_target = "primary" +media_target = "cosmos" +"#; + + let retried = valid.replace("max_retries = 0", "max_retries = 1"); + assert!(error_message(&retried).contains("requires max_retries = 0")); + + let wrong_media = valid.replace("media_target = \"cosmos\"", "media_target = \"primary\""); + assert!(error_message(&wrong_media).contains("must use a cosmos_media client")); + + let wrong_primary = valid.replace( + "primary_target = \"primary\"", + "primary_target = \"cosmos\"", + ); + assert!(error_message(&wrong_primary).contains("must use an LLM client")); + } + #[test] fn an_escalation_table_switches_the_classifier_route_to_escalation() -> ServerResult<()> { // Present: the classifier target judges the weak tier's reply each turn instead of diff --git a/examples/model-as-tool-media.toml b/examples/model-as-tool-media.toml new file mode 100644 index 000000000..dae32105d --- /dev/null +++ b/examples/model-as-tool-media.toml @@ -0,0 +1,41 @@ +schema_version = 1 + +# Both demo routes use the same frontier target. The only difference is whether +# local Cosmos is exposed to it through generate_media(prompt). + +[llm_clients.nvidia_anthropic] +format = "anthropic_messages" +base_url = "https://inference-api.nvidia.com" +api_key_env = "NVIDIA_API_KEY" + +[llm_clients.cosmos] +format = "cosmos_media" +base_url = "http://127.0.0.1:8000/v1" +max_retries = 0 +output_dir = ".switchyard/media" + +[targets.frontier] +id = "aws/anthropic/bedrock-claude-opus-4-8" +llm_client = "nvidia_anthropic" + +[targets.frontier.extra_body.output_config] +effort = "medium" + +[targets.cosmos] +id = "nvidia/Cosmos3-Nano" +llm_client = "cosmos" + +[routes.plain] +id = "switchyard/frontier" +type = "passthrough" +target = "frontier" +tool_calling = true +reasoning = true + +[routes.media] +id = "switchyard/media" +type = "model_as_tool" +primary_target = "frontier" +media_target = "cosmos" +tool_calling = true +reasoning = true diff --git a/switchyard/libsy/algorithms.py b/switchyard/libsy/algorithms.py index 92e9e747e..64fc33c2b 100644 --- a/switchyard/libsy/algorithms.py +++ b/switchyard/libsy/algorithms.py @@ -4,8 +4,9 @@ """Factories for Rust-owned libsy algorithms.""" from switchyard_rust.libsy import llm_task_classifier as llm_task_classifier +from switchyard_rust.libsy import model_as_tool as model_as_tool from switchyard_rust.libsy import noop as noop from switchyard_rust.libsy import random as random from switchyard_rust.libsy import stage_router as stage_router -__all__ = ["llm_task_classifier", "noop", "random", "stage_router"] +__all__ = ["llm_task_classifier", "model_as_tool", "noop", "random", "stage_router"] diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index 5dd607566..520290c7d 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -21,6 +21,7 @@ "Step", "TaskClassifierConfig", "llm_task_classifier", + "model_as_tool", "noop", "random", "stage_router", @@ -131,6 +132,8 @@ def llm_task_classifier( config: TaskClassifierConfig, ) -> Algorithm: ... + def model_as_tool(primary_target: str, media_target: str) -> Algorithm: ... + def stage_router( capable_target: str, efficient_target: str, diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index 00ca23b1c..bcbed8f21 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -110,6 +110,57 @@ async def test_random_streams_complex_steps_and_accepts_a_dictionary_response() assert response["outputs"][0]["content"] == [{"type": "text", "text": "fast"}] +async def test_model_as_tool_dispatches_selected_media_prompt() -> None: + class PrimaryClient(EchoClient): + async def call(self, request: dict[str, Any]) -> dict[str, Any]: + self.calls.append(request) + return { + "model": self.model, + "outputs": [ + { + "role": "assistant", + "content": [ + { + "type": "tool_call", + "id": "media-1", + "name": "generate_media", + "arguments": {"prompt": "A cinematic robot launch"}, + } + ], + "stop_reason": "tool_use", + } + ], + } + + primary = PrimaryClient("frontier") + media = EchoClient("nvidia/Cosmos3-Nano") + decisions, response = await run_algorithm( + algorithms.model_as_tool("frontier", "nvidia/Cosmos3-Nano"), + {"frontier": primary, "nvidia/Cosmos3-Nano": media}, + ) + + assert [decision.selected_model_id for decision in decisions] == [ + "frontier", + "nvidia/Cosmos3-Nano", + ] + assert primary.calls[0]["tools"][0]["name"] == "generate_media" + assert primary.calls[0]["tools"][0]["parameters"] == { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "A complete image generation prompt.", + } + }, + "required": ["prompt"], + "additionalProperties": False, + } + assert media.calls[0]["messages"][0]["content"] == [ + {"type": "text", "text": "A cinematic robot launch"} + ] + assert response["model"] == "nvidia/Cosmos3-Nano" + + async def test_into_parts_supports_decision_only_routing() -> None: algorithm = algorithms.random(["fast"])