diff --git a/crates/external_router/src/metrics.rs b/crates/external_router/src/metrics.rs index bac8e93196..822ce9bc9b 100644 --- a/crates/external_router/src/metrics.rs +++ b/crates/external_router/src/metrics.rs @@ -126,6 +126,7 @@ pub mod metrics_labels { pub const ENDPOINT_CHAT: &str = "chat"; pub const ENDPOINT_GENERATE: &str = "generate"; pub const ENDPOINT_RESPONSES: &str = "responses"; + pub const ENDPOINT_DECISIONS: &str = "decisions"; pub const ENDPOINT_COMPLETIONS: &str = "completions"; pub const ENDPOINT_RERANK: &str = "rerank"; pub const ENDPOINT_EMBEDDINGS: &str = "embeddings"; diff --git a/e2e_test/router/test_decisions.py b/e2e_test/router/test_decisions.py new file mode 100644 index 0000000000..ef5bc72281 --- /dev/null +++ b/e2e_test/router/test_decisions.py @@ -0,0 +1,155 @@ +"""OpenAI Decisions through a regular HTTP gateway and real SGLang inference. + +Uses raw HTTP because the installed OpenAI SDK can predate Decisions. The +SGLang adapter performs prefill scoring, so every response has zero output +usage. Run with E2E_ENGINE=sglang E2E_RUNTIME=sglang pytest +e2e_test/router/test_decisions.py -v. +""" + +from __future__ import annotations + +import math + +import httpx +import pytest + +_MODEL = "Qwen/Qwen3-4B-Instruct-2507" + + +def _post(gateway, body: dict) -> httpx.Response: + return httpx.post(f"{gateway.base_url}/v1/decisions", json=body, timeout=120.0) + + +def _assert_distribution(answer: dict) -> None: + probabilities = [item["probability"] for item in answer["probabilities"]] + assert probabilities + assert all(math.isfinite(p) and 0 <= p <= 1 for p in probabilities) + assert sum(probabilities) == pytest.approx(1.0, abs=1e-6) + assert math.isfinite(answer["confidence"]) + assert 0 <= answer["confidence"] <= 1 + + +@pytest.mark.engine("sglang") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.model(_MODEL) +@pytest.mark.parametrize("setup_backend", ["http"], indirect=True) +class TestDecisions: + def test_all_question_types_order_names_and_prefill_usage(self, setup_backend): + _, model, _, gateway = setup_backend + response = _post( + gateway, + { + "model": model, + "input": "The sky is blue. Two plus two is four. The review is excellent.", + "questions": [ + {"type": "predicate", "instructions": "Is two plus two four?"}, + { + "type": "choice", + "name": "repeated", + "instructions": "What color is the sky?", + "choices": [{"value": "blue"}, {"value": "red"}], + }, + { + "type": "score", + "name": "repeated", + "instructions": "Rate the review quality.", + "levels": [ + {"label": "poor", "description": "Bad quality"}, + {"label": "excellent", "description": "Very good quality"}, + ], + }, + { + "type": "choice", + "instructions": "Is the sky blue?", + "choices": [ + {"value": True, "description": "Yes, the sky is blue"}, + {"value": False, "description": "No, it is another color"}, + ], + }, + ], + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json") + result = response.json() + answers = result["answers"] + assert [a["type"] for a in answers] == ["predicate", "choice", "score", "choice"] + assert [a["name"] for a in answers] == [None, "repeated", "repeated", None] + assert 0.5 < answers[0]["probability"] <= 1 + assert answers[1]["choice"] == "blue" + assert [p["value"] for p in answers[1]["probabilities"]] == ["blue", "red"] + assert [p["value"] for p in answers[2]["probabilities"]] == [0, 1] + assert [p["label"] for p in answers[2]["probabilities"]] == ["poor", "excellent"] + # Scores are expectations over level indices, not the winning level. + assert answers[2]["score"] == pytest.approx( + sum(p["value"] * p["probability"] for p in answers[2]["probabilities"]), + abs=1e-6, + ) + assert type(answers[3]["choice"]) is bool + assert [p["value"] for p in answers[3]["probabilities"]] == [True, False] + assert all(type(p["value"]) is bool for p in answers[3]["probabilities"]) + for answer in answers[1:]: + _assert_distribution(answer) + usage = result["usage"] + assert usage["input_tokens"] > 0 + assert usage["output_tokens"] == 0 + assert usage["total_tokens"] == usage["input_tokens"] + assert usage["output_tokens_details"]["reasoning_tokens"] == 0 + assert set(usage["input_tokens_details"]) >= {"cached_tokens", "cache_write_tokens"} + + def test_ordered_user_messages_are_accepted(self, setup_backend): + _, model, _, gateway = setup_backend + response = _post( + gateway, + { + "model": model, + "input": [ + {"role": "user", "content": "The delivery arrived."}, + { + "role": "user", + "content": [{"type": "input_text", "text": "The package is intact."}], + }, + ], + "questions": [ + {"type": "predicate", "instructions": "Did the delivery arrive intact?"} + ], + }, + ) + assert response.status_code == 200, response.text + answer = response.json()["answers"][0] + assert answer["type"] == "predicate" + assert answer["name"] is None + assert 0.5 < answer["probability"] <= 1 + + @pytest.mark.parametrize( + "invalid", ["missing_model", "unknown_type", "numeric_choice", "image"] + ) + def test_invalid_or_unsupported_requests_are_client_errors(self, setup_backend, invalid): + _, model, _, gateway = setup_backend + body = { + "model": model, + "input": "A short text.", + "questions": [{"type": "predicate", "instructions": "Is this text?"}], + } + if invalid == "missing_model": + del body["model"] + elif invalid == "unknown_type": + body["questions"][0]["type"] = "unknown" + elif invalid == "numeric_choice": + body["questions"] = [ + {"type": "choice", "instructions": "Choose", "choices": [{"value": 1}]} + ] + else: + # Valid OpenAI image input; SGLang's pinned text-only adapter rejects it. + body["input"] = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": "data:image/png;base64,AA=="}], + } + ] + response = _post(gateway, body) + assert 400 <= response.status_code < 500, response.text + if invalid == "image": + assert response.status_code == 400, response.text + assert isinstance(response.json()["error"]["message"], str) diff --git a/model_gateway/src/routers/gateway.rs b/model_gateway/src/routers/gateway.rs index 4952a2903d..cf9c0ce948 100644 --- a/model_gateway/src/routers/gateway.rs +++ b/model_gateway/src/routers/gateway.rs @@ -22,6 +22,7 @@ use openai_protocol::{ chat::ChatCompletionRequest, classify::ClassifyRequest, completion::CompletionRequest, + decisions::DecisionsRequest, embedding::EmbeddingRequest, generate::GenerateRequest, interactions::InteractionsRequest, @@ -588,6 +589,43 @@ impl RouterTrait for Gateway { .await } + async fn route_decisions( + &self, + headers: Option<&HeaderMap>, + tenant_meta: &TenantRequestMeta, + body: DecisionsRequest, + model_id: &str, + ) -> Response { + let router = if self.enable_igw { + let snapshot = self.worker_registry.get_routing_snapshot(model_id); + if snapshot.is_empty() { + return route_error::model_not_found(model_id); + } + // Decisions is supported only by regular HTTP workers. General + // weighted dispatch could send the same model to PD or gRPC. + if snapshot.pool(RoutingPool::HttpRegular).is_empty() { + return route_error::not_implemented( + "decisions_not_supported", + "Decisions requires a regular HTTP worker for this model", + ); + } + self.routers.load().get(&router_ids::HTTP_REGULAR).cloned() + } else { + self.select_router_for_request(Some(model_id)) + }; + match router { + Some(router) => { + router + .route_decisions(headers, tenant_meta, body, model_id) + .await + } + None => route_error::not_implemented( + "decisions_not_supported", + "Decisions requires a regular HTTP router", + ), + } + } + async fn route_audio_transcriptions( &self, headers: Option<&HeaderMap>, @@ -879,6 +917,101 @@ mod tests { RouteRequestMeta::new(TenantKey::from("test-tenant")) } + #[derive(Debug)] + struct DecisionsStubRouter; + + #[async_trait] + impl RouterTrait for DecisionsStubRouter { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn route_decisions( + &self, + _headers: Option<&HeaderMap>, + _tenant_meta: &TenantRequestMeta, + _body: DecisionsRequest, + _model_id: &str, + ) -> Response { + StatusCode::OK.into_response() + } + + fn router_type(&self) -> &'static str { + "decisions" + } + } + + async fn decision_status(gateway: &Gateway, model: &str) -> StatusCode { + let body = serde_json::from_value(serde_json::json!({ + "model": model, "input": "evidence", + "questions": [{"type": "predicate", "instructions": "Is it evidence?"}] + })) + .unwrap(); + gateway + .route_decisions(None, &test_tenant_meta(), body, model) + .await + .status() + } + + #[tokio::test] + async fn decisions_use_only_regular_http_workers_in_mixed_fleets() { + for has_http_regular in [true, false] { + let gateway = test_gateway(true); + gateway.register_router(router_ids::HTTP_REGULAR, Arc::new(DecisionsStubRouter)); + gateway.register_router(router_ids::HTTP_PD, Arc::new(PdStubRouter)); + gateway.register_router(router_ids::GRPC_REGULAR, Arc::new(PdStubRouter)); + gateway.set_default_router(router_ids::GRPC_REGULAR); + let workers = [ + (ConnectionMode::Http, WorkerType::Regular), + (ConnectionMode::Http, WorkerType::Prefill), + (ConnectionMode::Http, WorkerType::Decode), + (ConnectionMode::Grpc, WorkerType::Regular), + (ConnectionMode::Zmq, WorkerType::Regular), + ]; + for (index, (mode, role)) in workers.into_iter().enumerate() { + if index == 0 && !has_http_regular { + continue; + } + gateway + .worker_registry + .register(Arc::new( + BasicWorkerBuilder::new(format!("http://worker-{index}:8080")) + .connection_mode(mode) + .worker_type(role) + .model(ModelCard::new("m").with_alias("alias")) + .build(), + )) + .unwrap(); + } + let expected = if has_http_regular { + StatusCode::OK + } else { + StatusCode::NOT_IMPLEMENTED + }; + for _ in 0..32 { + assert_eq!(decision_status(&gateway, "m").await, expected); + assert_eq!(decision_status(&gateway, "alias").await, expected); + } + assert_eq!( + decision_status(&gateway, "missing").await, + StatusCode::NOT_FOUND + ); + } + } + + #[tokio::test] + async fn decisions_respect_single_router_configuration() { + let gateway = test_gateway(false); + gateway.register_router(router_ids::HTTP_REGULAR, Arc::new(DecisionsStubRouter)); + assert_eq!(decision_status(&gateway, "m").await, StatusCode::OK); + gateway.register_router(router_ids::HTTP_PD, Arc::new(PdStubRouter)); + gateway.set_default_router(router_ids::HTTP_PD); + assert_eq!( + decision_status(&gateway, "m").await, + StatusCode::NOT_IMPLEMENTED + ); + } + #[tokio::test] async fn count_tokens_selects_http_workers_without_a_decode_leg() { for (role, router_id) in [ diff --git a/model_gateway/src/routers/grpc/utils/metrics.rs b/model_gateway/src/routers/grpc/utils/metrics.rs index f5afa3cf8d..8b11bc7752 100644 --- a/model_gateway/src/routers/grpc/utils/metrics.rs +++ b/model_gateway/src/routers/grpc/utils/metrics.rs @@ -12,6 +12,7 @@ pub(crate) fn route_to_endpoint(route: &str) -> &'static str { "/v1/completions" => metrics_labels::ENDPOINT_COMPLETIONS, "/v1/rerank" => metrics_labels::ENDPOINT_RERANK, "/v1/responses" => metrics_labels::ENDPOINT_RESPONSES, + "/v1/decisions" => metrics_labels::ENDPOINT_DECISIONS, "/v1/messages" => metrics_labels::ENDPOINT_MESSAGES, "/v1/audio/transcriptions" => metrics_labels::ENDPOINT_AUDIO_TRANSCRIPTIONS, _ => "other", diff --git a/model_gateway/src/routers/http/decisions.rs b/model_gateway/src/routers/http/decisions.rs new file mode 100644 index 0000000000..74327bcdbe --- /dev/null +++ b/model_gateway/src/routers/http/decisions.rs @@ -0,0 +1,932 @@ +//! SGLang 0.5.21 translation for the public OpenAI Decisions contract. +//! +//! System One scores the answer labels without generating reasoning. Its +//! confidence is SGLang's documented statistic; it is not asserted to be +//! numerically equivalent to OpenAI's confidence. +//! System One schedules within its own data-parallel group and does not honor +//! SMG's `data_parallel_rank` pinning, so per-rank cache affinity is unavailable. + +use axum::{ + body::{to_bytes, Body}, + http::{header::CONTENT_TYPE, HeaderValue}, + response::{IntoResponse, Response}, + Json, +}; +use openai_protocol::decisions::{ + DecisionAnswer, DecisionChoiceProbability, DecisionInput, DecisionInputContent, + DecisionInputPart, DecisionInputTokensDetails, DecisionOutputTokensDetails, DecisionQuestion, + DecisionResponse, DecisionScoreProbability, DecisionUsage, DecisionValue, DecisionsRequest, +}; +use serde::{de, Deserialize, Deserializer}; +use serde_json::{json, Map, Value}; + +use crate::routers::error; + +pub(super) const UPSTREAM_ROUTE: &str = "/v1/systemone"; + +pub(super) struct SglangDecisionAdapter { + request: Value, + questions: Vec, +} + +enum MappedQuestion { + Predicate { + name: Option, + }, + Choice { + name: Option, + choices: Vec, + }, + Score { + name: Option, + labels: Vec, + }, +} + +impl MappedQuestion { + fn name(&self) -> Option<&str> { + match self { + Self::Predicate { name } | Self::Choice { name, .. } | Self::Score { name, .. } => { + name.as_deref() + } + } + } +} + +impl SglangDecisionAdapter { + pub fn new(request: &DecisionsRequest) -> Result { + reject_extensions(&request.other, "request")?; + if let DecisionInput::Messages(messages) = &request.input { + for message in messages { + reject_extensions(&message.other, "input message")?; + if let DecisionInputContent::Parts(parts) = &message.content { + for part in parts { + match part { + DecisionInputPart::InputText { other, .. } => { + reject_extensions(other, "input text part")?; + } + DecisionInputPart::InputImage { .. } => { + return Err( + "SGLang 0.5.21 Decisions does not support image input".into() + ); + } + } + } + } + } + } + if request.input.text().trim().is_empty() { + return Err("SGLang Decisions requires nonempty text input".into()); + } + let mut questions = Map::new(); + for (index, question) in request.questions.iter().enumerate() { + if question.instructions().trim().is_empty() { + return Err(format!( + "SGLang Decisions question {index} requires nonempty instructions" + )); + } + let mapped = match question { + DecisionQuestion::Predicate { + instructions, + other, + .. + } => { + reject_extensions(other, "predicate question")?; + json!({"type":"noul", "instructions":instructions}) + } + DecisionQuestion::Choice { + instructions, + choices, + other, + .. + } => { + reject_extensions(other, "choice question")?; + if choices.is_empty() || choices.len() > 255 { + return Err(format!( + "SGLang Decisions question {index} requires 1 to 255 choices" + )); + } + let mut criteria = Map::new(); + for (option_index, choice) in choices.iter().enumerate() { + reject_extensions(&choice.other, "choice option")?; + let mut description = json!({"value":choice.value}); + if let Some(value) = &choice.description { + description["description"] = json!(value); + } + criteria.insert(format!("o{option_index}"), description); + } + json!({"type":"choice", "instructions":instructions, "criteria":criteria}) + } + DecisionQuestion::Score { + instructions, + levels, + other, + .. + } => { + reject_extensions(other, "score question")?; + if levels.is_empty() || levels.len() > 10 { + return Err(format!( + "SGLang Decisions question {index} requires 1 to 10 score levels" + )); + } + let mut criteria = Vec::with_capacity(levels.len()); + for level in levels { + reject_extensions(&level.other, "score level")?; + let mut description = json!({"label":level.label}); + if let Some(value) = &level.description { + description["description"] = json!(value); + } + criteria.push(description); + } + json!({"type":"score", "instructions":instructions, "criteria":criteria}) + } + }; + questions.insert(format!("q{index}"), mapped); + } + if questions.is_empty() { + return Err("SGLang Decisions requires at least one question".into()); + } + // Structured text evidence is rendered by System One as ordered JSON, + // preserving message/part boundaries and roles. Images are rejected above. + // safety_identifier is consumed at the gateway boundary, not a backend field. + Ok(Self { + request: json!({"model":request.model, "state":request.input, "questions":questions}), + questions: request + .questions + .iter() + .map(|question| { + let name = question.name().map(str::to_owned); + match question { + DecisionQuestion::Predicate { .. } => MappedQuestion::Predicate { name }, + DecisionQuestion::Choice { choices, .. } => MappedQuestion::Choice { + name, + choices: choices.iter().map(|choice| choice.value.clone()).collect(), + }, + DecisionQuestion::Score { levels, .. } => MappedQuestion::Score { + name, + labels: levels.iter().map(|level| level.label.clone()).collect(), + }, + } + }) + .collect(), + }) + } + + pub fn request(&self) -> &Value { + &self.request + } + + /// Drop evidence, instructions and descriptions after dispatch serialization. + /// Only the original names, typed choices and score labels remain for mapping. + pub fn release_request(&mut self) { + self.request = Value::Null; + } + + pub async fn convert_response(self, response: Response, max_payload_size: usize) -> Response { + let (mut parts, body) = response.into_parts(); + let converted = match to_bytes(body, max_payload_size).await { + Err(_) => error::bad_gateway( + "invalid_decisions_response", + "SGLang Decisions response exceeds the payload limit or could not be read", + ), + Ok(bytes) => { + let parsed = serde_json::from_slice::(&bytes).map(|value| value.0); + if parts.status.is_success() { + match parsed + .map_err(|err| format!("Invalid SGLang Decisions JSON: {err}")) + .and_then(|value| self.map_response(&value)) + { + Ok(value) => match serde_json::to_vec(&value) { + Ok(bytes) if bytes.len() <= max_payload_size => { + let mut response = Response::new(Body::from(bytes)); + *response.status_mut() = parts.status; + response + } + Ok(_) => error::bad_gateway( + "invalid_decisions_response", + "Converted Decisions response exceeds the payload limit", + ), + Err(err) => error::bad_gateway( + "invalid_decisions_response", + format!("Could not encode Decisions response: {err}"), + ), + }, + Err(message) => error::bad_gateway("invalid_decisions_response", message), + } + } else { + let value = parsed.ok(); + let message = value + .as_ref() + .and_then(|value| { + value + .pointer("/error/message") + .and_then(Value::as_str) + .or_else(|| value.get("detail").and_then(Value::as_str)) + }) + .map(str::to_owned) + .unwrap_or_else(|| { + value + .as_ref() + .and_then(|value| value.get("detail")) + .map(Value::to_string) + .unwrap_or_else(|| String::from_utf8_lossy(&bytes).into_owned()) + }); + let code = value + .as_ref() + .and_then(|value| value.pointer("/error/code")) + .and_then(Value::as_str) + .unwrap_or("upstream_decisions_error"); + let mut response = error::create_error(parts.status, code, message); + if let Some(value) = value.filter(|value| { + value + .pointer("/error/message") + .and_then(Value::as_str) + .is_some() + }) { + // Keep the backend's OpenAI error type, param and extensions. + // The helper response still supplies trusted metric extensions. + *response.body_mut() = Json(value).into_response().into_body(); + } + response + } + } + }; + // Rewriting the body invalidates lengths, encodings and entity validators. + // Keep routing/request identifiers, Retry-After, and other safe headers. + for name in [ + "content-length", + "content-encoding", + "transfer-encoding", + "etag", + "last-modified", + "content-md5", + "digest", + "content-digest", + "repr-digest", + "content-range", + "accept-ranges", + ] { + parts.headers.remove(name); + } + parts + .headers + .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + let (converted_parts, body) = converted.into_parts(); + parts.status = converted_parts.status; + parts.extensions.extend(converted_parts.extensions); + for (name, value) in &converted_parts.headers { + if name != CONTENT_TYPE && name != "content-length" { + parts.headers.insert(name.clone(), value.clone()); + } + } + Response::from_parts(parts, body) + } + + fn map_response(&self, value: &Value) -> Result { + let model = string_field(value, "model")?.to_owned(); + let upstream_answers = object_field(value, "answers")?; + if upstream_answers.len() != self.questions.len() { + return Err("SGLang Decisions returned an unexpected number of answers".into()); + } + let mut answers = Vec::with_capacity(self.questions.len()); + for (index, question) in self.questions.iter().enumerate() { + let id = format!("q{index}"); + let answer = upstream_answers + .get(&id) + .ok_or_else(|| format!("SGLang Decisions omitted answer {id}"))?; + let name = question.name().map(str::to_owned); + let kind = string_field(answer, "type")?; + let mapped = match question { + MappedQuestion::Predicate { .. } if kind == "noul" => DecisionAnswer::Predicate { + name, + probability: probability_field(answer, "noul")?, + }, + MappedQuestion::Choice { choices, .. } if kind == "choice" => { + let scores = object_field(answer, "probabilities")?; + if scores.len() != choices.len() { + return Err(format!( + "SGLang Decisions returned invalid probabilities for {id}" + )); + } + let chosen = string_field(answer, "choice")?; + let mut choice = None; + let mut chosen_probability = None; + let mut probabilities = Vec::with_capacity(choices.len()); + for (option_index, option) in choices.iter().enumerate() { + let key = format!("o{option_index}"); + if chosen == key { + choice = Some(option.clone()); + } + let probability = + probability_value(scores.get(&key).ok_or_else(|| { + format!("SGLang Decisions omitted probability {id}/{key}") + })?)?; + if chosen == key { + chosen_probability = Some(probability); + } + probabilities.push(DecisionChoiceProbability { + value: option.clone(), + probability, + }); + } + validate_probability_sum(probabilities.iter().map(|value| value.probability))?; + let top = probabilities + .iter() + .map(|value| value.probability) + .fold(0.0, f64::max); + if chosen_probability.is_some_and(|probability| probability < top) { + return Err(format!( + "SGLang Decisions returned a choice that is not most probable for {id}" + )); + } + DecisionAnswer::Choice { + name, + choice: choice.ok_or_else(|| { + format!("SGLang Decisions returned unknown choice {chosen} for {id}") + })?, + confidence: probability_field(answer, "confidence")?, + probabilities, + } + } + MappedQuestion::Score { labels, .. } if kind == "score" => { + let scores = object_field(answer, "probabilities")?; + if scores.len() != labels.len() { + return Err(format!( + "SGLang Decisions returned invalid probabilities for {id}" + )); + } + let score = number_field(answer, "score")?; + if !(0.0..=(labels.len() - 1) as f64).contains(&score) { + return Err(format!( + "SGLang Decisions returned an out-of-range score for {id}" + )); + } + let mut probabilities = Vec::with_capacity(labels.len()); + for (level_index, level) in labels.iter().enumerate() { + let key = level_index.to_string(); + probabilities.push(DecisionScoreProbability { + value: level_index as i64, + label: level.clone(), + probability: probability_value(scores.get(&key).ok_or_else(|| { + format!("SGLang Decisions omitted probability {id}/{key}") + })?)?, + }); + } + validate_probability_sum(probabilities.iter().map(|value| value.probability))?; + let expected: f64 = probabilities + .iter() + .map(|value| value.value as f64 * value.probability) + .sum(); + if (score - expected).abs() + > PROBABILITY_TOLERANCE * ((labels.len() - 1) as f64).max(1.0) + { + return Err(format!("SGLang Decisions returned a score inconsistent with its probabilities for {id}")); + } + DecisionAnswer::Score { + name, + score, + confidence: probability_field(answer, "confidence")?, + probabilities, + } + } + _ => { + return Err(format!( + "SGLang Decisions returned an unexpected answer type for {id}" + )) + } + }; + answers.push(mapped); + } + let usage = value.get("usage").ok_or("SGLang Decisions omitted usage")?; + let input_tokens = usage + .get("input_tokens") + .and_then(Value::as_u64) + .ok_or("SGLang Decisions returned invalid input_tokens")?; + let output_tokens = usage + .get("output_tokens") + .and_then(Value::as_u64) + .ok_or("SGLang Decisions returned invalid output_tokens")?; + if output_tokens != 0 { + return Err("SGLang Decisions scoring unexpectedly reported output tokens".into()); + } + let total_tokens = input_tokens + .checked_add(output_tokens) + .ok_or("SGLang Decisions token count overflow")?; + Ok(DecisionResponse { + model, + answers, + usage: DecisionUsage { + input_tokens, + output_tokens, + total_tokens, + // System One does not report cache counters. Zero denotes + // unreported cache usage, not a measured absence of cache activity. + input_tokens_details: DecisionInputTokensDetails { + cached_tokens: 0, + cache_write_tokens: 0, + }, + // Label scoring generates no reasoning or output tokens. + output_tokens_details: DecisionOutputTokensDetails { + reasoning_tokens: 0, + }, + }, + }) + } +} + +fn reject_extensions(other: &Map, context: &str) -> Result<(), String> { + if let Some(field) = other.keys().next() { + Err(format!( + "SGLang Decisions cannot honor field {field:?} on {context}" + )) + } else { + Ok(()) + } +} + +fn object_field<'a>(value: &'a Value, field: &str) -> Result<&'a Map, String> { + value + .get(field) + .and_then(Value::as_object) + .ok_or_else(|| format!("SGLang Decisions returned invalid or missing {field}")) +} + +fn string_field<'a>(value: &'a Value, field: &str) -> Result<&'a str, String> { + value + .get(field) + .and_then(Value::as_str) + .ok_or_else(|| format!("SGLang Decisions returned invalid or missing {field}")) +} + +fn number_field(value: &Value, field: &str) -> Result { + value + .get(field) + .and_then(Value::as_f64) + .filter(|number| number.is_finite()) + .ok_or_else(|| format!("SGLang Decisions returned invalid or missing {field}")) +} + +fn probability_field(value: &Value, field: &str) -> Result { + probability_value( + value + .get(field) + .ok_or_else(|| format!("SGLang Decisions omitted {field}"))?, + ) +} + +// Allow float32 softmax roundoff without accepting inconsistent distributions. +const PROBABILITY_TOLERANCE: f64 = 1e-6; + +fn validate_probability_sum(probabilities: impl Iterator) -> Result<(), String> { + if (probabilities.sum::() - 1.0).abs() > PROBABILITY_TOLERANCE { + Err("SGLang Decisions probabilities do not sum to one".into()) + } else { + Ok(()) + } +} + +fn probability_value(value: &Value) -> Result { + value + .as_f64() + .filter(|number| number.is_finite() && (0.0..=1.0).contains(number)) + .ok_or_else(|| "SGLang Decisions returned an invalid probability or confidence".into()) +} + +/// Reject duplicate JSON keys before they can silently replace an answer or +/// candidate probability. The normal Value deserializer retains only the last. +struct UniqueValue(Value); + +impl<'de> Deserialize<'de> for UniqueValue { + fn deserialize>(deserializer: D) -> Result { + struct Visitor; + impl<'de> de::Visitor<'de> for Visitor { + type Value = UniqueValue; + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("JSON without duplicate object keys") + } + fn visit_bool(self, value: bool) -> Result { + Ok(UniqueValue(Value::Bool(value))) + } + fn visit_i64(self, value: i64) -> Result { + Ok(UniqueValue(json!(value))) + } + fn visit_u64(self, value: u64) -> Result { + Ok(UniqueValue(json!(value))) + } + fn visit_f64(self, value: f64) -> Result { + serde_json::Number::from_f64(value) + .map(|value| UniqueValue(Value::Number(value))) + .ok_or_else(|| E::custom("non-finite number")) + } + fn visit_str(self, value: &str) -> Result { + Ok(UniqueValue(Value::String(value.to_owned()))) + } + fn visit_string(self, value: String) -> Result { + Ok(UniqueValue(Value::String(value))) + } + fn visit_unit(self) -> Result { + Ok(UniqueValue(Value::Null)) + } + fn visit_seq>( + self, + mut access: A, + ) -> Result { + let mut values = Vec::new(); + while let Some(value) = access.next_element::()? { + values.push(value.0); + } + Ok(UniqueValue(Value::Array(values))) + } + fn visit_map>( + self, + mut access: A, + ) -> Result { + let mut values = Map::new(); + while let Some(key) = access.next_key::()? { + if values.contains_key(&key) { + return Err(de::Error::custom(format!("duplicate key {key:?}"))); + } + values.insert(key, access.next_value::()?.0); + } + Ok(UniqueValue(Value::Object(values))) + } + } + deserializer.deserialize_any(Visitor) + } +} + +#[cfg(test)] +mod tests { + use axum::{body::to_bytes, http::StatusCode, response::IntoResponse, Json}; + use serde_json::json; + + use super::*; + + fn request(input: Value, questions: Value) -> DecisionsRequest { + let mut fields = Map::new(); + fields.insert("model".into(), json!("test")); + fields.insert("input".into(), input); + fields.insert("questions".into(), questions); + serde_json::from_value(Value::Object(fields)).unwrap() + } + + fn all_questions() -> Value { + json!([ + {"type":"predicate","instructions":"Is it ready?","name":"same"}, + {"type":"choice","instructions":"Select","choices":[{"value":true},{"value":"true","description":"Literal text"}],"name":"same"}, + {"type":"score","instructions":"Rate","levels":[{"label":"low"},{"label":"high","description":"Fully ready"}]} + ]) + } + + fn upstream() -> Value { + json!({"model":"served","answers":{ + "q2":{"type":"score","score":0.75,"confidence":0.5,"probabilities":{"0":0.25,"1":0.75},"legend":{"0":{"label":"low"},"1":{"label":"high","description":"Fully ready"}},"x_label_mass":0.9}, + "q1":{"type":"choice","choice":"o1","confidence":0.6,"probabilities":{"o1":0.8,"o0":0.2},"x_label_mass":0.9}, + "q0":{"type":"noul","noul":0.8,"x_label_mass":0.9} + },"usage":{"input_tokens":22,"output_tokens":0}}) + } + + async fn body(response: Response) -> Value { + serde_json::from_slice(&to_bytes(response.into_body(), 100_000).await.unwrap()).unwrap() + } + + #[test] + fn maps_typed_values_and_duplicate_names_to_private_ids() { + let adapter = + SglangDecisionAdapter::new(&request(json!("Evidence"), all_questions())).unwrap(); + assert_eq!( + adapter.request(), + &json!({"model":"test","state":"Evidence","questions":{ + "q0":{"type":"noul","instructions":"Is it ready?"}, + "q1":{"type":"choice","instructions":"Select","criteria":{"o0":{"value":true},"o1":{"value":"true","description":"Literal text"}}}, + "q2":{"type":"score","instructions":"Rate","criteria":[{"label":"low"},{"label":"high","description":"Fully ready"}]} + }}) + ); + } + + #[test] + fn structured_text_evidence_retains_message_and_part_order() { + let input = json!([ + {"role":"user","type":"message","content":[{"type":"input_text","text":"first"},{"type":"input_text","text":"second"}]}, + {"role":"user","content":"third"} + ]); + let adapter = SglangDecisionAdapter::new(&request(input.clone(), all_questions())).unwrap(); + assert_eq!(adapter.request()["state"], input); + } + + #[test] + fn rejects_images_and_extensions_instead_of_dropping_them() { + let image = json!([{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AQ=="}]}]); + let error = SglangDecisionAdapter::new(&request(image, all_questions())) + .err() + .unwrap(); + assert!(error.contains("image")); + for value in [ + json!({"model":"test","input":"x","questions":[{"type":"predicate","instructions":"x"}],"temperature":1}), + json!({"model":"test","input":[{"role":"user","content":"x","extra":true}],"questions":[{"type":"predicate","instructions":"x"}]}), + json!({"model":"test","input":[{"role":"user","content":[{"type":"input_text","text":"x","extra":true}]}],"questions":[{"type":"predicate","instructions":"x"}]}), + json!({"model":"test","input":"x","questions":[{"type":"predicate","instructions":"x","extra":true}]}), + json!({"model":"test","input":"x","questions":[{"type":"choice","instructions":"x","choices":[{"value":"a","extra":true}]}]}), + json!({"model":"test","input":"x","questions":[{"type":"score","instructions":"x","levels":[{"label":"a","extra":true}]}]}), + ] { + let request = serde_json::from_value(value).unwrap(); + assert!( + SglangDecisionAdapter::new(&request) + .err() + .unwrap() + .contains("extra") + || SglangDecisionAdapter::new(&request) + .err() + .unwrap() + .contains("temperature") + ); + } + } + + #[test] + fn enforces_sglang_bounds_and_accepts_singletons() { + for (kind, field, item, limit) in [ + ("choice", "choices", json!({"value":"a"}), 255), + ("score", "levels", json!({"label":"a"}), 10), + ] { + for size in [0, limit + 1] { + let mut question = json!({"type":kind,"instructions":"x"}); + question[field] = json!(vec![item.clone(); size]); + assert!( + SglangDecisionAdapter::new(&request(json!("x"), json!([question]))).is_err() + ); + } + let mut question = json!({"type":kind,"instructions":"x"}); + question[field] = json!([item]); + assert!(SglangDecisionAdapter::new(&request(json!("x"), json!([question]))).is_ok()); + } + assert!(SglangDecisionAdapter::new(&request(json!(" \n"), all_questions())).is_err()); + assert!(SglangDecisionAdapter::new(&request( + json!("x"), + json!([{ "type":"predicate","instructions":" " }]) + )) + .is_err()); + } + + #[tokio::test] + async fn converts_answers_in_request_order_with_typed_values_and_usage() { + let adapter = + SglangDecisionAdapter::new(&request(json!("Evidence"), all_questions())).unwrap(); + let response = adapter + .convert_response(Json(upstream()).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + body(response).await, + json!({"model":"served","answers":[ + {"type":"predicate","name":"same","probability":0.8}, + {"type":"choice","name":"same","choice":"true","confidence":0.6,"probabilities":[{"value":true,"probability":0.2},{"value":"true","probability":0.8}]}, + {"type":"score","name":null,"score":0.75,"confidence":0.5,"probabilities":[{"value":0,"label":"low","probability":0.25},{"value":1,"label":"high","probability":0.75}]} + ],"usage":{"input_tokens":22,"output_tokens":0,"total_tokens":22,"input_tokens_details":{"cached_tokens":0,"cache_write_tokens":0},"output_tokens_details":{"reasoning_tokens":0}}}) + ); + } + + #[tokio::test] + async fn preserves_headers_and_status_while_removing_body_validators() { + let adapter = + SglangDecisionAdapter::new(&request(json!("Evidence"), all_questions())).unwrap(); + let response = Response::builder() + .status(StatusCode::CREATED) + .header("x-request-id", "request-42") + .header("x-smg-routed-worker", "worker") + .header("content-length", "999") + .header("content-encoding", "gzip") + .header("etag", "old") + .header("content-md5", "old") + .body(Body::from(upstream().to_string())) + .unwrap(); + let response = adapter.convert_response(response, 100_000).await; + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()["x-request-id"], "request-42"); + assert_eq!(response.headers()["x-smg-routed-worker"], "worker"); + for header in ["content-length", "content-encoding", "etag", "content-md5"] { + assert!(!response.headers().contains_key(header)); + } + assert_eq!(response.headers()["content-type"], "application/json"); + } + + #[tokio::test] + async fn upstream_validation_errors_remain_errors_with_useful_messages() { + for fixture in [ + json!({"error":{"message":"unsupported reasoning template","type":"invalid_request_error","code":"invalid_model"}}), + json!({"detail":"question q0 label is not a single token"}), + ] { + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response( + (StatusCode::BAD_REQUEST, Json(fixture)).into_response(), + 100_000, + ) + .await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let value = body(response).await; + assert!( + value["error"]["message"] + .as_str() + .unwrap() + .contains("unsupported") + || value["error"]["message"] + .as_str() + .unwrap() + .contains("single token") + ); + assert!(value.get("answers").is_none()); + } + } + + #[tokio::test] + async fn malformed_successes_fail_as_bad_gateway() { + let mutations = [ + ("/answers/q0/noul", json!(1.1)), + ("/answers/q0/type", json!("predicate")), + ("/answers/q1/choice", json!("unknown")), + ("/answers/q1/confidence", json!(-0.1)), + ("/answers/q1/probabilities/o0", Value::Null), + ("/answers/q1/probabilities", json!({"o0":0.2,"other":0.8})), + ("/answers/q2/score", json!(2)), + ("/answers/q2/score", json!("NaN")), + ("/answers/q2/probabilities", json!({"0":0.2})), + ("/answers/q0", Value::Null), + ("/usage/input_tokens", json!(-1)), + ("/usage/output_tokens", json!(1)), + ("/model", Value::Null), + ]; + for (path, replacement) in mutations { + let mut value = upstream(); + *value.pointer_mut(path).unwrap() = replacement; + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Json(value).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY, "{path}"); + assert!(body(response).await["error"]["message"].is_string()); + } + for text in [ + "{", + "null", + r#"{"model":"served","answers":{"q0":{"type":"noul","noul":0.2,"noul":0.9}},"usage":{"input_tokens":1,"output_tokens":0}}"#, + ] { + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Response::new(Body::from(text)), 100_000) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + } + + #[tokio::test] + async fn conversion_enforces_the_response_buffer_limit() { + let adapter = SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Json(upstream()).into_response(), 10) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn duplicate_backend_keys_cannot_overwrite_valid_answer_probabilities() { + for text in [ + upstream() + .to_string() + .replace("\"noul\":0.8", "\"noul\":0.2,\"noul\":0.8"), + upstream() + .to_string() + .replace("\"o0\":0.2", "\"o0\":0.4,\"o0\":0.2"), + ] { + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Response::new(Body::from(text)), 100_000) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert!(body(response).await["error"]["message"] + .as_str() + .unwrap() + .contains("duplicate key")); + } + } + + #[tokio::test] + async fn singleton_answers_preserve_original_strings_and_score_labels() { + let questions = json!([ + {"type":"choice","instructions":"Select","choices":[{"value":"\nA: true"}]}, + {"type":"score","instructions":"Rate","levels":[{"label":"A: low"}]} + ]); + let adapter = SglangDecisionAdapter::new(&request(json!("x"), questions)).unwrap(); + assert_eq!( + adapter.request()["questions"]["q0"]["criteria"], + json!({"o0":{"value":"\nA: true"}}) + ); + let fixture = json!({"model":"served","answers":{ + "q0":{"type":"choice","choice":"o0","confidence":1,"probabilities":{"o0":1},"x_label_mass":0.6}, + "q1":{"type":"score","score":0,"confidence":1,"probabilities":{"0":1},"legend":{"0":{"label":"A: low"}},"x_label_mass":0.7} + },"usage":{"input_tokens":1,"output_tokens":0}}); + let response = adapter + .convert_response(Json(fixture).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + body(response).await["answers"], + json!([ + {"type":"choice","name":null,"choice":"\nA: true","confidence":1.0,"probabilities":[{"value":"\nA: true","probability":1.0}]}, + {"type":"score","name":null,"score":0.0,"confidence":1.0,"probabilities":[{"value":0,"label":"A: low","probability":1.0}]} + ]) + ); + } + + #[test] + fn gateway_safety_identifier_is_not_forwarded_as_a_backend_extension() { + let mut value = json!({"model":"test","input":"x","questions":[{"type":"predicate","instructions":"x"}],"safety_identifier":"user-42"}); + let request = serde_json::from_value(value.take()).unwrap(); + let adapter = SglangDecisionAdapter::new(&request).unwrap(); + assert_eq!( + adapter.request(), + &json!({"model":"test","state":"x","questions":{"q0":{"type":"noul","instructions":"x"}}}) + ); + } + + #[tokio::test] + async fn rejects_inconsistent_distributions_scores_and_choice_winners() { + for (path, replacement) in [ + ("/answers/q1/probabilities", json!({"o0":0.8,"o1":0.8})), + ("/answers/q2/probabilities", json!({"0":0.5,"1":0.75})), + ("/answers/q2/score", json!(0.1)), + ("/answers/q1/choice", json!("o0")), + ] { + let mut fixture = upstream(); + *fixture.pointer_mut(path).unwrap() = replacement; + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Json(fixture).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY, "{path}"); + } + } + + #[tokio::test] + async fn accepts_tied_winners_and_probability_roundoff() { + for winner in ["o0", "o1"] { + let mut fixture = upstream(); + fixture["answers"]["q1"]["probabilities"] = json!({"o0":0.5,"o1":0.5}); + fixture["answers"]["q1"]["choice"] = json!(winner); + fixture["answers"]["q2"]["probabilities"] = json!({"0":0.25,"1":0.75000001}); + let adapter = + SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response(Json(fixture).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::OK); + } + } + + #[tokio::test] + async fn preserves_existing_openai_error_envelopes() { + let fixture = json!({"error":{"message":"unsupported template","type":"invalid_request_error","code":"unsupported_model","param":"model","backend_context":{"retry":false}},"trace":"abc"}); + let adapter = SglangDecisionAdapter::new(&request(json!("x"), all_questions())).unwrap(); + let response = adapter + .convert_response( + (StatusCode::BAD_REQUEST, Json(fixture.clone())).into_response(), + 100_000, + ) + .await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(body(response).await, fixture); + } + + #[tokio::test] + async fn bounds_the_encoded_public_response_including_original_labels() { + let questions = json!([{ "type":"choice", "instructions":"Select", "choices":[{"value":"x".repeat(2000)}]}]); + let fixture = json!({"model":"served","answers":{"q0":{"type":"choice","choice":"o0","confidence":1,"probabilities":{"o0":1},"x_label_mass":0.5}},"usage":{"input_tokens":1,"output_tokens":0}}); + let limit = fixture.to_string().len() + 20; + let adapter = SglangDecisionAdapter::new(&request(json!("x"), questions)).unwrap(); + let response = adapter + .convert_response(Json(fixture).into_response(), limit) + .await; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn releasing_dispatch_evidence_retains_answer_mapping() { + let mut adapter = + SglangDecisionAdapter::new(&request(json!("Evidence"), all_questions())).unwrap(); + adapter.release_request(); + assert!(adapter.request().is_null()); + let response = adapter + .convert_response(Json(upstream()).into_response(), 100_000) + .await; + assert_eq!(response.status(), StatusCode::OK); + let answers = body(response).await["answers"].clone(); + assert_eq!(answers[0]["name"], "same"); + assert_eq!(answers[1]["choice"], "true"); + assert_eq!(answers[2]["probabilities"][1]["label"], "high"); + } +} diff --git a/model_gateway/src/routers/http/mod.rs b/model_gateway/src/routers/http/mod.rs index b1effddd21..0e085c652c 100644 --- a/model_gateway/src/routers/http/mod.rs +++ b/model_gateway/src/routers/http/mod.rs @@ -1,5 +1,6 @@ //! HTTP router implementations +pub(crate) mod decisions; pub mod pd_router; pub mod pd_types; pub(crate) mod request_body; diff --git a/model_gateway/src/routers/http/request_body.rs b/model_gateway/src/routers/http/request_body.rs index 549b62aba2..c772c48346 100644 --- a/model_gateway/src/routers/http/request_body.rs +++ b/model_gateway/src/routers/http/request_body.rs @@ -43,13 +43,34 @@ pub(crate) fn serialize_request_body( canonical_model: Option<&str>, worker: &dyn Worker, raw_len: Option, +) -> Result, RequestBodyError> { + serialize_request_body_impl(typed_req, canonical_model, worker, raw_len, true) +} + +/// Preserve explicitly supplied backend fields while applying model aliases +/// and worker preparation. Decisions extensions have no shared defaults. +pub(crate) fn serialize_request_body_preserving_fields( + typed_req: &T, + canonical_model: Option<&str>, + worker: &dyn Worker, + raw_len: Option, +) -> Result, RequestBodyError> { + serialize_request_body_impl(typed_req, canonical_model, worker, raw_len, false) +} + +fn serialize_request_body_impl( + typed_req: &T, + canonical_model: Option<&str>, + worker: &dyn Worker, + raw_len: Option, + strip_defaults: bool, ) -> Result, RequestBodyError> { let bytes = serialize_json_sized(typed_req, raw_len).map_err(RequestBodyError::Serialize)?; // A body the raw editor cannot parse (in practice: non-objects) takes // the Value pipeline rather than skipping the hooks. let Ok(mut body) = RawBody::parse(&bytes) else { - return value_request_body(typed_req, canonical_model, worker, raw_len); + return value_request_body(typed_req, canonical_model, worker, raw_len, strip_defaults); }; // Edits only swap the model id, append the DP rank, or strip fields, so // the first pass plus the two inserted values bounds the reserialization. @@ -68,13 +89,23 @@ pub(crate) fn serialize_request_body( if worker.mutates_request() { let rank = match worker.dp_rank() { Some(rank) if worker.uses_builtin_prepare_request() => rank, - _ => return value_request_body(typed_req, canonical_model, worker, raw_len), + _ => { + return value_request_body( + typed_req, + canonical_model, + worker, + raw_len, + strip_defaults, + ) + } }; let rank = to_raw_value(&rank).map_err(RequestBodyError::Serialize)?; extra += DATA_PARALLEL_RANK.len() + rank.get().len() + 4; body.insert(DATA_PARALLEL_RANK, rank); } - body.strip_default_sglang_fields(); + if strip_defaults { + body.strip_default_sglang_fields(); + } if body.mutated() { let mut out = Vec::with_capacity(bytes.len() + extra); serde_json::to_writer(&mut out, &body).map_err(RequestBodyError::Serialize)?; @@ -91,6 +122,7 @@ fn value_request_body( canonical_model: Option<&str>, worker: &dyn Worker, raw_len: Option, + strip_defaults: bool, ) -> Result, RequestBodyError> { let mut json_val = request_to_value(typed_req, raw_len).map_err(RequestBodyError::Serialize)?; if let Some(canonical_model) = canonical_model { @@ -99,7 +131,9 @@ fn value_request_body( let mut json_val = worker .prepare_request(json_val) .map_err(RequestBodyError::Prepare)?; - strip_default_sglang_fields(&mut json_val); + if strip_defaults { + strip_default_sglang_fields(&mut json_val); + } serialize_json_sized(&json_val, raw_len).map_err(RequestBodyError::Serialize) } @@ -777,4 +811,63 @@ mod tests { "and so did the built-in one" ); } + + #[test] + fn preserving_fields_retains_extensions_with_alias_and_dp_rank() { + let req = json!({ + "model": "alias-model", + "skip_special_tokens": false, + "separate_reasoning": true, + "chat_template_kwargs": null, + "native": {"keep": [false, null]} + }); + let body = serialize_request_body_preserving_fields( + &req, + Some("canonical-model"), + &dp_worker(), + None, + ) + .unwrap(); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + parsed, + json!({ + "model": "canonical-model", + "skip_special_tokens": false, + "separate_reasoning": true, + "chat_template_kwargs": null, + "native": {"keep": [false, null]}, + "data_parallel_rank": 3 + }) + ); + } + + #[test] + fn preserving_fields_retains_extensions_through_custom_prepare() { + let req = json!({ + "model": "alias-model", + "skip_special_tokens": false, + "separate_reasoning": true, + "chat_template_kwargs": null + }); + let body = serialize_request_body_preserving_fields( + &req, + Some("canonical-model"), + &CustomPrepare(dp_worker()), + None, + ) + .unwrap(); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + parsed, + json!({ + "model": "canonical-model", + "skip_special_tokens": false, + "separate_reasoning": true, + "chat_template_kwargs": null, + "custom": true, + "data_parallel_rank": 3 + }) + ); + } } diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index f5d4835be6..21c733fd8d 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -22,6 +22,7 @@ use openai_protocol::{ classify::ClassifyRequest, common::GenerationRequest, completion::CompletionRequest, + decisions::DecisionsRequest, embedding::EmbeddingRequest, generate::GenerateRequest, messages::{CountMessageTokensRequest, CreateMessageRequest}, @@ -74,15 +75,18 @@ use crate::{ gateway::Gateway, grpc::utils::{error_type_from_status, route_to_endpoint}, http::{ - request_body::{serialize_request_body, RequestBodyError}, + decisions::{SglangDecisionAdapter, UPSTREAM_ROUTE}, + request_body::{ + serialize_request_body, serialize_request_body_preserving_fields, RequestBodyError, + }, request_stream::{CappedBodyStream, StreamProgress}, }, BodyPolicy, RouterTrait, }, wasm::module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, worker::{ - AttachedBody, ConnectionMode, RoutingPool, Worker, WorkerLoadGuard, WorkerRegistry, - WorkerType, + AttachedBody, ConnectionMode, RoutingPool, RuntimeType, Worker, WorkerLoadGuard, + WorkerRegistry, WorkerType, }, }; @@ -551,6 +555,36 @@ impl Router { return shed; } + // Select the adapter for the chosen worker. Native Decisions workers + // keep the original typed serialization path, including extensions. + let mut decision_adapter = if route == "/v1/decisions" + && worker.metadata().spec.runtime_type == RuntimeType::Sglang + { + let adapter = lease.with_view(|view| { + let request = serde_json::to_value(view.request) + .and_then(serde_json::from_value::) + .map_err(|error| { + error::bad_request( + "serialization_failed", + format!("Failed to serialize Decisions request: {error}"), + ) + })?; + SglangDecisionAdapter::new(&request) + .map_err(|message| error::bad_request("unsupported_decisions_request", message)) + }); + match adapter { + Ok(adapter) => Some(adapter), + Err(response) => return response, + } + } else { + None + }; + let upstream_route = if decision_adapter.is_some() { + UPSTREAM_ROUTE + } else { + route + }; + let load_guard = lease.with_view(|view| { WorkerLoadGuard::with_key( worker.clone(), @@ -575,18 +609,37 @@ impl Router { }) }); let response = match lease.serialize_with(|view| { - serialize_request_body(view.request, canonical_model, worker.as_ref(), raw_body_len) + if let Some(adapter) = &decision_adapter { + serialize_request_body( + adapter.request(), + canonical_model, + worker.as_ref(), + raw_body_len, + ) + } else if route == "/v1/decisions" { + serialize_request_body_preserving_fields( + view.request, + canonical_model, + worker.as_ref(), + raw_body_len, + ) + } else { + serialize_request_body(view.request, canonical_model, worker.as_ref(), raw_body_len) + } }) { Ok(body) => { // Past this point dispatch needs only the serialized bytes; // the lease frees the parsed request and its routing // derivatives now when retries are disabled. lease.release_dispatch(); + if let Some(adapter) = decision_adapter.as_mut() { + adapter.release_request(); + } let mode = StreamRelayMode { is_stream, rechunk }; self.send_serialized_request( headers, body, - route, + upstream_route, worker.as_ref(), mode, load_guard, @@ -603,6 +656,16 @@ impl Router { ), }; + // Judge the converted result so malformed successes become backend + // failures and follow the same outcome accounting and retry policy. + let response = if let Some(adapter) = decision_adapter { + adapter + .convert_response(response, self.max_payload_size) + .await + } else { + response + }; + events::RequestReceivedEvent {}.emit(); let status = response.status(); @@ -2054,6 +2117,17 @@ impl RouterTrait for Router { .await } + async fn route_decisions( + &self, + headers: Option<&HeaderMap>, + _tenant_meta: &TenantRequestMeta, + body: DecisionsRequest, + model_id: &str, + ) -> Response { + self.route_typed_request(headers, body, "/v1/decisions", model_id) + .await + } + async fn route_audio_transcriptions( &self, headers: Option<&HeaderMap>, diff --git a/model_gateway/src/routers/mod.rs b/model_gateway/src/routers/mod.rs index 8a0f33abb9..5bfcbc2079 100644 --- a/model_gateway/src/routers/mod.rs +++ b/model_gateway/src/routers/mod.rs @@ -13,6 +13,7 @@ use openai_protocol::{ chat::ChatCompletionRequest, classify::ClassifyRequest, completion::CompletionRequest, + decisions::DecisionsRequest, embedding::EmbeddingRequest, generate::GenerateRequest, interactions::InteractionsRequest, @@ -192,6 +193,20 @@ pub trait RouterTrait: Send + Sync + Debug { (StatusCode::NOT_IMPLEMENTED, "Classify not implemented").into_response() } + /// Route OpenAI Decisions requests on regular HTTP workers. + async fn route_decisions( + &self, + _headers: Option<&HeaderMap>, + _tenant_meta: &TenantRequestMeta, + _body: DecisionsRequest, + _model_id: &str, + ) -> Response { + error::not_implemented( + "decisions_not_supported", + "Decisions is supported only by the regular HTTP router", + ) + } + /// Route audio transcription requests (OpenAI-compatible /v1/audio/transcriptions). /// /// Unlike the JSON-bodied endpoints, `/v1/audio/transcriptions` uses diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index bb1668cd24..3f65ae4731 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -18,6 +18,7 @@ use openai_protocol::{ chat::ChatCompletionRequest, classify::ClassifyRequest, completion::CompletionRequest, + decisions::DecisionsRequest, embedding::EmbeddingRequest, generate::GenerateRequest, interactions::InteractionsRequest, @@ -266,6 +267,23 @@ async fn v1_interactions( .await } +async fn v1_decisions( + State(state): State>, + headers: HeaderMap, + Extension(tenant_meta): Extension, + cancel: middleware::scheduler::PreemptionGuard, + ValidatedJson(body): ValidatedJson, +) -> Response { + let model = body.model.clone(); + cancel + .guard( + state + .router + .route_decisions(Some(&headers), &tenant_meta, body, &model), + ) + .await +} + async fn v1_embeddings( State(state): State>, headers: HeaderMap, @@ -879,6 +897,7 @@ pub fn build_app( .route("/v1/messages/count_tokens", post(v1_messages_count_tokens)) .route("/v1/interactions", post(v1_interactions)) .route("/v1/classify", post(v1_classify)) + .route("/v1/decisions", post(v1_decisions)) // Per-request buffer-vs-stream decision for typed-JSON bodies; // declined requests pass to the handlers untouched. .route_layer(axum::middleware::from_fn_with_state( diff --git a/model_gateway/tests/decisions_http_test.rs b/model_gateway/tests/decisions_http_test.rs new file mode 100644 index 0000000000..f8a2fdf374 --- /dev/null +++ b/model_gateway/tests/decisions_http_test.rs @@ -0,0 +1,321 @@ +//! Decisions contract at the production server/Gateway/HTTP-router boundary. +#![expect(clippy::unwrap_used, reason = "test failures should panic")] + +mod common; + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use axum::{ + body::{to_bytes, Body}, + extract::{Request, State}, + http::{ + header::{CONTENT_LENGTH, CONTENT_TYPE}, + HeaderMap, StatusCode, + }, + response::Response, + routing::post, + Router, +}; +use common::test_app::{create_test_app_context, create_test_app_with_context}; +use serde_json::{json, Value}; +use smg::{ + routers::{ + factory::router_ids, gateway::Gateway, http::router::Router as HttpRouter, RouterTrait, + }, + worker::{BasicWorkerBuilder, ModelCard, RuntimeType}, +}; +use tokio::{net::TcpListener, task::JoinHandle}; +use tower::ServiceExt; + +const MODEL: &str = "decision-model"; +const ALIAS: &str = "decision-alias"; + +#[derive(Clone)] +struct UpstreamState { + requests: Arc>>, + status: StatusCode, + body: String, +} +struct Upstream { + url: String, + state: UpstreamState, + task: JoinHandle<()>, +} +impl Drop for Upstream { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn record_request(State(state): State, req: Request) -> Response { + let (parts, body) = req.into_parts(); + let body = to_bytes(body, 1024 * 1024).await.unwrap(); + state.requests.lock().unwrap().push(( + parts.uri.path().to_string(), + parts.headers, + serde_json::from_slice(&body).unwrap(), + )); + Response::builder() + .status(state.status) + .header(CONTENT_TYPE, "application/json") + .header(CONTENT_LENGTH, state.body.len()) + .header("x-request-id", "upstream-decision-id") + .body(Body::from(state.body)) + .unwrap() +} +#[expect( + clippy::disallowed_methods, + reason = "test server is aborted by the owning Upstream on drop" +)] +async fn upstream(status: StatusCode, body: Value) -> Upstream { + let state = UpstreamState { + requests: Arc::new(Mutex::new(Vec::new())), + status, + body: body.to_string(), + }; + let app = Router::new() + .route("/v1/decisions", post(record_request)) + .route("/v1/systemone", post(record_request)) + .with_state(state.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + Upstream { url, state, task } +} +async fn app_for(upstream: &Upstream, runtime: RuntimeType) -> Router { + let ctx = create_test_app_context().await; + ctx.worker_registry + .register(Arc::new( + BasicWorkerBuilder::new(&upstream.url) + .runtime_type(runtime) + .model(ModelCard::new(MODEL).with_alias(ALIAS)) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + )) + .unwrap(); + let router = Arc::new(HttpRouter::new(&ctx).await.unwrap()); + let gateway = Arc::new(Gateway::new(ctx.worker_registry.clone())); + gateway.register_router(router_ids::HTTP_REGULAR, router); + create_test_app_with_context(gateway, ctx) +} +fn request(body: Value) -> Request { + Request::builder() + .method("POST") + .uri("/v1/decisions") + .header(CONTENT_TYPE, "application/json") + .header("x-request-id", "client-decision-id") + .body(Body::from(body.to_string())) + .unwrap() +} +async fn json_body(response: Response) -> Value { + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() +} +fn predicate_request() -> Value { + json!({"model": MODEL, "input": "The sky is blue.", "questions": [{"type": "predicate", "instructions": "Is the sky blue?"}]}) +} +fn native_response() -> Value { + json!({"model": MODEL, "answers": [{"type": "predicate", "name": null, "probability": 0.9, "native_answer": true}], + "usage": {"input_tokens": 7, "output_tokens": 0, "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 2, "cache_write_tokens": 3}, + "output_tokens_details": {"reasoning_tokens": 0}}, "native_response": {"keep": 1}}) +} +fn sglang_response() -> Value { + // Intentionally not in request order: correlation must use qN, not map order/name. + json!({"model": "served-model", "answers": { + "q2": {"type": "score", "score": 0.75, "confidence": 0.5, + "legend": {"0": {"label": "low"}, "1": {"label": "high", "description": "excellent"}}, + "probabilities": {"0": 0.25, "1": 0.75}, "x_label_mass": 0.8}, + "q0": {"type": "noul", "noul": 0.8, "x_label_mass": 0.4}, + "q1": {"type": "choice", "choice": "o1", "confidence": 0.6, + "probabilities": {"o0": 0.2, "o1": 0.8}, "x_label_mass": 0.7} + }, "usage": {"input_tokens": 22, "output_tokens": 0}}) +} +fn all_questions_request() -> Value { + json!({"model": ALIAS, "input": [{"role": "user", "content": "first evidence"}, + {"role": "user", "content": [{"type": "input_text", "text": "second evidence"}]}], + "questions": [ + {"type": "predicate", "instructions": "Is this positive?"}, + {"type": "choice", "name": "same", "instructions": "Choose", "choices": [ + {"value": true}, {"value": "true", "description": "string"}]}, + {"type": "score", "name": "same", "instructions": "Score", "levels": [ + {"label": "low"}, {"label": "high", "description": "excellent"}]} + ]}) +} + +#[tokio::test] +async fn native_roundtrip_preserves_extensions_and_resolves_model_alias() { + let expected = native_response(); + let upstream = upstream(StatusCode::OK, expected.clone()).await; + let app = app_for(&upstream, RuntimeType::Generic).await; + let mut body = predicate_request(); + body["model"] = json!(ALIAS); + body["input"] = json!([{"role": "user", "type": "message", "native_message": 1, + "content": [ + {"type": "input_text", "text": "Inspect", "native_part": true}, + {"type": "input_image", "image_url": "data:image/png;base64,AA==", "detail": "original"}, + {"type": "input_text", "text": "carefully"} + ]}]); + body["safety_identifier"] = json!("opaque-user"); + body["native_request"] = json!({"keep": [1, true]}); + body["skip_special_tokens"] = json!(false); + body["separate_reasoning"] = json!(true); + body["chat_template_kwargs"] = Value::Null; + body["questions"][0]["native_question"] = json!("preserved"); + let response = app.oneshot(request(body.clone())).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(json_body(response).await, expected); + let requests = upstream.state.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, "/v1/decisions"); + body["model"] = json!(MODEL); + assert_eq!(requests[0].2, body); + assert_eq!(requests[0].1["x-request-id"], "client-decision-id"); +} + +#[tokio::test] +async fn sglang_translates_questions_and_restores_order_names_and_typed_values() { + let upstream = upstream(StatusCode::OK, sglang_response()).await; + let app = app_for(&upstream, RuntimeType::Sglang).await; + let body = all_questions_request(); + let response = app.oneshot(request(body.clone())).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()[CONTENT_TYPE], "application/json"); + // Middleware may replace x-request-id with the incoming ID; it must survive conversion. + assert!(response.headers().contains_key("x-request-id")); + let length = response + .headers() + .get(CONTENT_LENGTH) + .map(|v| v.to_str().unwrap().parse::().unwrap()); + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + if let Some(length) = length { + assert_eq!(length, bytes.len(), "stale upstream content length"); + } + let result: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!( + result["answers"], + json!([ + {"type": "predicate", "name": null, "probability": 0.8}, + {"type": "choice", "name": "same", "choice": "true", "confidence": 0.6, + "probabilities": [{"value": true, "probability": 0.2}, {"value": "true", "probability": 0.8}]}, + {"type": "score", "name": "same", "score": 0.75, "confidence": 0.5, + "probabilities": [{"value": 0, "label": "low", "probability": 0.25}, {"value": 1, "label": "high", "probability": 0.75}]} + ]) + ); + assert_eq!( + result["usage"], + json!({"input_tokens": 22, "output_tokens": 0, "total_tokens": 22, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, "output_tokens_details": {"reasoning_tokens": 0}}) + ); + let requests = upstream.state.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, "/v1/systemone"); + assert_eq!( + requests[0].2, + json!({"model": MODEL, "state": body["input"], "questions": { + "q0": {"type": "noul", "instructions": "Is this positive?"}, + "q1": {"type": "choice", "instructions": "Choose", "criteria": {"o0": {"value": true}, "o1": {"value": "true", "description": "string"}}}, + "q2": {"type": "score", "instructions": "Score", "criteria": [{"label": "low"}, {"label": "high", "description": "excellent"}]} + }}) + ); +} + +#[tokio::test] +async fn missing_model_is_rejected_before_dispatch() { + let upstream = upstream(StatusCode::OK, native_response()).await; + let app = app_for(&upstream, RuntimeType::Generic).await; + let mut body = predicate_request(); + body.as_object_mut().unwrap().remove("model"); + let response = app.oneshot(request(body)).await.unwrap(); + assert!(response.status().is_client_error()); + assert!(upstream.state.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn sglang_image_is_rejected_without_dispatch() { + let upstream = upstream(StatusCode::OK, json!({})).await; + let app = app_for(&upstream, RuntimeType::Sglang).await; + let mut body = predicate_request(); + body["input"] = json!([{"role": "user", "content": [{"type": "input_text", "text": "Inspect this"}, + {"type": "input_image", "image_url": "data:image/png;base64,AA=="}]}]); + let response = app.oneshot(request(body)).await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert!(json_body(response).await["error"]["message"].is_string()); + assert!(upstream.state.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn sglang_backend_errors_pass_through_without_success_conversion() { + let expected = + json!({"error": {"message": "backend validation failed", "type": "invalid_request_error"}}); + let upstream = upstream(StatusCode::UNPROCESSABLE_ENTITY, expected.clone()).await; + let app = app_for(&upstream, RuntimeType::Sglang).await; + let response = app.oneshot(request(predicate_request())).await.unwrap(); + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!(json_body(response).await, expected); +} + +#[tokio::test] +async fn sglang_malformed_success_is_a_gateway_error() { + let upstream = upstream(StatusCode::OK, json!({"not_a_decision": true})).await; + let app = app_for(&upstream, RuntimeType::Sglang).await; + let response = app.oneshot(request(predicate_request())).await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert!(json_body(response).await["error"]["message"].is_string()); +} + +#[derive(Debug)] +struct UnsupportedRouter(&'static str); +#[async_trait] +impl RouterTrait for UnsupportedRouter { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn router_type(&self) -> &'static str { + self.0 + } +} +#[tokio::test] +async fn unsupported_router_transports_return_not_implemented() { + for family in ["pd", "grpc", "grpc_pd"] { + let ctx = create_test_app_context().await; + let app = create_test_app_with_context(Arc::new(UnsupportedRouter(family)), ctx); + let response = app.oneshot(request(predicate_request())).await.unwrap(); + assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED, "{family}"); + } +} + +#[tokio::test] +async fn sglang_unsupported_extensions_are_rejected_instead_of_dropped() { + let upstream = upstream(StatusCode::OK, sglang_response()).await; + let app = app_for(&upstream, RuntimeType::Sglang).await; + for pointer in [ + "", + "/input/0", + "/questions/0", + "/questions/1/choices/0", + "/questions/2/levels/0", + ] { + let mut body = all_questions_request(); + body.pointer_mut(pointer) + .unwrap() + .as_object_mut() + .unwrap() + .insert( + "unsupported_extension".into(), + json!({"changes_semantics": true}), + ); + let response = app.clone().oneshot(request(body)).await.unwrap(); + assert_eq!( + response.status(), + StatusCode::BAD_REQUEST, + "extension at {pointer}" + ); + assert!(json_body(response).await["error"]["message"].is_string()); + } + assert!(upstream.state.requests.lock().unwrap().is_empty()); +}