diff --git a/Cargo.lock b/Cargo.lock index 1d812f47..76284a52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -287,6 +287,7 @@ version = "0.3.0" dependencies = [ "aisix-core", "aisix-gateway", + "aisix-provider-anthropic", "aisix-provider-openai", "async-stream", "async-trait", @@ -332,6 +333,7 @@ version = "0.3.0" dependencies = [ "aisix-core", "aisix-gateway", + "aisix-provider-anthropic", "async-stream", "async-trait", "bytes", diff --git a/crates/aisix-gateway/src/bridge.rs b/crates/aisix-gateway/src/bridge.rs index 3a21f102..9297f935 100644 --- a/crates/aisix-gateway/src/bridge.rs +++ b/crates/aisix-gateway/src/bridge.rs @@ -100,7 +100,25 @@ pub struct BridgeContext { pub provider_key: std::sync::Arc, /// Deadline for the entire upstream call. Bridges are expected to /// honour this by cancelling any in-flight HTTP request. + /// + /// On a streaming dispatch this is the **streaming** budget, which + /// bounds the connect phase and the gap between chunks rather than + /// the whole completion. A bridge that answers a streaming request + /// with a non-streaming upstream leg must use + /// [`non_streaming_deadline`](Self::non_streaming_deadline) instead. pub deadline: Option, + /// The end-to-end budget for a non-streaming upstream call, carried + /// alongside `deadline` on streaming dispatches. + /// + /// A structured-output request on the synthetic-tool route cannot be + /// streamed — the JSON only exists once the tool call is complete — + /// so those bridges run the upstream leg non-streaming and render + /// the result as chunks. Measured against the streaming budget, a + /// completion that takes longer than one chunk gap is supposed to + /// would be cut off; this is the budget that call is actually + /// entitled to. `None` on a non-streaming dispatch, where `deadline` + /// already is it. + pub non_streaming_deadline: Option, /// The authenticated caller, for `${request.api_key.*}` header /// templates. Default (all-empty) on calls with no caller behind /// them — a background job poll, an internal embedding lookup. @@ -160,6 +178,7 @@ impl BridgeContext { model, provider_key, deadline: None, + non_streaming_deadline: None, caller: CallerIdentity::default(), client_headers: None, model_id: String::new(), @@ -172,6 +191,25 @@ impl BridgeContext { self } + /// Record the end-to-end budget a non-streaming call would have got, + /// for the streaming dispatches whose `deadline` is the smaller + /// streaming budget. See + /// [`non_streaming_deadline`](Self::non_streaming_deadline). + pub fn with_non_streaming_deadline(mut self, deadline: Option) -> Self { + self.non_streaming_deadline = deadline; + self + } + + /// The deadline an upstream leg that is *not* streaming should run + /// under, whichever kind of dispatch this context came from. + pub fn non_streaming_ctx(&self) -> Self { + let mut ctx = self.clone(); + if let Some(deadline) = self.non_streaming_deadline { + ctx.deadline = Some(deadline); + } + ctx + } + /// Attach the caller identity and inbound headers the outbound-header /// pipeline reads. Dispatch paths with a real client request call this; /// leaving it off means no client header is ever forwarded and diff --git a/crates/aisix-gateway/src/lib.rs b/crates/aisix-gateway/src/lib.rs index 2e012a31..145485a0 100644 --- a/crates/aisix-gateway/src/lib.rs +++ b/crates/aisix-gateway/src/lib.rs @@ -12,6 +12,10 @@ //! that dispatches `ChatFormat` to the right `Bridge`. //! - [`sse`] — a provider-agnostic SSE line decoder. Bridges that stream //! over SSE feed it raw bytes and pull typed events back out. +//! - [`structured_output`] — the `response_format` translation pieces +//! every bridge shares: the synthetic JSON tool and its reverse +//! translation, the fake stream that carries it, and the two schema +//! normalisations. //! - [`credential`] — cache keys for credential-derived upstream tokens. //! - [`upstream_http`] — connection-layer settings every provider client //! shares (connect timeout, TCP keepalive, pool expiry) plus the @@ -28,6 +32,7 @@ pub mod chat; pub mod credential; pub mod hub; pub mod sse; +pub mod structured_output; pub mod upstream_headers; pub mod upstream_http; pub mod upstream_tls; @@ -46,6 +51,11 @@ pub use chat::{ pub use credential::credential_fingerprint; pub use hub::{upstream_protocol, Hub, UPSTREAM_PROTOCOL_UNKNOWN}; pub use sse::{SseDecoder, SseEvent}; +pub use structured_output::{ + apply_schema_limits, close_object_schemas, json_schema_from_response_format, + response_into_fake_stream_chunks, seal_object_schemas, unwrap_json_tool_call, SchemaLimits, + ANTHROPIC_SCHEMA_LIMITS, GEMINI_OPENAPI_SCHEMA_LIMITS, JSON_TOOL_DESCRIPTION, JSON_TOOL_NAME, +}; pub use upstream_headers::{ apply_request_headers, client_header_forwardable, header_forward_blocked, resolve_default_headers, resolve_extra_headers, CallerIdentity, ForwardedClientHeaders, diff --git a/crates/aisix-gateway/src/structured_output.rs b/crates/aisix-gateway/src/structured_output.rs new file mode 100644 index 00000000..2bf735f2 --- /dev/null +++ b/crates/aisix-gateway/src/structured_output.rs @@ -0,0 +1,1232 @@ +//! Structured outputs: the parts every provider bridge translating an +//! OpenAI `response_format` needs to agree on. +//! +//! Providers reach JSON by two different routes and the gateway uses +//! both. Where the upstream has a native schema control (Anthropic's +//! `output_config.format`, Bedrock Converse's `outputConfig.textFormat`, +//! Gemini's `responseJsonSchema`) the schema goes straight onto the +//! wire. Where it does not, the schema rides a **synthetic tool** the +//! model is asked to call — [`JSON_TOOL_NAME`] — whose input *is* the +//! answer. The tool route needs the same two translations on the way +//! back everywhere it is used, so they live here rather than in one +//! provider crate: [`unwrap_json_tool_call`] turns the call back into +//! content, and [`response_into_fake_stream_chunks`] renders the +//! completed response as a stream, because a tool call cannot be +//! streamed before it is complete. +//! +//! [`seal_object_schemas`] and [`close_object_schemas`] are the two +//! schema normalisations these paths need; they differ only in whether +//! every declared property is forced into `required`. + +use crate::chat::{ChatChunk, ChatDelta, ChatResponse, FinishReason, Role}; + +/// Name of the synthetic tool the tool route asks the model to call. +/// The response decoder recognises it by this name to translate the +/// call back into plain JSON content, so the two sides must agree. +pub const JSON_TOOL_NAME: &str = "json_tool_call"; + +/// Description carried on the synthetic tool. +pub const JSON_TOOL_DESCRIPTION: &str = + "Respond by calling this tool with your answer as JSON matching its input schema."; + +/// Pull the JSON schema out of an OpenAI `response_format`, verbatim. +/// +/// Returns `None` for anything that is not a `json_schema` carrying a +/// non-null schema — `{"type":"json_object"}` and `{"type":"text"}` +/// included, neither of which names a schema to translate. +pub fn json_schema_from_response_format( + response_format: &serde_json::Value, +) -> Option { + if response_format.get("type").and_then(|t| t.as_str())? != "json_schema" { + return None; + } + response_format + .get("json_schema")? + .get("schema") + .filter(|s| !s.is_null()) + .cloned() +} + +/// Recursively set `additionalProperties: false` on every object schema, +/// leaving `required` exactly as the caller wrote it. +/// +/// This is what Anthropic and Bedrock require: both reject an object +/// that does not close, and both list `required` as an ordinary, +/// optional JSON Schema keyword — a property left out of it stays +/// optional and simply sorts after the required ones in the output. +/// Forcing every property into `required` would silently promote a +/// caller's optional field to mandatory, which changes what the model +/// is allowed to answer. +pub fn seal_object_schemas(schema: &mut serde_json::Value) { + walk_object_schemas(schema, false); +} + +/// Recursively make every object schema satisfy OpenAI **strict** mode: +/// `additionalProperties: false`, and every declared property listed in +/// `required`. Strict mode defines optionality through a nullable type +/// rather than through `required`, so the promotion is part of the +/// contract there — unlike [`seal_object_schemas`]. +pub fn close_object_schemas(schema: &mut serde_json::Value) { + walk_object_schemas(schema, true); +} + +/// Every member of a schema object that itself holds a schema, or a +/// collection of them. +/// +/// Both walkers in this module run off this one list. They used to carry +/// their own, which is how they came to disagree and how each came to +/// miss the applicator keywords entirely: a `maximum` under an +/// `if`/`then` branch stayed on the wire, and the object a +/// `Dict[str, Model]` field compiles to — `additionalProperties` in its +/// schema form — was never sealed. +/// +/// Three member shapes are handled: a map whose *values* are schemas +/// (`properties`, `$defs`, `patternProperties`, `dependentSchemas`), a +/// list of schemas (`anyOf`, `prefixItems`), and a single schema +/// (`items`, `not`, `contains`, …) — with `items` also taking the +/// draft-07 tuple form, so both shapes are tried for every one of those. +/// A map's keys are never schemas and never keywords: they are names the +/// caller chose. +fn for_each_subschema( + obj: &mut serde_json::Map, + visit: &mut impl FnMut(&mut serde_json::Value), +) { + const SCHEMA_MAPS: &[&str] = &[ + "properties", + "patternProperties", + "dependentSchemas", + "$defs", + "definitions", + ]; + const SCHEMA_LISTS: &[&str] = &["anyOf", "oneOf", "allOf", "prefixItems"]; + const SCHEMA_VALUES: &[&str] = &[ + "items", + "not", + "if", + "then", + "else", + "contains", + "propertyNames", + "additionalProperties", + "unevaluatedProperties", + "unevaluatedItems", + ]; + + for key in SCHEMA_MAPS { + if let Some(entries) = obj.get_mut(*key).and_then(|m| m.as_object_mut()) { + for entry in entries.values_mut() { + visit(entry); + } + } + } + for key in SCHEMA_LISTS { + if let Some(entries) = obj.get_mut(*key).and_then(|l| l.as_array_mut()) { + for entry in entries { + visit(entry); + } + } + } + for key in SCHEMA_VALUES { + match obj.get_mut(*key) { + // The tuple form of `items`. + Some(serde_json::Value::Array(entries)) => { + for entry in entries { + visit(entry); + } + } + // `additionalProperties: false` and friends are booleans, + // not schemas. + Some(value) if value.is_object() => visit(value), + _ => {} + } + } +} + +/// Whether a schema node describes an object and therefore has to be +/// sealed. `type` is not always the bare string `"object"`: the +/// canonical strict-mode spelling of an optional nested object is the +/// union `["object", "null"]`, and a node carrying `properties` with no +/// `type` at all is still an object schema. Missing either leaves that +/// node — and everything under it, since the walk would not recurse — +/// open, which the providers that require sealing reject outright. +fn is_object_schema(obj: &serde_json::Map) -> bool { + match obj.get("type") { + Some(serde_json::Value::String(ty)) => ty == "object", + Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some("object")), + // No `type`, but `properties` can only describe an object. + None => obj.contains_key("properties"), + _ => false, + } +} + +fn walk_object_schemas(schema: &mut serde_json::Value, require_every_property: bool) { + let Some(obj) = schema.as_object_mut() else { + return; + }; + if is_object_schema(obj) { + if let Some(properties) = obj.get("properties").and_then(|p| p.as_object()) { + let required: Vec = + properties.keys().map(|k| k.as_str().into()).collect(); + obj.insert("additionalProperties".to_string(), false.into()); + if require_every_property { + obj.insert("required".to_string(), required.into()); + } + } + } + for_each_subschema(obj, &mut |sub| { + walk_object_schemas(sub, require_every_property) + }); +} + +/// The subset of JSON Schema one provider's constrained decoder accepts. +/// +/// Every provider that compiles a schema into a decoding grammar +/// supports only a subset of JSON Schema and returns a 400 for anything +/// outside it — so a schema that worked against an OpenAI upstream +/// fails outright once the gateway starts forwarding it. Rather than +/// hand that error to a caller who did nothing wrong, each edge narrows +/// the schema to what its provider takes, and says in the schema itself +/// what it had to drop. +pub struct SchemaLimits { + /// Scalar constraint keywords the provider rejects. Each is removed + /// and recorded in that node's `description`, so the constraint is + /// still stated to the model even though it is no longer enforced by + /// the decoder. + pub noted_constraints: &'static [&'static str], + /// Keywords the provider rejects that say nothing a sentence can + /// carry — structural combinators and applicators. Removed quietly. + pub dropped_keywords: &'static [&'static str], + /// `minItems` values the provider accepts. `None` = all of them. + pub allowed_min_items: Option<&'static [u64]>, + /// Rewrite `oneOf` into `anyOf`. The providers here document + /// `anyOf` and not `oneOf`; for constraining *output* the + /// difference (exactly-one vs at-least-one) does not bind, since a + /// document the model produces matches whichever branch it followed. + /// Renaming keeps the alternatives, which dropping would not. + pub relax_one_of: bool, + /// Inline internal `$ref`s and remove the definition blocks they + /// point at. For providers whose schema dialect has no `$ref` at + /// all; the ones that document internal references keep theirs. + pub inline_internal_refs: bool, +} + +/// What Anthropic's structured outputs accept, per the "JSON Schema +/// limitations" section of their structured-outputs guide. Bedrock +/// documents the same subset for both its Converse `outputConfig` and +/// the Anthropic Messages `/invoke` body, so both edges use this. +/// +/// Internal `$ref` / `$defs` / `definitions` are supported by both and +/// are left in place. Recursive schemas and external `$ref`s are not, +/// and nothing this can do would make them legal, so they are left for +/// the upstream to reject. +pub const ANTHROPIC_SCHEMA_LIMITS: SchemaLimits = SchemaLimits { + noted_constraints: &[ + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "minLength", + "maxLength", + "maxItems", + "uniqueItems", + ], + dropped_keywords: &[], + allowed_min_items: Some(&[0, 1]), + relax_one_of: true, + inline_internal_refs: false, +}; + +/// What Gemini's older `responseSchema` dialect accepts. It is an +/// OpenAPI 3.0 `Schema` object, not JSON Schema: unknown members are +/// rejected by name, there is no `$ref`, and the applicator keywords +/// have no equivalent. Numeric and string bounds *are* part of that +/// dialect, so unlike Anthropic they survive. +pub const GEMINI_OPENAPI_SCHEMA_LIMITS: SchemaLimits = SchemaLimits { + noted_constraints: &[ + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "uniqueItems", + ], + // Everything the `Schema` type has no member for, taken from + // Google's own discovery document rather than from prose: + // `GET https://generativelanguage.googleapis.com/$discovery/rest?version=v1beta`, + // whose `schemas.Schema.properties` is exactly `anyOf`, `default`, + // `description`, `enum`, `example`, `format`, `items`, `maxItems`, + // `maxLength`, `maxProperties`, `maximum`, `minItems`, `minLength`, + // `minProperties`, `minimum`, `nullable`, `pattern`, `properties`, + // `propertyOrdering`, `required`, `title`, `type`. Anything else is + // rejected by name, so it cannot simply ride along. + dropped_keywords: &[ + "additionalProperties", + "allOf", + "not", + "if", + "then", + "else", + "const", + "contains", + "patternProperties", + "prefixItems", + "propertyNames", + "dependentSchemas", + "dependentRequired", + "unevaluatedProperties", + "unevaluatedItems", + "readOnly", + "writeOnly", + "deprecated", + "contentEncoding", + "contentMediaType", + "$schema", + "$id", + "$comment", + "$anchor", + ], + allowed_min_items: None, + relax_one_of: true, + inline_internal_refs: true, +}; + +/// Narrow `schema` to what `limits` says the provider accepts. +pub fn apply_schema_limits(schema: &mut serde_json::Value, limits: &SchemaLimits) { + if limits.inline_internal_refs { + inline_internal_refs(schema); + } + narrow_schema_node(schema, limits); +} + +fn narrow_schema_node(schema: &mut serde_json::Value, limits: &SchemaLimits) { + let Some(obj) = schema.as_object_mut() else { + return; + }; + + // Collect the constraints being removed in the order they are + // declared on the limits, so the note reads the same every time. + let mut notes: Vec = Vec::new(); + for key in limits.noted_constraints { + if let Some(value) = obj.remove(*key) { + notes.push(format!("{key}: {}", render_constraint(&value))); + } + } + if let Some(allowed) = limits.allowed_min_items { + let out_of_range = obj + .get("minItems") + .and_then(serde_json::Value::as_u64) + .is_some_and(|v| !allowed.contains(&v)); + if out_of_range { + if let Some(value) = obj.remove("minItems") { + notes.push(format!("minItems: {}", render_constraint(&value))); + } + } + } + if !notes.is_empty() { + let note = notes.join(", "); + let merged = match obj.get("description").and_then(|d| d.as_str()) { + Some(existing) if !existing.is_empty() => format!("{existing} ({note})"), + _ => note, + }; + obj.insert("description".to_string(), merged.into()); + } + + for key in limits.dropped_keywords { + obj.remove(*key); + } + if limits.relax_one_of { + if let Some(branches) = obj.remove("oneOf") { + obj.entry("anyOf").or_insert(branches); + } + } + + // Every position that holds a schema, from the one list both walkers + // read. A blind walk over the members would read the keys of + // `properties` as keywords, so a caller whose document has a field + // called `minimum` or `const` would lose that field and gain a + // `description` built out of its own property names. + for_each_subschema(obj, &mut |sub| narrow_schema_node(sub, limits)); +} + +/// Render a constraint value for the description note. Strings keep +/// their quotes off; everything else is its compact JSON form. +fn render_constraint(value: &serde_json::Value) -> String { + match value { + serde_json::Value::String(s) => s.clone(), + other => other.to_string(), + } +} + +/// Replace every internal `$ref` with the definition it names and drop +/// the definition blocks, for dialects that have no `$ref`. +/// +/// A `$ref` this cannot resolve — external, or recursive past +/// [`MAX_REF_DEPTH`] — is left exactly as it came in. Nothing this +/// function could do would make such a schema legal, so the upstream's +/// own rejection is the honest outcome. +fn inline_internal_refs(schema: &mut serde_json::Value) { + // `$defs` and `definitions` are separate namespaces — a schema may + // define the same name in both — so the map is keyed by the pointer + // that reaches each one, not by the bare name. + let mut defs: serde_json::Map = serde_json::Map::new(); + for block in ["$defs", "definitions"] { + if let Some(entries) = schema.get(block).and_then(|d| d.as_object()) { + for (name, definition) in entries { + defs.insert(format!("{block}/{name}"), definition.clone()); + } + } + } + if defs.is_empty() { + return; + } + // Inlining runs on a copy under a global expansion budget, and the + // schema is only adopted if it finished inside it. A recursive + // definition with several alternatives multiplies at every level — + // a few hundred bytes of schema can expand into hundreds of + // megabytes — and this runs synchronously while a caller waits, so + // the budget has to bound the total work, not just the depth of one + // chain. On overrun nothing is rewritten: the `$ref`s go upstream + // as the caller wrote them and the provider rejects what it cannot + // resolve, which is the same outcome as any other unresolvable + // reference here. + // Each definition's serialised size, measured once. The count alone + // does not bound the work: a definition carrying a few hundred + // kilobytes of `description` reaches a hundred megabytes well inside + // any sane expansion count, and that document is then serialised + // again on its way upstream. + let mut budget = RefBudget { + expansions: MAX_REF_EXPANSIONS, + bytes: MAX_REF_BYTES, + sizes: defs + .iter() + .map(|(name, definition)| (name.clone(), definition.to_string().len())) + .collect(), + }; + let mut working = schema.clone(); + if !substitute_refs(&mut working, &defs, 0, &mut budget) { + tracing::debug!("leaving $ref in place: inlining exceeded its expansion budget"); + return; + } + if let Some(obj) = working.as_object_mut() { + obj.remove("$defs"); + obj.remove("definitions"); + } + *schema = working; +} + +/// How many times one `$ref` chain is followed before giving up. Bounds +/// the depth of a single chain; [`MAX_REF_EXPANSIONS`] bounds the whole +/// job, which is what a recursive definition with several alternatives +/// actually blows through. +const MAX_REF_DEPTH: usize = 8; + +/// Total `$ref` expansions allowed for one schema. Comfortably above +/// any hand-written or generated schema — a large typed model produces +/// tens — and far below the point where expansion costs real time. +const MAX_REF_EXPANSIONS: usize = 2_000; + +/// Total bytes of definition allowed to be spliced in. The count bounds +/// how many times a definition is copied; this bounds how large the +/// copies are, which is the half that decides how much memory the +/// expanded schema — and its re-serialisation on the way upstream — +/// occupies. +const MAX_REF_BYTES: usize = 4 * 1024 * 1024; + +/// What one inlining pass is allowed to spend. +struct RefBudget { + expansions: usize, + bytes: usize, + /// Serialised size of each definition, by the key that reaches it. + sizes: std::collections::HashMap, +} + +impl RefBudget { + /// Charge one expansion of `name`. `false` once either half is gone. + fn charge(&mut self, name: &str) -> bool { + let Some(expansions) = self.expansions.checked_sub(1) else { + return false; + }; + let size = self.sizes.get(name).copied().unwrap_or(0); + let Some(bytes) = self.bytes.checked_sub(size) else { + return false; + }; + self.expansions = expansions; + self.bytes = bytes; + true + } +} + +/// Expand every resolvable internal `$ref` in `node`. Returns `false` +/// when `budget` ran out, in which case `node` is left partly rewritten +/// and the caller must discard it. +fn substitute_refs( + node: &mut serde_json::Value, + defs: &serde_json::Map, + depth: usize, + budget: &mut RefBudget, +) -> bool { + if let Some(array) = node.as_array_mut() { + for item in array { + if !substitute_refs(item, defs, depth, budget) { + return false; + } + } + return true; + } + let Some(obj) = node.as_object_mut() else { + return true; + }; + if let Some(reference) = obj.get("$ref").and_then(|r| r.as_str()) { + let Some(name) = internal_ref_name(reference) else { + return true; // external reference: not ours to resolve + }; + let Some(definition) = defs.get(name) else { + return true; + }; + if depth >= MAX_REF_DEPTH { + return true; // recursive chain: leave the `$ref` in place + } + if !budget.charge(name) { + return false; + } + let mut expanded = definition.clone(); + if !substitute_refs(&mut expanded, defs, depth + 1, budget) { + return false; + } + *node = expanded; + return true; + } + for value in obj.values_mut() { + if !substitute_refs(value, defs, depth, budget) { + return false; + } + } + true +} + +/// The `/` key a `#/$defs/Name` or `#/definitions/Name` +/// pointer resolves to. `None` for anything else, which includes every +/// external reference. +fn internal_ref_name(reference: &str) -> Option<&str> { + let path = reference.strip_prefix("#/")?; + let (block, name) = path.split_once('/')?; + if !matches!(block, "$defs" | "definitions") || name.is_empty() || name.contains('/') { + return None; + } + Some(path) +} + +/// Undo the tool route: turn the model's call to the synthetic +/// [`JSON_TOOL_NAME`] tool back into the plain JSON content the caller +/// asked for. Only ever applied to a response whose request carried the +/// synthetic tool, so a caller's own tool of that name is never touched. +/// +/// The call's arguments are already the JSON-encoded tool input, which +/// is exactly the document the schema describes. +/// +/// When it is the only call the JSON **replaces** the content: the +/// caller asked for a document they can parse, and a model that +/// narrated before calling the tool ("Sure, here you go:") would +/// otherwise leave them with a string that is not JSON. Any prose is +/// dropped, `tool_calls` with it, and a tool-use finish reason is +/// demoted to `stop` — what a client that never offered a tool must +/// see. Prose is likeliest exactly where the tool could not be forced +/// (a caller's own `tool_choice`, extended thinking, a Converse family +/// with no `toolChoice`), so this is not a rare shape. +/// +/// When the model called real tools alongside it, the caller *did* ask +/// for tool calls and is parsing the response themselves, so those +/// calls and their finish reason are left untouched and the JSON is +/// appended to whatever text came with them. +pub fn unwrap_json_tool_call(resp: &mut ChatResponse) { + let mut json_parts: Vec = Vec::new(); + let mut real_calls_remain = false; + if let Some(serde_json::Value::Array(calls)) = resp.message.extra.get_mut("tool_calls") { + calls.retain(|call| { + let name = call + .pointer("/function/name") + .and_then(|n| n.as_str()) + .unwrap_or_default(); + if name != JSON_TOOL_NAME { + return true; + } + if let Some(args) = call.pointer("/function/arguments").and_then(|a| a.as_str()) { + json_parts.push(args.to_string()); + } + false + }); + real_calls_remain = !calls.is_empty(); + } + if json_parts.is_empty() { + return; + } + let json = json_parts.join("\n"); + if !real_calls_remain { + resp.message.extra.remove("tool_calls"); + resp.finish_reason = FinishReason::Stop; + resp.message.content = Some(json); + return; + } + resp.message.content = Some(match resp.message.content.take() { + Some(text) if !text.is_empty() => format!("{text}\n{json}"), + _ => json, + }); +} + +/// Render a complete response as the chunk sequence a streaming client +/// expects: role, content, finish, usage. +/// +/// The tool route cannot stream — the JSON only exists once the tool +/// call is complete — so a bridge runs that request non-streaming and +/// fake-streams the result through here. Keeping the usage on its own +/// terminal chunk matches what a real upstream emits, so the proxy's +/// accounting and every downstream encoder see an ordinary stream. +pub fn response_into_fake_stream_chunks(resp: ChatResponse) -> Vec { + let ChatResponse { + id, + model, + message, + finish_reason, + usage, + } = resp; + let chunk = |delta, finish_reason, usage| ChatChunk { + id: id.clone(), + model: model.clone(), + delta, + finish_reason, + usage, + }; + // The non-streaming `tool_calls` shape carries no `index`, but the + // streaming one must: OpenAI SDKs accumulate by it, and this repo's + // Anthropic SSE re-encoder reads it to key each `content_block`, + // folding every index-less call onto block 0. Number them densely + // in arrival order, leaving any index a decoder already assigned. + let tool_calls = message + .extra + .get("tool_calls") + .and_then(|c| c.as_array()) + .map(|calls| { + calls + .iter() + .enumerate() + .map(|(i, call)| { + let mut call = call.clone(); + if let Some(obj) = call.as_object_mut() { + obj.entry("index").or_insert(i.into()); + } + call + }) + .collect() + }); + vec![ + chunk( + ChatDelta { + role: Some(Role::Assistant), + ..ChatDelta::default() + }, + None, + None, + ), + chunk( + ChatDelta { + content: Some(message.content.unwrap_or_default()), + tool_calls, + ..ChatDelta::default() + }, + None, + None, + ), + chunk(ChatDelta::default(), Some(finish_reason), None), + chunk(ChatDelta::default(), None, Some(usage)), + ] +} + +#[cfg(test)] +mod tests { + use super::*; + + fn person_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "pet": { + "type": "object", + "properties": {"kind": {"type": "string"}}, + "required": ["kind"], + }, + }, + "required": ["name"], + }) + } + + #[test] + fn sealing_closes_every_object_and_leaves_required_alone() { + let mut schema = person_schema(); + seal_object_schemas(&mut schema); + assert_eq!(schema["additionalProperties"], false); + assert_eq!(schema["properties"]["pet"]["additionalProperties"], false); + // `pet` was optional and stays optional — the whole point of the + // seal-only variant. + assert_eq!(schema["required"], serde_json::json!(["name"])); + assert_eq!( + schema["properties"]["pet"]["required"], + serde_json::json!(["kind"]) + ); + } + + #[test] + fn strict_closing_promotes_every_property_to_required() { + let mut schema = person_schema(); + close_object_schemas(&mut schema); + assert_eq!(schema["additionalProperties"], false); + assert_eq!(schema["required"], serde_json::json!(["name", "pet"])); + } + + #[test] + fn sealing_reaches_arrays_branches_and_definitions() { + let mut schema = serde_json::json!({ + "type": "array", + "items": {"type": "object", "properties": {"a": {"type": "string"}}}, + "anyOf": [{"type": "object", "properties": {"b": {"type": "string"}}}], + "$defs": {"d": {"type": "object", "properties": {"c": {"type": "string"}}}}, + }); + seal_object_schemas(&mut schema); + assert_eq!(schema["items"]["additionalProperties"], false); + assert_eq!(schema["anyOf"][0]["additionalProperties"], false); + assert_eq!(schema["$defs"]["d"]["additionalProperties"], false); + } + + #[test] + fn sealing_recognises_union_typed_and_untyped_object_nodes() { + // `["object","null"]` is how strict mode spells an optional + // nested object, and a node with `properties` and no `type` is + // still an object. Missing either leaves the whole subtree open + // and the provider rejects the request. + let mut schema = serde_json::json!({ + "type": "object", + "properties": { + "nullable": { + "type": ["object", "null"], + "properties": {"a": {"type": "string"}}, + }, + "untyped": {"properties": {"b": {"type": "string"}}}, + }, + }); + seal_object_schemas(&mut schema); + assert_eq!( + schema["properties"]["nullable"]["additionalProperties"], + false + ); + assert_eq!( + schema["properties"]["nullable"]["properties"]["a"]["type"], + "string" + ); + assert_eq!( + schema["properties"]["untyped"]["additionalProperties"], + false + ); + } + + // ── provider schema subsets ─────────────────────────────────── + + #[test] + fn anthropic_limits_strip_every_unsupported_constraint_and_say_so() { + let mut schema = serde_json::json!({ + "type": "object", + "properties": { + "age": { + "type": "integer", + "description": "the age", + "minimum": 1, + "maximum": 120, + "exclusiveMinimum": 0, + "exclusiveMaximum": 121, + "multipleOf": 1, + }, + "name": {"type": "string", "minLength": 2, "maxLength": 20}, + "tags": { + "type": "array", + "items": {"type": "string"}, + "minItems": 3, + "maxItems": 9, + "uniqueItems": true, + }, + }, + }); + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + + let age = &schema["properties"]["age"]; + for keyword in [ + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + ] { + assert!(age.get(keyword).is_none(), "{keyword} must be stripped"); + } + // An existing description keeps its own text and gains the note. + assert_eq!( + age["description"], + concat!( + "the age (minimum: 1, maximum: 120, exclusiveMinimum: 0, ", + "exclusiveMaximum: 121, multipleOf: 1)" + ) + ); + + let name = &schema["properties"]["name"]; + assert!(name.get("minLength").is_none()); + assert!(name.get("maxLength").is_none()); + // No description to begin with: the note becomes one. + assert_eq!(name["description"], "minLength: 2, maxLength: 20"); + + let tags = &schema["properties"]["tags"]; + assert!(tags.get("maxItems").is_none()); + assert!(tags.get("uniqueItems").is_none()); + // `minItems` is supported only at 0 and 1, so 3 goes too. + assert!(tags.get("minItems").is_none()); + assert_eq!( + tags["description"], + "maxItems: 9, uniqueItems: true, minItems: 3" + ); + } + + #[test] + fn anthropic_limits_keep_the_min_items_values_the_provider_takes() { + for kept in [0, 1] { + let mut schema = + serde_json::json!({"type": "array", "items": {"type": "string"}, "minItems": kept}); + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + assert_eq!(schema["minItems"], kept, "minItems {kept} is supported"); + assert!(schema.get("description").is_none()); + } + } + + #[test] + fn one_of_is_relaxed_to_any_of_rather_than_dropped() { + // Neither provider documents `oneOf`; dropping it would take the + // alternatives with it, so the branches move to `anyOf`. + let mut schema = serde_json::json!({ + "oneOf": [{"type": "string"}, {"type": "integer"}], + }); + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + assert!(schema.get("oneOf").is_none()); + assert_eq!( + schema["anyOf"], + serde_json::json!([{"type": "string"}, {"type": "integer"}]) + ); + } + + #[test] + fn anthropic_limits_leave_internal_references_in_place() { + // Anthropic and Bedrock both document internal `$ref`; only the + // dialects without one need inlining. + let mut schema = serde_json::json!({ + "type": "object", + "properties": {"pet": {"$ref": "#/$defs/Pet"}}, + "$defs": {"Pet": {"type": "object", "properties": {"kind": {"type": "string"}}}}, + }); + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + assert_eq!(schema["properties"]["pet"]["$ref"], "#/$defs/Pet"); + assert!(schema["$defs"]["Pet"].is_object()); + } + + #[test] + fn gemini_limits_inline_internal_references_and_drop_the_blocks() { + let mut schema = serde_json::json!({ + "type": "object", + "properties": { + "pet": {"$ref": "#/$defs/Pet"}, + "other": {"$ref": "#/definitions/Pet"}, + }, + "$defs": {"Pet": {"type": "object", "properties": {"kind": {"type": "string"}}}}, + "definitions": {"Pet": {"type": "string"}}, + }); + apply_schema_limits(&mut schema, &GEMINI_OPENAPI_SCHEMA_LIMITS); + assert_eq!(schema["properties"]["pet"]["type"], "object"); + assert_eq!( + schema["properties"]["pet"]["properties"]["kind"]["type"], + "string" + ); + assert_eq!(schema["properties"]["other"]["type"], "string"); + assert!(schema.get("$defs").is_none()); + assert!(schema.get("definitions").is_none()); + } + + #[test] + fn an_external_or_recursive_reference_is_left_for_the_upstream_to_reject() { + // Nothing inlining can do makes either legal, so the request + // goes as written and the provider says why. + let mut schema = serde_json::json!({ + "type": "object", + "properties": {"remote": {"$ref": "https://example.com/Pet.json"}}, + "$defs": {"Pet": {"type": "string"}}, + }); + apply_schema_limits(&mut schema, &GEMINI_OPENAPI_SCHEMA_LIMITS); + assert_eq!( + schema["properties"]["remote"]["$ref"], + "https://example.com/Pet.json" + ); + + let mut recursive = serde_json::json!({ + "$ref": "#/$defs/Node", + "$defs": { + "Node": { + "type": "object", + "properties": {"child": {"$ref": "#/$defs/Node"}}, + }, + }, + }); + apply_schema_limits(&mut recursive, &GEMINI_OPENAPI_SCHEMA_LIMITS); + // Expansion stops at the depth cap rather than looping; the + // innermost `$ref` survives and Vertex rejects it. + let json = recursive.to_string(); + assert!(json.contains("#/$defs/Node"), "{json}"); + } + + #[test] + fn gemini_limits_drop_the_applicators_the_dialect_has_no_member_for() { + let mut schema = serde_json::json!({ + "type": "object", + "properties": { + "a": {"type": "string", "const": "x", "multipleOf": 2, "uniqueItems": true}, + }, + "allOf": [{"type": "object"}], + "not": {"type": "null"}, + "if": {"type": "object"}, + "then": {"type": "object"}, + "else": {"type": "object"}, + "patternProperties": {"^a": {"type": "string"}}, + "prefixItems": [{"type": "string"}], + }); + apply_schema_limits(&mut schema, &GEMINI_OPENAPI_SCHEMA_LIMITS); + for dropped in [ + "allOf", + "not", + "if", + "then", + "else", + "patternProperties", + "prefixItems", + ] { + assert!(schema.get(dropped).is_none(), "{dropped} must be dropped"); + } + let a = &schema["properties"]["a"]; + assert!(a.get("const").is_none()); + assert!(a.get("multipleOf").is_none()); + assert_eq!(a["description"], "multipleOf: 2, uniqueItems: true"); + // Bounds ARE part of the OpenAPI dialect, so they survive. + let mut bounded = serde_json::json!({"type": "integer", "minimum": 1, "maximum": 9}); + apply_schema_limits(&mut bounded, &GEMINI_OPENAPI_SCHEMA_LIMITS); + assert_eq!(bounded["minimum"], 1); + assert_eq!(bounded["maximum"], 9); + } + + #[test] + fn a_property_named_like_a_keyword_is_not_mistaken_for_one() { + // `properties` is a map of caller-chosen names, not of schema + // keywords. Walking it blindly deletes a field called `minimum` + // and builds a `description` out of the caller's own field + // names — so a perfectly ordinary document schema comes out + // missing members. + let document = serde_json::json!({ + "type": "object", + "properties": { + "minimum": {"type": "number"}, + "maximum": {"type": "number"}, + "maxLength": {"type": "integer"}, + "const": {"type": "string"}, + "if": {"type": "boolean"}, + "allOf": {"type": "string"}, + "uniqueItems": {"type": "boolean"}, + }, + }); + for (label, limits) in [ + ("anthropic", &ANTHROPIC_SCHEMA_LIMITS), + ("gemini", &GEMINI_OPENAPI_SCHEMA_LIMITS), + ] { + let mut schema = document.clone(); + apply_schema_limits(&mut schema, limits); + assert_eq!( + schema["properties"], document["properties"], + "{label}: every property must survive untouched" + ); + assert!( + schema.get("description").is_none(), + "{label}: no note belongs on a node with no constraints" + ); + } + } + + #[test] + fn tuple_form_items_are_reached_by_both_walkers() { + // Draft-07 spells a positional array as `items: [schema, …]`. + // That is a list of schemas, not a schema, so a walker that + // treats it as one skips every element — leaving those objects + // open and their constraints on the wire. + let document = serde_json::json!({ + "type": "array", + "items": [ + {"type": "object", "properties": {"a": {"type": "string", "maxLength": 4}}}, + {"type": "integer", "minimum": 2}, + ], + }); + + let mut sealed = document.clone(); + seal_object_schemas(&mut sealed); + assert_eq!(sealed["items"][0]["additionalProperties"], false); + + let mut narrowed = document.clone(); + apply_schema_limits(&mut narrowed, &ANTHROPIC_SCHEMA_LIMITS); + assert!(narrowed["items"][0]["properties"]["a"] + .get("maxLength") + .is_none()); + assert_eq!( + narrowed["items"][0]["properties"]["a"]["description"], + "maxLength: 4" + ); + assert!(narrowed["items"][1].get("minimum").is_none()); + assert_eq!(narrowed["items"][1]["description"], "minimum: 2"); + } + + #[test] + fn a_branching_recursive_schema_is_left_alone_rather_than_expanded() { + // Each level multiplies by the number of alternatives, so a few + // hundred bytes can expand into hundreds of megabytes — on the + // request path, with a caller waiting. The budget bounds the + // whole job, and on overrun nothing is rewritten. + let branches: Vec = (0..7) + .map(|i| serde_json::json!({format!("child{i}"): {"$ref": "#/$defs/Node"}})) + .collect(); + let mut properties = serde_json::Map::new(); + for branch in &branches { + for (k, v) in branch.as_object().unwrap() { + properties.insert(k.clone(), v.clone()); + } + } + let schema = serde_json::json!({ + "$ref": "#/$defs/Node", + "$defs": {"Node": {"type": "object", "properties": properties}}, + }); + + let mut narrowed = schema.clone(); + let started = std::time::Instant::now(); + apply_schema_limits(&mut narrowed, &GEMINI_OPENAPI_SCHEMA_LIMITS); + let elapsed = started.elapsed(); + + assert!( + elapsed < std::time::Duration::from_secs(2), + "inlining must not run away: took {elapsed:?}" + ); + // Untouched: the `$ref` and its definition both survive, so the + // upstream gets the schema as written and says why it cannot + // take it. + assert_eq!(narrowed["$ref"], "#/$defs/Node"); + assert!(narrowed["$defs"]["Node"].is_object()); + assert!( + narrowed.to_string().len() < 4_096, + "nothing should have been expanded" + ); + } + + #[test] + fn an_ordinary_nested_schema_still_inlines_under_the_budget() { + // The budget must not be so tight that real schemas stop + // inlining — the generated ones nest a handful of models deep. + let mut schema = serde_json::json!({ + "type": "object", + "properties": { + "a": {"$ref": "#/$defs/Inner"}, + "b": {"$ref": "#/$defs/Inner"}, + "c": {"type": "array", "items": {"$ref": "#/$defs/Inner"}}, + }, + "$defs": { + "Inner": {"type": "object", "properties": {"leaf": {"$ref": "#/$defs/Leaf"}}}, + "Leaf": {"type": "string"}, + }, + }); + apply_schema_limits(&mut schema, &GEMINI_OPENAPI_SCHEMA_LIMITS); + assert!(schema.get("$defs").is_none()); + assert_eq!( + schema["properties"]["a"]["properties"]["leaf"]["type"], + "string" + ); + assert_eq!( + schema["properties"]["c"]["items"]["properties"]["leaf"]["type"], + "string" + ); + } + + /// Every position a sub-schema can sit in, with the node to place + /// there and the pointer that reaches it once placed. Both walkers + /// run off one list, so a position missing from that list is missing + /// from both — which is how the applicator keywords came to be + /// skipped by each of them at once. + fn subschema_positions() -> Vec<(&'static str, serde_json::Value, &'static str)> { + let object = serde_json::json!({"type": "object", "properties": {"n": {"type": "string"}}}); + vec![ + ( + "properties", + serde_json::json!({"child": object}), + "/properties/child", + ), + ("items", object.clone(), "/items"), + ("items (tuple)", serde_json::json!([object]), "/items/0"), + ("prefixItems", serde_json::json!([object]), "/prefixItems/0"), + ("anyOf", serde_json::json!([object]), "/anyOf/0"), + ("allOf", serde_json::json!([object]), "/allOf/0"), + ("$defs", serde_json::json!({"D": object}), "/$defs/D"), + ( + "definitions", + serde_json::json!({"D": object}), + "/definitions/D", + ), + ( + "additionalProperties", + object.clone(), + "/additionalProperties", + ), + ( + "patternProperties", + serde_json::json!({"^a": object}), + "/patternProperties/^a", + ), + ( + "dependentSchemas", + serde_json::json!({"a": object}), + "/dependentSchemas/a", + ), + ("not", object.clone(), "/not"), + ("if", object.clone(), "/if"), + ("then", object.clone(), "/then"), + ("else", object.clone(), "/else"), + ("contains", object.clone(), "/contains"), + ("propertyNames", object.clone(), "/propertyNames"), + ( + "unevaluatedProperties", + object.clone(), + "/unevaluatedProperties", + ), + ("unevaluatedItems", object, "/unevaluatedItems"), + ] + } + + #[test] + fn sealing_reaches_every_sub_schema_position() { + for (name, value, pointer) in subschema_positions() { + let mut schema = serde_json::json!({}); + schema[name.split_whitespace().next().unwrap()] = value; + seal_object_schemas(&mut schema); + assert_eq!( + schema + .pointer(pointer) + .and_then(|s| s.get("additionalProperties")), + Some(&serde_json::Value::Bool(false)), + "{name}: the schema under {pointer} was never sealed" + ); + } + } + + #[test] + fn narrowing_reaches_every_sub_schema_position() { + for (name, value, pointer) in subschema_positions() { + let key = name.split_whitespace().next().unwrap(); + // Put a constraint the Anthropic subset rejects on the leaf, + // so reaching the node is observable. + let mut value = value; + let leaf_pointer = pointer.trim_start_matches(&format!("/{key}")).to_string(); + let leaf = match leaf_pointer.as_str() { + "" => &mut value, + p => value.pointer_mut(p).unwrap(), + }; + leaf["properties"]["n"]["maxLength"] = 7.into(); + + let mut schema = serde_json::json!({}); + schema[key] = value; + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + + let leaf = schema + .pointer(&format!("{pointer}/properties/n")) + .unwrap_or_else(|| panic!("{name}: {pointer} vanished")); + assert!( + leaf.get("maxLength").is_none(), + "{name}: maxLength under {pointer} reached the wire" + ); + assert_eq!(leaf["description"], "maxLength: 7", "{name}"); + } + } + + #[test] + fn a_definition_too_large_to_splice_is_left_as_a_reference() { + // The expansion COUNT alone does not bound the work. This schema + // is nowhere near the count budget — twenty references, no + // recursion — but each one splices in 400KB, so the document + // that would go upstream is megabytes of duplicated text, and it + // is serialised again on the way out. + let big = "x".repeat(400 * 1024); + let mut properties = serde_json::Map::new(); + for i in 0..20 { + properties.insert(format!("f{i}"), serde_json::json!({"$ref": "#/$defs/Big"})); + } + let schema = serde_json::json!({ + "type": "object", + "properties": properties, + "$defs": {"Big": {"type": "string", "description": big}}, + }); + + let mut narrowed = schema.clone(); + apply_schema_limits(&mut narrowed, &GEMINI_OPENAPI_SCHEMA_LIMITS); + + // Sent exactly as the caller wrote it, references intact; the + // provider says why it cannot take them. + assert_eq!(narrowed["properties"]["f0"]["$ref"], "#/$defs/Big"); + assert!(narrowed["$defs"]["Big"].is_object()); + assert!( + narrowed.to_string().len() < 2 * big.len(), + "nothing should have been spliced in" + ); + } + + #[test] + fn a_handful_of_large_references_still_inlines() { + // The byte budget must not be so tight that an ordinary schema + // with a couple of well-documented models stops inlining. + let text = "x".repeat(8 * 1024); + let schema = serde_json::json!({ + "type": "object", + "properties": { + "a": {"$ref": "#/$defs/Doc"}, + "b": {"$ref": "#/$defs/Doc"}, + }, + "$defs": {"Doc": {"type": "string", "description": text}}, + }); + let mut narrowed = schema; + apply_schema_limits(&mut narrowed, &GEMINI_OPENAPI_SCHEMA_LIMITS); + assert!(narrowed.get("$defs").is_none()); + assert_eq!(narrowed["properties"]["a"]["type"], "string"); + assert_eq!(narrowed["properties"]["b"]["type"], "string"); + } + + #[test] + fn only_a_json_schema_response_format_yields_a_schema() { + let schema = serde_json::json!({"type": "object"}); + assert_eq!( + json_schema_from_response_format(&serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema, "strict": true}, + })), + Some(schema) + ); + for other in [ + serde_json::json!({"type": "json_object"}), + serde_json::json!({"type": "text"}), + serde_json::json!({"type": "json_schema", "json_schema": {"name": "answer"}}), + serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": null}, + }), + ] { + assert_eq!(json_schema_from_response_format(&other), None, "{other}"); + } + } +} diff --git a/crates/aisix-provider-anthropic/src/bridge.rs b/crates/aisix-provider-anthropic/src/bridge.rs index ea2bb32e..8332cb6f 100644 --- a/crates/aisix-provider-anthropic/src/bridge.rs +++ b/crates/aisix-provider-anthropic/src/bridge.rs @@ -24,9 +24,11 @@ use futures::StreamExt; use reqwest::{header, Client, StatusCode}; use std::time::{Duration, Instant}; +use aisix_gateway::structured_output::{response_into_fake_stream_chunks, unwrap_json_tool_call}; + use crate::wire::{ build_request, inject_cache_breakpoints, response_into_chat_response, split_system, - AnthropicResponse, AnthropicStreamEvent, StreamState, + structured_output_for, AnthropicResponse, AnthropicStreamEvent, StreamState, StructuredOutput, }; /// Matches the API header that Anthropic bakes backwards-compat into. @@ -360,6 +362,10 @@ impl Bridge for AnthropicBridge { let (system, messages) = split_system(req).map_err(|e| BridgeError::InvalidUpstreamConfig(e.to_string()))?; let mut body = build_request(req, upstream, system, messages, false); + let synthetic_json_tool = matches!( + structured_output_for(req, upstream), + StructuredOutput::Tool(_) + ); maybe_inject_cache_breakpoints(&mut body, ctx); let url = cached_endpoint_url( &ctx.provider_key_id, @@ -396,7 +402,11 @@ impl Bridge for AnthropicBridge { .json() .await .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; - Ok(response_into_chat_response(parsed)) + let mut chat = response_into_chat_response(parsed); + if synthetic_json_tool { + unwrap_json_tool_call(&mut chat); + } + Ok(chat) }) .await } @@ -409,6 +419,23 @@ impl Bridge for AnthropicBridge { let key = api_key(ctx)?; let upstream = upstream_model(ctx)?; + // The tool path's JSON only exists once the synthetic tool call + // has been assembled, so it cannot be streamed as it arrives. + // Run the request non-streaming and fake-stream the translated + // result: the client sees an ordinary chunk sequence, and usage + // rides its own terminal chunk exactly as on a real stream. + if matches!( + structured_output_for(req, upstream), + StructuredOutput::Tool(_) + ) { + // The leg is not streaming, so it runs under the budget a + // non-streaming call would have got — the streaming budget + // this context carries bounds a chunk gap, not a completion. + let chunks = + response_into_fake_stream_chunks(self.chat(req, &ctx.non_streaming_ctx()).await?); + return Ok(Box::pin(futures::stream::iter(chunks.into_iter().map(Ok)))); + } + let (system, messages) = split_system(req).map_err(|e| BridgeError::InvalidUpstreamConfig(e.to_string()))?; let mut body = build_request(req, upstream, system, messages, true); @@ -553,6 +580,77 @@ mod tests { .await; } + #[tokio::test] + async fn a_small_stream_budget_does_not_cut_the_fake_stream_leg() { + // On a streaming dispatch the deadline is the streaming budget, + // which bounds a chunk gap rather than a whole completion. The + // tool route's upstream leg is not streaming, so it runs under + // the end-to-end budget carried beside it. + let server = MockServer::start().await; + Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/v1/messages")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(std::time::Duration::from_millis(250)) + .set_body_json(serde_json::json!({ + "id": "msg_json", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": "json_tool_call", + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + })), + ) + .mount(&server) + .await; + + let model: Model = serde_json::from_value(serde_json::json!({ + "display_name": "my-claude", + "provider": "anthropic", + // Older family: takes the tool route, which cannot stream. + "model_name": "claude-3-5-haiku-20241022", + "provider_key_id": "11111111-1111-1111-1111-111111111111", + })) + .unwrap(); + let ctx = BridgeContext::new("req-1", Arc::new(model), sample_provider_key(&server.uri())) + .with_deadline(std::time::Duration::from_millis(50)) + .with_non_streaming_deadline(Some(std::time::Duration::from_secs(30))); + + let mut req = ChatFormat::new("my-claude", vec![ChatMessage::user("who is Ada")]); + req.stream = Some(true); + req.extra.insert( + "response_format".into(), + serde_json::json!({ + "type": "json_schema", + "json_schema": { + "name": "person", + "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "strict": true, + }, + }), + ); + + let stream = AnthropicBridge::new() + .chat_stream(&req, &ctx) + .await + .expect("the fake-stream leg must not be cut by the chunk-gap budget"); + let chunks: Vec = futures::StreamExt::collect::>(stream) + .await + .into_iter() + .map(Result::unwrap) + .collect(); + assert_eq!( + chunks[1].delta.content.as_deref(), + Some(r#"{"name":"Ada"}"#) + ); + } + #[tokio::test] async fn non_streaming_injects_breakpoints_when_enabled() { let server = MockServer::start().await; diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index bba78e1c..c3fbeca0 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/crates/aisix-provider-anthropic/src/wire.rs @@ -24,6 +24,10 @@ use std::borrow::Cow; use aisix_core::MappedEffort; +use aisix_gateway::structured_output::{ + apply_schema_limits, json_schema_from_response_format, seal_object_schemas, + ANTHROPIC_SCHEMA_LIMITS, JSON_TOOL_DESCRIPTION, JSON_TOOL_NAME, +}; use aisix_gateway::{ BridgeError, ChatChunk, ChatDelta, ChatFormat, ChatMessage, ChatResponse, FinishReason, Role, UsageStats, @@ -446,7 +450,7 @@ pub fn build_request<'a>( // no tool survives translation: upstream rejects the field without // an accompanying `tools` list (AISIX-Cloud#1614). let mut extras = req.extra.clone(); - let tools = extras + let mut tools = extras .remove("tools") .and_then(translate_openai_tools_to_anthropic); let requested_tool_choice = extras.remove("tool_choice"); @@ -462,6 +466,7 @@ pub fn build_request<'a>( .as_ref() .and_then(serde_json::Value::as_bool) == Some(false); + let client_set_tool_choice = tool_choice_states_a_preference(requested_tool_choice.as_ref()); let mut tool_choice = tools .as_ref() .and(requested_tool_choice) @@ -483,14 +488,58 @@ pub fn build_request<'a>( } } translate_reasoning_effort_to_anthropic(&mut extras); - // `response_format` is the OpenAI spelling of structured outputs and - // has no top-level Anthropic counterpart, so it would ride `extra` - // onto the body and be rejected as an unknown parameter. It reaches - // this bridge from a chat caller and from the `/v1/responses` - // translation of `text.format`; both are dropped here, as every - // OpenAI-only knob with no provider-neutral equivalent is. - if extras.remove("response_format").is_some() { - tracing::debug!("dropping response_format: no Anthropic counterpart on this path"); + // `response_format` is the OpenAI spelling of structured outputs. It + // has no top-level Anthropic counterpart, so it is always consumed + // here — riding `extra` onto the body would be rejected as an unknown + // parameter. What it becomes instead depends on the target model; see + // [`StructuredOutput`]. It reaches this bridge from a chat caller and + // from the `/v1/responses` translation of `text.format`. + let structured_output = structured_output_for(req, upstream_model); + extras.remove("response_format"); + match structured_output { + StructuredOutput::None => {} + StructuredOutput::Native(schema) => { + let format = serde_json::json!({"type": "json_schema", "schema": schema}); + match extras.get_mut("output_config") { + // `output_config` is a carrier shared with `effort` and + // `task_budget`; merge beside whatever is already there. + // A `format` the caller sent natively is the more + // specific statement of the same setting and wins. + Some(serde_json::Value::Object(config)) => { + config.entry("format").or_insert(format); + } + // Not an object: Anthropic rejects the shape either way, + // and replacing it would lose what the caller meant. + Some(_) => {} + None => { + extras.insert( + "output_config".to_string(), + serde_json::json!({"format": format}), + ); + } + } + } + StructuredOutput::Tool(schema) => { + tools.get_or_insert_with(Vec::new).push(serde_json::json!({ + "name": JSON_TOOL_NAME, + "description": JSON_TOOL_DESCRIPTION, + "input_schema": schema, + })); + // Forcing the tool is what makes the reply JSON rather than a + // suggestion the model may ignore. Two things outrank it: a + // `tool_choice` the caller set themselves, and extended + // thinking, which Anthropic rejects outright beside a forced + // choice. Both leave the synthetic tool on offer under the + // model's own `auto`. + let thinking_enabled = extras + .get("thinking") + .and_then(|t| t.get("type")) + .and_then(|t| t.as_str()) + .is_some_and(|t| t != "disabled"); + if !client_set_tool_choice && !thinking_enabled { + tool_choice = Some(serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME})); + } + } } AnthropicRequest { model: upstream_model, @@ -506,6 +555,137 @@ pub fn build_request<'a>( } } +/// Whether the client stated a `tool_choice` of their own, which the +/// structured-output tool route must yield to. +/// +/// Any value counts, `"auto"` included. `auto` is not an absence of +/// intent: it is the client saying the model decides, and a client +/// running an agent loop sends it alongside its own tools on every +/// turn. Forcing the synthetic tool there would mean those tools could +/// never be called for as long as `response_format` is set — the loop +/// would simply stop working. The gateway forces only when the client +/// left the choice unstated entirely. +/// +/// An explicit JSON `null` counts as unstated: it is the wire spelling +/// of "unset" that SDKs emit for an absent optional, and nothing +/// downstream makes a choice out of it either — the translation maps it +/// to no `tool_choice` at all. Reading it as a preference would leave a +/// request that asks for JSON, forces nothing and states nothing, so +/// the model answers in prose. +pub fn tool_choice_states_a_preference(tool_choice: Option<&serde_json::Value>) -> bool { + tool_choice.is_some_and(|choice| !choice.is_null()) +} + +/// Where a request's OpenAI `response_format` lands on the Anthropic wire. +/// +/// Anthropic has two ways to get JSON out of a model and they are not +/// interchangeable: `output_config.format` constrains decoding but only +/// the newest Claude families accept it, while a forced tool call works +/// on every model that supports tools at all — including the non-Claude +/// models served behind Anthropic-compatible endpoints. The target +/// model's name picks between them; see +/// [`supports_native_structured_output`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StructuredOutput { + /// Nothing goes on the wire: the caller sent no `response_format`, or + /// sent one carrying no schema (`{"type":"json_object"}`, which has + /// no Anthropic counterpart on either path — Anthropic's JSON + /// controls are schema-driven). + None, + /// `output_config.format` — the model's own structured-output field. + Native(serde_json::Value), + /// A synthetic [`JSON_TOOL_NAME`] tool whose input *is* the answer; + /// the reply is translated back into content by + /// [`unwrap_json_tool_call`]. + Tool(serde_json::Value), +} + +/// Decide what the request's `response_format` becomes for +/// `upstream_model`. Pure, so the bridge can ask the same question again +/// on the streaming path without rebuilding the body. +pub fn structured_output_for(req: &ChatFormat, upstream_model: &str) -> StructuredOutput { + let Some(schema) = req + .extra + .get("response_format") + .and_then(response_format_schema) + else { + return StructuredOutput::None; + }; + if supports_native_structured_output(upstream_model) { + StructuredOutput::Native(schema) + } else { + StructuredOutput::Tool(schema) + } +} + +/// Pull the JSON schema out of an OpenAI `response_format`, sealed over +/// its properties. Anthropic requires every object in the schema to carry +/// `additionalProperties: false` on both paths — the native field rejects +/// an open object outright, and a tool `input_schema` that leaves one +/// open invites the model to invent members — so the schema is sealed +/// regardless of the caller's `strict` flag. +/// +/// `required` is left exactly as the caller wrote it. Anthropic treats it +/// as an ordinary JSON Schema keyword: a property left out stays optional +/// and merely sorts after the required ones in the output. Promoting +/// every property, the way OpenAI strict mode does, would make a caller's +/// optional field mandatory on this provider and nowhere else. +/// [`anthropic_output_format_to_response_format`] going the other way is +/// the one direction that does promote, because the `response_format` it +/// emits declares `strict: true`. +fn response_format_schema(response_format: &serde_json::Value) -> Option { + let mut schema = json_schema_from_response_format(response_format)?; + seal_object_schemas(&mut schema); + // Anthropic compiles the schema into a decoding grammar and 400s on + // any keyword outside its documented subset, so the constraints it + // cannot take are moved into the descriptions the model reads. + apply_schema_limits(&mut schema, &ANTHROPIC_SCHEMA_LIMITS); + Some(schema) +} + +/// Whether `model` names a Claude family that accepts Anthropic's native +/// structured-output control, `output_config.format`. That is Claude 4.5 +/// and later: `claude-{sonnet,opus,haiku}-4-5`, every `claude-*-4-6` and +/// above, and every `claude-*-5*`. +/// +/// The gateway holds only the operator-supplied upstream model name — it +/// has no capability map — so the family version is read off the name. +/// Anthropic has used two orderings (`claude-3-5-haiku-…` and +/// `claude-sonnet-4-5-…`), so the version is the first one- or two-digit +/// segment rather than a fixed position; the trailing release date is +/// eight digits and so can never be mistaken for a minor, which is what +/// keeps `claude-sonnet-4-20250514` at 4.0. `@` splits alongside `-` for +/// the `claude-sonnet-4-5@20250929` spelling. +/// +/// Everything this returns `false` for — older Claude families, unparsable +/// names, and every non-Claude name reached through an +/// Anthropic-compatible endpoint — takes the tool path, which needs no +/// capability beyond tool calling. +pub fn supports_native_structured_output(model: &str) -> bool { + claude_family_version(model).is_some_and(|version| version >= (4, 5)) +} + +fn claude_family_version(model: &str) -> Option<(u32, u32)> { + let lowered = model.trim().to_ascii_lowercase(); + let segments: Vec<&str> = lowered.split(['-', '@']).collect(); + if segments.first()? != &"claude" { + return None; + } + // A version segment is one or two digits; anything longer is a + // release date (`20250514`) or a build id, never a family number. + fn is_version(s: &str) -> bool { + (1..=2).contains(&s.len()) && s.bytes().all(|b| b.is_ascii_digit()) + } + let major_at = segments.iter().position(|s| is_version(s))?; + let major: u32 = segments[major_at].parse().ok()?; + let minor = segments + .get(major_at + 1) + .filter(|s| is_version(s)) + .and_then(|s| s.parse().ok()) + .unwrap_or(0); + Some((major, minor)) +} + /// Rewrite an OpenAI-shape `reasoning_effort` into Anthropic's effort /// control, the mirror of [`reasoning_effort_for`]. The field is always /// consumed: forwarding it verbatim reaches `/v1/messages` as an unknown @@ -1039,20 +1219,25 @@ const ANTHROPIC_DEFAULT_EFFORT: &str = "high"; /// `response_format` shape. Anthropic's structured outputs are /// constrained-decoded, so the OpenAI side is emitted with /// `strict: true` to keep that guarantee rather than degrading it to a -/// best-effort hint; strict mode in turn requires every object schema to -/// close over its properties (LiteLLM normalises the schema the same -/// way). Returns `None` for any other shape, which is then dropped. +/// best-effort hint. +/// +/// The schema itself is carried **verbatim**. Strict mode's requirement +/// that every declared property be listed in `required` is applied by +/// the OpenAI request builder, at the edge where `strict: true` actually +/// goes on the wire — doing it here would rewrite the caller's schema +/// for every downstream, and this translation also feeds the Anthropic +/// and Bedrock edges, where an optional property must stay optional. +/// Returns `None` for any other shape, which is then dropped. fn anthropic_output_format_to_response_format( output_format: serde_json::Value, ) -> Option { if output_format.get("type").and_then(|t| t.as_str())? != "json_schema" { return None; } - let mut schema = output_format.get("schema")?.clone(); + let schema = output_format.get("schema")?.clone(); if schema.is_null() { return None; } - close_object_schemas(&mut schema); Some(serde_json::json!({ "type": "json_schema", "json_schema": { @@ -1063,45 +1248,6 @@ fn anthropic_output_format_to_response_format( })) } -/// Recursively make every object schema satisfy OpenAI strict mode: -/// `additionalProperties: false`, and every declared property listed in -/// `required`. -fn close_object_schemas(schema: &mut serde_json::Value) { - let Some(obj) = schema.as_object_mut() else { - return; - }; - if obj.get("type").and_then(|t| t.as_str()) == Some("object") { - if let Some(properties) = obj.get("properties").and_then(|p| p.as_object()) { - let required: Vec = - properties.keys().map(|k| k.as_str().into()).collect(); - obj.insert("additionalProperties".to_string(), false.into()); - obj.insert("required".to_string(), required.into()); - } - if let Some(properties) = obj.get_mut("properties").and_then(|p| p.as_object_mut()) { - for property in properties.values_mut() { - close_object_schemas(property); - } - } - } - if let Some(items) = obj.get_mut("items") { - close_object_schemas(items); - } - for key in ["anyOf", "oneOf", "allOf"] { - if let Some(branches) = obj.get_mut(key).and_then(|b| b.as_array_mut()) { - for branch in branches { - close_object_schemas(branch); - } - } - } - for key in ["$defs", "definitions"] { - if let Some(defs) = obj.get_mut(key).and_then(|d| d.as_object_mut()) { - for def in defs.values_mut() { - close_object_schemas(def); - } - } - } -} - /// Non-streaming response shape from `/v1/messages`. #[derive(Debug, Deserialize)] pub struct AnthropicResponse { @@ -2655,6 +2801,10 @@ fn content_block_stop_event(index: usize) -> AnthropicSseEvent { #[cfg(test)] mod tests { + use aisix_gateway::structured_output::{ + response_into_fake_stream_chunks, unwrap_json_tool_call, + }; + use super::*; const BILLING_LINE: &str = @@ -4835,13 +4985,27 @@ mod tests { let rf = extra.get("response_format").expect("response_format set"); assert_eq!(rf["type"], serde_json::json!("json_schema")); assert_eq!(rf["json_schema"]["strict"], serde_json::json!(true)); - let schema = &rf["json_schema"]["schema"]; - // Strict mode closes every object level, not just the root. - assert_eq!(schema["additionalProperties"], serde_json::json!(false)); - assert_eq!(schema["required"], serde_json::json!(["city", "days"])); - let item = &schema["properties"]["days"]["items"]; - assert_eq!(item["additionalProperties"], serde_json::json!(false)); - assert_eq!(item["required"], serde_json::json!(["high"])); + // The schema is carried verbatim. Strict mode's closing is + // applied by the OpenAI request builder, the edge where + // `strict: true` actually goes on the wire — this normalised + // request also reaches the Anthropic, Bedrock and Gemini edges, + // where the caller's `required` is theirs to keep. + assert_eq!( + rf["json_schema"]["schema"], + serde_json::json!({ + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": { + "type": "array", + "items": { + "type": "object", + "properties": {"high": {"type": "number"}}, + }, + }, + }, + }) + ); assert!(!extra.contains_key("output_config")); } @@ -4864,8 +5028,8 @@ mod tests { .clone(); translate_extras_to_openai_shape(&mut extra, MappedEffort::AsWritten); assert_eq!( - extra["response_format"]["json_schema"]["schema"]["required"], - serde_json::json!(["legacy"]) + extra["response_format"]["json_schema"]["schema"]["properties"], + serde_json::json!({"legacy": {"type": "string"}}) ); assert!(!extra.contains_key("output_format")); } @@ -6006,4 +6170,503 @@ mod tests { serde_json::to_value(parse_inbound_request(&body).unwrap()).unwrap(), ); } + + // ── structured outputs: chat `response_format` → Anthropic ──────── + + /// A `response_format` asking for a schema, the shape both a chat + /// caller and the `/v1/responses` translation of `text.format` send. + fn json_schema_format(schema: serde_json::Value) -> serde_json::Value { + serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema, "strict": true}, + }) + } + + fn person_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "pet": {"type": "object", "properties": {"kind": {"type": "string"}}}, + }, + }) + } + + fn request_with_response_format(response_format: serde_json::Value) -> ChatFormat { + let mut req = ChatFormat::new("m", vec![ChatMessage::user("who are you")]); + req.extra.insert("response_format".into(), response_format); + req + } + + fn build<'a>(req: &'a ChatFormat, upstream_model: &'a str) -> AnthropicRequest<'a> { + let (system, messages) = split_system(req).unwrap(); + build_request(req, upstream_model, system, messages, false) + } + + #[test] + fn native_structured_output_gate_admits_4_5_and_later_only() { + // Anthropic has used two name orderings and appends a release + // date; the gate reads the family version out of both without a + // capability map. Everything it rejects takes the tool path. + for name in [ + "claude-sonnet-4-5", + "claude-sonnet-4-5-20250929", + "claude-sonnet-4-5@20250929", + "claude-opus-4-5", + "claude-haiku-4-5", + "claude-opus-4-6", + "claude-opus-4-7", + "claude-opus-4-8-20260101", + "claude-fable-5-1", + "claude-mythos-5", + "claude-opus-5", + "claude-sonnet-5", + "CLAUDE-SONNET-4-5", + ] { + assert!( + supports_native_structured_output(name), + "{name} should take the native path" + ); + } + for name in [ + // The bare "4" family: the trailing eight-digit release date + // is not a minor version. + "claude-sonnet-4-20250514", + "claude-opus-4-1", + "claude-opus-4-1-20250805", + "claude-3-7-sonnet-20250219", + "claude-3-5-haiku-20241022", + "claude-3-opus-20240229", + "claude-2.1", + "claude-instant-1.2", + "claude", + // Non-Claude names reached through an Anthropic-compatible + // endpoint, and Bedrock/Vertex-prefixed spellings. + "glm-4.5", + "deepseek-chat", + "anthropic.claude-sonnet-4-5-v1:0", + "", + ] { + assert!( + !supports_native_structured_output(name), + "{name} should take the tool path" + ); + } + } + + #[test] + fn native_path_emits_output_config_format_and_closes_the_schema() { + let req = request_with_response_format(json_schema_format(person_schema())); + let built = build(&req, "claude-sonnet-4-5"); + let format = &built.extra["output_config"]["format"]; + assert_eq!(format["type"], "json_schema"); + // Every object in the schema, nested ones included, is closed — + // Anthropic rejects an open object. + assert_eq!(format["schema"]["additionalProperties"], false); + assert_eq!( + format["schema"]["properties"]["pet"]["additionalProperties"], + false + ); + // The OpenAI spelling never reaches the body, and the native + // path adds no tool. + assert!(!built.extra.contains_key("response_format")); + assert!(built.tools.is_none()); + assert!(built.tool_choice.is_none()); + } + + #[test] + fn an_optional_property_stays_optional_on_both_paths() { + // Anthropic lists `required` as an ordinary JSON Schema keyword + // and documents optional properties explicitly, so a caller's + // optional field must not be promoted to mandatory the way + // OpenAI strict mode promotes it. + let schema = serde_json::json!({ + "type": "object", + "properties": {"name": {"type": "string"}, "nickname": {"type": "string"}}, + "required": ["name"], + }); + let req = request_with_response_format(json_schema_format(schema)); + + let native = build(&req, "claude-sonnet-4-5"); + assert_eq!( + native.extra["output_config"]["format"]["schema"]["required"], + serde_json::json!(["name"]) + ); + assert_eq!( + native.extra["output_config"]["format"]["schema"]["additionalProperties"], + false + ); + + let tool = build(&req, "claude-3-5-haiku-20241022"); + let input_schema = &tool.tools.as_ref().unwrap()[0]["input_schema"]; + assert_eq!(input_schema["required"], serde_json::json!(["name"])); + assert_eq!(input_schema["additionalProperties"], false); + } + + #[test] + fn constraints_anthropic_rejects_move_into_the_description_on_both_paths() { + // Anthropic compiles the schema into a decoding grammar and 400s + // on any keyword outside its documented subset, so a schema a + // generator produced from typed models would fail outright. The + // constraints are stated to the model instead. + let schema = serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string", "description": "full name", "maxLength": 20}, + "age": {"type": "integer", "minimum": 1}, + }, + }); + let req = request_with_response_format(json_schema_format(schema)); + + for (model, on_the_wire) in [ + ("claude-sonnet-4-5", None), + ("claude-3-5-haiku-20241022", Some(JSON_TOOL_NAME)), + ] { + let built = build(&req, model); + let sent = match on_the_wire { + None => built.extra["output_config"]["format"]["schema"].clone(), + Some(_) => built.tools.as_ref().unwrap()[0]["input_schema"].clone(), + }; + assert!( + sent["properties"]["name"].get("maxLength").is_none(), + "{model}: maxLength must not reach the wire" + ); + assert_eq!( + sent["properties"]["name"]["description"], "full name (maxLength: 20)", + "{model}" + ); + assert!( + sent["properties"]["age"].get("minimum").is_none(), + "{model}" + ); + assert_eq!(sent["properties"]["age"]["description"], "minimum: 1"); + } + } + + #[test] + fn native_format_merges_beside_a_translated_effort() { + // `output_config` is a shared carrier: the format must land + // beside the effort `reasoning_effort` translates into, not + // replace it. + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert("reasoning_effort".into(), "high".into()); + let built = build(&req, "claude-opus-4-7"); + assert_eq!(built.extra["output_config"]["effort"], "high"); + assert_eq!( + built.extra["output_config"]["format"]["type"], + "json_schema" + ); + } + + #[test] + fn native_format_yields_to_one_the_caller_sent_natively() { + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert( + "output_config".into(), + serde_json::json!({"format": {"type": "json_schema", "schema": {"type": "string"}}}), + ); + let built = build(&req, "claude-sonnet-4-5"); + assert_eq!( + built.extra["output_config"]["format"]["schema"], + serde_json::json!({"type": "string"}) + ); + } + + #[test] + fn json_object_without_a_schema_emits_nothing_on_either_path() { + // Anthropic's JSON controls are schema-driven on both paths, so + // a schemaless `json_object` has nothing to translate into. It + // is still consumed — forwarding it would 400 upstream. + for model in ["claude-sonnet-4-5", "claude-3-5-haiku-20241022"] { + let req = request_with_response_format(serde_json::json!({"type": "json_object"})); + let built = build(&req, model); + assert!(!built.extra.contains_key("response_format")); + assert!(!built.extra.contains_key("output_config")); + assert!(built.tools.is_none()); + assert!(built.tool_choice.is_none()); + } + } + + #[test] + fn tool_path_appends_the_synthetic_tool_and_forces_it() { + let req = request_with_response_format(json_schema_format(person_schema())); + let built = build(&req, "claude-3-5-haiku-20241022"); + let tools = built.tools.as_ref().expect("synthetic tool on the wire"); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["name"], JSON_TOOL_NAME); + assert!(tools[0]["description"].as_str().unwrap().contains("JSON")); + assert_eq!(tools[0]["input_schema"]["additionalProperties"], false); + assert_eq!( + built.tool_choice, + Some(serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME})) + ); + assert!(!built.extra.contains_key("response_format")); + assert!(!built.extra.contains_key("output_config")); + } + + #[test] + fn tool_path_keeps_the_callers_own_tools() { + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert( + "tools".into(), + serde_json::json!([{ + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + }]), + ); + let built = build(&req, "glm-4.5"); + let tools = built.tools.as_ref().unwrap(); + assert_eq!(tools.len(), 2); + assert_eq!(tools[0]["name"], "get_weather"); + assert_eq!(tools[1]["name"], JSON_TOOL_NAME); + } + + #[test] + fn a_tool_choice_the_caller_sent_outranks_the_forced_json_tool() { + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert( + "tools".into(), + serde_json::json!([{ + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + }]), + ); + req.extra.insert( + "tool_choice".into(), + serde_json::json!({"type": "function", "function": {"name": "get_weather"}}), + ); + let built = build(&req, "glm-4.5"); + assert_eq!( + built.tool_choice, + Some(serde_json::json!({"type": "tool", "name": "get_weather"})) + ); + } + + #[test] + fn any_tool_choice_the_client_sent_outranks_forcing_the_json_tool() { + // `auto` included. A client running an agent loop sends it + // beside its own tools every turn; forcing the synthetic tool + // there would mean those tools could never be called for as long + // as `response_format` is set. The tool is still offered, so the + // model can reach the JSON on its own. + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert( + "tools".into(), + serde_json::json!([{ + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + }]), + ); + for stated in [ + serde_json::json!("auto"), + serde_json::json!("required"), + serde_json::json!("none"), + serde_json::json!({"type": "function", "function": {"name": "get_weather"}}), + ] { + req.extra.insert("tool_choice".into(), stated.clone()); + let built = build(&req, "glm-4.5"); + assert_ne!( + built.tool_choice, + Some(serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME})), + "tool_choice {stated} must not be overridden" + ); + assert!( + built + .tools + .as_ref() + .unwrap() + .iter() + .any(|t| t["name"] == JSON_TOOL_NAME), + "the synthetic tool is still on offer for {stated}" + ); + } + + // With no choice stated at all, the gateway forces — and an + // explicit JSON `null` is the wire spelling of unstated, which + // SDKs emit for an absent optional. + for unstated in [None, Some(serde_json::Value::Null)] { + match unstated { + Some(v) => req.extra.insert("tool_choice".into(), v), + None => req.extra.remove("tool_choice"), + }; + let built = build(&req, "glm-4.5"); + assert_eq!( + built.tool_choice, + Some(serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME})), + "an unstated tool_choice must not suppress the forcing" + ); + } + } + + #[test] + fn extended_thinking_leaves_the_synthetic_tool_on_auto() { + // Anthropic rejects a forced tool choice beside extended + // thinking, so the tool is offered rather than forced. + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra.insert( + "thinking".into(), + serde_json::json!({"type": "enabled", "budget_tokens": 2048}), + ); + let built = build(&req, "claude-3-7-sonnet-20250219"); + assert_eq!(built.tools.as_ref().unwrap()[0]["name"], JSON_TOOL_NAME); + assert!(built.tool_choice.is_none()); + + // Thinking the caller switched off is no obstacle. + let mut req = request_with_response_format(json_schema_format(person_schema())); + req.extra + .insert("thinking".into(), serde_json::json!({"type": "disabled"})); + let built = build(&req, "claude-3-7-sonnet-20250219"); + assert_eq!( + built.tool_choice, + Some(serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME})) + ); + } + + /// The bridge's own decode of an upstream reply that called the + /// synthetic tool, plus any real tool calls the model made too. + fn synthetic_tool_reply(extra_blocks: serde_json::Value) -> ChatResponse { + let mut content = vec![serde_json::json!({ + "type": "tool_use", + "id": "toolu_json", + "name": JSON_TOOL_NAME, + "input": {"name": "Ada"}, + })]; + content.extend(extra_blocks.as_array().unwrap().iter().cloned()); + let body = serde_json::json!({ + "id": "msg_json_01", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": content, + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + }); + response_into_chat_response(serde_json::from_value(body).unwrap()) + } + + #[test] + fn unwrapping_the_only_synthetic_call_yields_a_plain_json_completion() { + let mut resp = synthetic_tool_reply(serde_json::json!([])); + unwrap_json_tool_call(&mut resp); + assert_eq!(resp.message.content.as_deref(), Some(r#"{"name":"Ada"}"#)); + assert!(!resp.message.extra.contains_key("tool_calls")); + // A client that never offered a tool must not be told the model + // stopped to call one. + assert_eq!(resp.finish_reason, FinishReason::Stop); + } + + #[test] + fn a_prose_preamble_never_survives_into_the_json_answer() { + // The tool is often offered rather than forced (a caller's own + // `tool_choice`, extended thinking, a family with no forced + // choice), and a model that narrates before calling it would + // otherwise hand the caller a string that is not JSON. + let body = serde_json::json!({ + "id": "msg_preamble", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [ + {"type": "text", "text": "Sure, here you go:"}, + {"type": "tool_use", "id": "toolu_json", "name": JSON_TOOL_NAME, + "input": {"name": "Ada"}}, + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + }); + let mut resp = response_into_chat_response(serde_json::from_value(body).unwrap()); + unwrap_json_tool_call(&mut resp); + let content = resp.message.content.as_deref().unwrap(); + assert_eq!(content, r#"{"name":"Ada"}"#); + serde_json::from_str::(content).expect("content parses as JSON"); + } + + #[test] + fn fake_streamed_tool_calls_carry_a_dense_index() { + // The streaming shape needs `index`; the non-streaming decode + // this is built from does not emit one, and the SSE re-encoder + // folds every index-less call onto content block 0. + let mut resp = synthetic_tool_reply(serde_json::json!([ + {"type": "tool_use", "id": "toolu_a", "name": "get_weather", "input": {"city": "SF"}}, + {"type": "tool_use", "id": "toolu_b", "name": "get_time", "input": {"tz": "UTC"}}, + ])); + unwrap_json_tool_call(&mut resp); + let chunks = response_into_fake_stream_chunks(resp); + let calls = chunks[1].delta.tool_calls.as_ref().unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0]["index"], 0); + assert_eq!(calls[1]["index"], 1); + assert_eq!(calls[0]["function"]["name"], "get_weather"); + assert_eq!(calls[1]["function"]["name"], "get_time"); + } + + #[test] + fn unwrapping_beside_a_real_call_keeps_the_real_call_and_its_finish_reason() { + let mut resp = synthetic_tool_reply(serde_json::json!([ + {"type": "text", "text": "checking"}, + {"type": "tool_use", "id": "toolu_w", "name": "get_weather", "input": {"city": "SF"}}, + ])); + unwrap_json_tool_call(&mut resp); + let calls = resp.message.extra["tool_calls"].as_array().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0]["function"]["name"], "get_weather"); + assert_eq!( + resp.message.content.as_deref(), + Some("checking\n{\"name\":\"Ada\"}") + ); + assert_eq!(resp.finish_reason, FinishReason::ToolCalls); + } + + #[test] + fn unwrapping_leaves_a_response_without_the_synthetic_call_alone() { + let body = serde_json::json!({ + "id": "msg_plain", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [ + {"type": "tool_use", "id": "toolu_w", "name": "get_weather", "input": {}}, + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }); + let mut resp = response_into_chat_response(serde_json::from_value(body).unwrap()); + let before = serde_json::to_value(&resp).unwrap(); + unwrap_json_tool_call(&mut resp); + assert_eq!(serde_json::to_value(&resp).unwrap(), before); + } + + #[test] + fn fake_stream_emits_role_content_finish_and_usage_in_order() { + let mut resp = synthetic_tool_reply(serde_json::json!([])); + unwrap_json_tool_call(&mut resp); + let usage = resp.usage.clone(); + let chunks = response_into_fake_stream_chunks(resp); + assert_eq!(chunks.len(), 4); + assert!(chunks.iter().all(|c| c.id == "msg_json_01")); + assert_eq!(chunks[0].delta.role, Some(Role::Assistant)); + assert!(chunks[0].delta.content.is_none()); + assert_eq!( + chunks[1].delta.content.as_deref(), + Some(r#"{"name":"Ada"}"#) + ); + assert!(chunks[1].delta.tool_calls.is_none()); + assert_eq!(chunks[2].finish_reason, Some(FinishReason::Stop)); + assert!(chunks[0..3].iter().all(|c| c.usage.is_none())); + assert_eq!(chunks[3].usage, Some(usage)); + assert!(chunks[3].finish_reason.is_none()); + } + + #[test] + fn fake_stream_carries_real_tool_calls_through() { + let mut resp = synthetic_tool_reply(serde_json::json!([ + {"type": "tool_use", "id": "toolu_w", "name": "get_weather", "input": {"city": "SF"}}, + ])); + unwrap_json_tool_call(&mut resp); + let chunks = response_into_fake_stream_chunks(resp); + let calls = chunks[1].delta.tool_calls.as_ref().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0]["function"]["name"], "get_weather"); + assert_eq!(chunks[2].finish_reason, Some(FinishReason::ToolCalls)); + } } diff --git a/crates/aisix-provider-azure-openai/Cargo.toml b/crates/aisix-provider-azure-openai/Cargo.toml index 1cbd672d..cd7d92ba 100644 --- a/crates/aisix-provider-azure-openai/Cargo.toml +++ b/crates/aisix-provider-azure-openai/Cargo.toml @@ -30,6 +30,10 @@ bytes.workspace = true http.workspace = true [dev-dependencies] +# Only for the strict-structured-output test: the `/v1/messages` inbound +# translation lives in the Anthropic crate, and this edge must emit for +# it exactly what the OpenAI edge does. +aisix-provider-anthropic = { path = "../aisix-provider-anthropic" } # `net` + `io-util` back the hand-rolled SSE upstream in # `chat_stream_delivers_a_long_stream_whose_gaps_stay_within_budget`, which # needs to emit frames on a schedule — wiremock can only delay a response as diff --git a/crates/aisix-provider-azure-openai/src/bridge.rs b/crates/aisix-provider-azure-openai/src/bridge.rs index 1656241b..4fdbec1a 100644 --- a/crates/aisix-provider-azure-openai/src/bridge.rs +++ b/crates/aisix-provider-azure-openai/src/bridge.rs @@ -29,6 +29,7 @@ use reqwest::{header, Client, StatusCode}; use serde_json::Value; use std::time::{Duration, Instant}; +use aisix_provider_openai::close_strict_response_format_schema; use aisix_provider_openai::overrides::{ apply_content_list_to_string, apply_default_body_fields, apply_param_constraints, apply_param_renames, apply_stream_done_marker_policy, extract_reasoning_field, @@ -594,6 +595,10 @@ fn prepare_outbound_body( ) -> Result { let mut body = serde_json::to_value(typed) .map_err(|e| BridgeError::Config(format!("serialize request body: {e}")))?; + // Azure serves the OpenAI wire, so it owes the same strict-mode + // schema closing the OpenAI edge applies — one function, not a + // second copy, because these two bodies have to stay identical. + close_strict_response_format_schema(&mut body); if let Some(r) = request { apply_param_renames(&mut body, &r.param_renames); if let Some(constraints) = &r.param_constraints { @@ -939,6 +944,76 @@ fn parse_stream_chunk( mod tests { use super::*; + /// Azure serves the OpenAI wire, so an Anthropic-shaped + /// `/v1/messages` request translated onto an Azure deployment must + /// arrive with exactly the schema the OpenAI edge would send. Azure + /// keeps its own copy of the outbound-body pipeline, which is + /// precisely how it came to be missing the strict-mode closing. + #[test] + fn a_translated_messages_request_reaches_azure_byte_for_byte_as_before() { + use aisix_gateway::{ChatFormat, ChatMessage}; + use aisix_provider_anthropic::wire::translate_extras_to_openai_shape; + + let mut extra = serde_json::json!({ + "output_config": { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": { + "type": "array", + "items": { + "type": "object", + "properties": {"high": {"type": "number"}}, + }, + }, + }, + }, + } + } + }) + .as_object() + .unwrap() + .clone(); + translate_extras_to_openai_shape(&mut extra, aisix_core::MappedEffort::AsWritten); + + let mut req = ChatFormat::new("m", vec![ChatMessage::user("weather?")]); + req.extra = extra; + let messages = messages_from(&req); + let typed = build_request(&req, "ci-chat", &messages, false); + let body = prepare_outbound_body(&typed, None, None).unwrap(); + + assert_eq!( + body["response_format"], + serde_json::json!({ + "type": "json_schema", + "json_schema": { + "name": "structured_output", + "strict": true, + "schema": { + "type": "object", + "additionalProperties": false, + "required": ["city", "days"], + "properties": { + "city": {"type": "string"}, + "days": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["high"], + "properties": {"high": {"type": "number"}}, + }, + }, + }, + }, + }, + }) + ); + } + /// AISIX-Cloud#1222 scenario 3: an in-band `data: {"error":{...}}` /// frame inside the committed 200 stream surfaces as the typed /// in-band error, not a serde decode failure. diff --git a/crates/aisix-provider-bedrock/src/bridge.rs b/crates/aisix-provider-bedrock/src/bridge.rs index 9bb4f2c8..02ecb3e2 100644 --- a/crates/aisix-provider-bedrock/src/bridge.rs +++ b/crates/aisix-provider-bedrock/src/bridge.rs @@ -16,6 +16,9 @@ //! set) is forwarded as the SDK's `endpoint_url` so operators can //! point at a private deployment / VPC endpoint. +use aisix_gateway::structured_output::{ + response_into_fake_stream_chunks, unwrap_json_tool_call, JSON_TOOL_DESCRIPTION, JSON_TOOL_NAME, +}; use aisix_gateway::{ Bridge, BridgeContext, BridgeError, ChatChunk, ChatChunkStream, ChatDelta, ChatFormat, ChatMessage, ChatResponse, EmbeddingObject, EmbeddingRequest, EmbeddingResponse, @@ -33,7 +36,8 @@ use aws_sdk_bedrockruntime::primitives::Blob; use aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError; use aws_sdk_bedrockruntime::types::{ AnyToolChoice, ContentBlock, ContentBlockDelta, ContentBlockStart, ConversationRole, - ConverseStreamOutput, InferenceConfiguration, Message as BedrockMessage, SpecificToolChoice, + ConverseStreamOutput, InferenceConfiguration, JsonSchemaDefinition, Message as BedrockMessage, + OutputConfig, OutputFormat, OutputFormatStructure, OutputFormatType, SpecificToolChoice, StopReason as SdkStopReason, SystemContentBlock, Tool, ToolChoice, ToolConfiguration, ToolInputSchema, ToolResultBlock, ToolResultContentBlock, ToolSpecification, ToolUseBlock, }; @@ -44,8 +48,9 @@ use serde::Deserialize; use std::time::{Duration, Instant}; use aisix_provider_anthropic::wire::{ - build_request, response_into_chat_response, split_system, - translate_reasoning_effort_to_anthropic, AnthropicResponse, + build_request, response_into_chat_response, split_system, structured_output_for, + tool_choice_states_a_preference, translate_reasoning_effort_to_anthropic, AnthropicResponse, + StructuredOutput, }; // Per-`ProviderKey` request override pipeline (#302 §5 / #340). The JSON-body @@ -722,6 +727,28 @@ impl Bridge for BedrockBridge { (optionally prefixed with a cross-region inference profile like us. / eu. / apac.)" )) })?; + // The tool route's JSON only exists once the synthetic tool call + // has been assembled, so it cannot be streamed as it arrives. + // Run the request non-streaming — through `chat`, so each + // publisher keeps the wire it answers on — and fake-stream the + // translated result: the client sees an ordinary chunk sequence, + // and usage rides its own terminal chunk exactly as on a real + // stream. + if matches!( + bedrock_structured_output(req, upstream_id), + StructuredOutput::Tool(_) + ) { + // The leg is not streaming, so it runs under the budget a + // non-streaming call would have got — the streaming budget + // this context carries bounds a chunk gap, not a completion. + let chunks = + response_into_fake_stream_chunks(self.chat(req, &ctx.non_streaming_ctx()).await?); + return Ok(Box::pin(async_stream::stream! { + for chunk in chunks { + yield Ok(chunk); + } + })); + } // Phase G productionization (#302 Step 3): unified Converse // stream path for all publishers — same SDK call (.converse_stream) // owns the AWS event-stream binary frame decoding internally, @@ -863,7 +890,17 @@ impl BedrockBridge { let (system, messages) = split_system(req).map_err(|e| BridgeError::InvalidUpstreamConfig(format!("{e}")))?; - let anthropic_req = build_request(req, upstream_id, system, messages, false); + // The Anthropic request builder reads the Claude family version + // off the model name to pick a structured-output shape, and + // `anthropic.claude-…` is not a name it can read. Hand it the + // Claude name the Bedrock id carries; the `model` field it lands + // in is stripped below, since /invoke keys the model off the URL. + let anthropic_model = bedrock_claude_model_name(upstream_id).unwrap_or(upstream_id); + let synthetic_json_tool = matches!( + structured_output_for(req, anthropic_model), + StructuredOutput::Tool(_) + ); + let anthropic_req = build_request(req, anthropic_model, system, messages, false); let mut body_value = serde_json::to_value(&anthropic_req) .map_err(|e| BridgeError::Config(format!("serialize Anthropic request body: {e}")))?; // Apply the per-ProviderKey body override pipeline (#340) BEFORE the @@ -897,7 +934,11 @@ impl BedrockBridge { let parsed: AnthropicResponse = serde_json::from_slice(resp.body().as_ref()) .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; - Ok(response_into_chat_response(parsed)) + let mut chat = response_into_chat_response(parsed); + if synthetic_json_tool { + unwrap_json_tool_call(&mut chat); + } + Ok(chat) } /// Resolve credentials + build an SDK client. Pulled out of @@ -982,11 +1023,28 @@ impl BedrockBridge { if let Some(fields) = build_converse_additional_model_request_fields(req, upstream_id) { call = call.additional_model_request_fields(fields); } + // A caller's `response_format` becomes one of two request shapes + // here, picked by whether the model constrains its own decoding. + let structured = bedrock_structured_output(req, upstream_id); + let mut json_tool_schema = None; + match &structured { + StructuredOutput::None => {} + StructuredOutput::Native(schema) => { + if let Some(cfg) = build_output_config(req, schema) { + call = call.output_config(cfg); + } + } + StructuredOutput::Tool(schema) => json_tool_schema = Some(schema), + } // #560: forward OpenAI `tools` / `tool_choice` into Converse's // `toolConfig`. Without this every Converse publisher silently // drops tool calling and improvises the call as prose // (finish_reason=stop, tool_calls=[]). - if let Some(tc) = build_tool_config(req) { + if let Some(tc) = build_tool_config( + req, + json_tool_schema, + BedrockPublisher::from_model_id(upstream_id), + ) { call = call.tool_config(tc); } @@ -994,7 +1052,11 @@ impl BedrockBridge { .send() .await .map_err(|e| map_converse_sdk_error(e, started, deadline))?; - Ok(converse_output_into_chat_response(resp, upstream_id)) + let mut chat = converse_output_into_chat_response(resp, upstream_id); + if json_tool_schema.is_some() { + unwrap_json_tool_call(&mut chat); + } + Ok(chat) } /// Dispatch Bedrock chat via the unified Converse stream API. @@ -1055,9 +1117,18 @@ impl BedrockBridge { if let Some(fields) = build_converse_additional_model_request_fields(req, upstream_id) { call = call.additional_model_request_fields(fields); } + // Only the native structured-output shape reaches this path — + // `chat_stream` diverts the tool route before it gets here, + // because a tool call cannot be streamed before it is complete. + if let StructuredOutput::Native(schema) = bedrock_structured_output(req, upstream_id) { + if let Some(cfg) = build_output_config(req, &schema) { + call = call.output_config(cfg); + } + } // #560: forward tools on the stream path too (all publishers, // incl. Anthropic, stream through Converse). - if let Some(tc) = build_tool_config(req) { + if let Some(tc) = build_tool_config(req, None, BedrockPublisher::from_model_id(upstream_id)) + { call = call.tool_config(tc); } @@ -1798,6 +1869,128 @@ fn build_converse_additional_model_request_fields( } } +/// Decide what a request's OpenAI `response_format` becomes on Bedrock. +/// +/// The decision is the Anthropic crate's — same gate, same schema +/// sealing — asked with the Claude model name Bedrock's id *carries* +/// rather than the id itself, which that gate cannot read. Every model +/// that is not a Claude one keeps [`StructuredOutput::Tool`], which +/// needs no capability beyond tool calling; AWS's native structured +/// output on Converse covers the same Claude families Anthropic's own +/// API does. +fn bedrock_structured_output(req: &ChatFormat, upstream_id: &str) -> StructuredOutput { + let decided = structured_output_for( + req, + bedrock_claude_model_name(upstream_id).unwrap_or(upstream_id), + ); + // The tool route only exists where the model can call a tool at + // all. Attaching a `toolConfig` to a publisher whose Converse + // implementation has none fails the whole request — and it would + // fail a request carrying no tools of its own, purely because the + // caller asked for JSON. Leaving the field unhonoured is the lesser + // outcome, and is what these models did before. + if matches!(decided, StructuredOutput::Tool(_)) + && !converse_supports_tool_use(BedrockPublisher::from_model_id(upstream_id)) + { + tracing::debug!( + model = %upstream_id, + "dropping response_format: this Bedrock publisher supports neither native structured output nor tool use" + ); + return StructuredOutput::None; + } + decided +} + +/// Whether this publisher's Converse implementation supports tool use at +/// all, per AWS's supported-model table. Titan Text has none, and +/// `Other` is the set this bridge has not classified — several of which +/// (DeepSeek R1 among them) also reject `toolConfig`. +/// +/// Only the *synthetic* tool is gated on this. A caller who sent their +/// own `tools` still gets them forwarded: an explicit unsupported tool +/// request is theirs to see rejected, which is what #560 shipped. +fn converse_supports_tool_use(publisher: Option) -> bool { + matches!( + publisher, + Some(BedrockPublisher::Anthropic) + | Some(BedrockPublisher::AmazonNova) + | Some(BedrockPublisher::Meta) + | Some(BedrockPublisher::Mistral) + | Some(BedrockPublisher::Cohere) + ) +} + +/// Pull the Anthropic model name out of a Bedrock model id: +/// `anthropic.claude-sonnet-4-5-20250929-v1:0` and its region-prefixed +/// (`us.anthropic.…`) and ARN (`arn:…:inference-profile/us.anthropic.…`) +/// spellings all yield `claude-sonnet-4-5-20250929-v1:0`. +/// +/// `None` for every other id — another publisher, or a +/// provisioned-model / application-inference-profile ARN whose last +/// segment is an opaque id. An opaque id names a model the gateway +/// cannot identify, so it takes the tool path, which works regardless. +fn bedrock_claude_model_name(model_id: &str) -> Option<&str> { + // An inference-profile or provisioned-model ARN carries the model id + // in its last path segment; a bare id has no `/` and is unchanged. + let tail = model_id.rsplit('/').next().unwrap_or(model_id); + let (tag, rest) = strip_region_prefix(tail).split_once('.')?; + if !tag.eq_ignore_ascii_case("anthropic") { + return None; + } + rest.get(..6) + .is_some_and(|head| head.eq_ignore_ascii_case("claude")) + .then_some(rest) +} + +/// The `outputConfig` carrying a schema Bedrock constrains decoding to. +/// +/// Converse takes the schema as a JSON **string**, not as an object — +/// `jsonSchema.schema` is typed `String` in the API. `name` is the +/// caller's own schema name where they sent one, since it is what they +/// see quoted back in a schema-compilation error. +fn build_output_config(req: &ChatFormat, schema: &serde_json::Value) -> Option { + let json_schema = req.extra.get("response_format")?.get("json_schema"); + let mut definition = JsonSchemaDefinition::builder().schema(schema.to_string()); + definition = definition.name( + json_schema + .and_then(|j| j.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or(DEFAULT_JSON_SCHEMA_NAME), + ); + if let Some(description) = json_schema + .and_then(|j| j.get("description")) + .and_then(|d| d.as_str()) + { + definition = definition.description(description); + } + let format = OutputFormat::builder() + .r#type(OutputFormatType::JsonSchema) + .structure(OutputFormatStructure::JsonSchema(definition.build().ok()?)) + .build() + .ok()?; + Some(OutputConfig::builder().text_format(format).build()) +} + +/// Schema name sent when the caller did not name theirs. `name` is +/// required on the wire even though the SDK models it as optional. +const DEFAULT_JSON_SCHEMA_NAME: &str = "structured_output"; + +/// Whether Converse honours an explicit `toolChoice` for this +/// publisher. Per AWS, forcing a specific tool is Anthropic Claude and +/// Amazon Nova only; every other publisher rejects the field, so the +/// synthetic tool is merely offered there and the model is trusted to +/// take the only tool on the table. +/// +/// This gates only the forcing the *gateway* adds. A `tool_choice` the +/// caller set themselves still goes through unchanged — an explicit +/// unsupported force is theirs to see rejected. +fn converse_honours_forced_tool_choice(publisher: Option) -> bool { + matches!( + publisher, + Some(BedrockPublisher::Anthropic) | Some(BedrockPublisher::AmazonNova) + ) +} + /// Build a Converse [`ToolConfiguration`] from the OpenAI `tools` / /// `tool_choice` the caller sent — they arrive in [`ChatFormat::extra`] /// (the gateway captures unknown top-level request fields there). Returns @@ -1819,15 +2012,49 @@ fn build_converse_additional_model_request_fields( /// References: /// , /// -fn build_tool_config(req: &ChatFormat) -> Option { +/// +/// `json_tool_schema` is the structured-output tool route: a schema to +/// append as the synthetic [`JSON_TOOL_NAME`] tool and force the model +/// onto. Forcing is what makes the reply JSON rather than a suggestion +/// the model may ignore, and three things outrank it — a `tool_choice` +/// the caller set themselves, extended thinking (which Anthropic rejects +/// beside a forced choice), and a publisher whose Converse +/// implementation has no `toolChoice` at all. Each leaves the synthetic +/// tool on offer under the model's own `auto`. +fn build_tool_config( + req: &ChatFormat, + json_tool_schema: Option<&serde_json::Value>, + publisher: Option, +) -> Option { // OpenAI `tool_choice:"none"` means "don't call any tool this turn". // Converse has no `none`, so send no toolConfig — the model answers // in prose, honouring the user-visible contract (no tool call). Tool // visibility is lost, but that matches the intent of "none". - if req.extra.get("tool_choice").and_then(|v| v.as_str()) == Some("none") { + // + // A `response_format` asking for JSON contradicts it: the JSON has + // to come out of a tool call on this route. The schema wins — the + // caller's `response_format` is the more specific statement about + // what the answer must be — and the contradiction still costs the + // forcing, so the model is only offered the tool. + let caller_forbade_tools = + req.extra.get("tool_choice").and_then(|v| v.as_str()) == Some("none"); + if caller_forbade_tools && json_tool_schema.is_none() { return None; } - let tools_json = req.extra.get("tools").and_then(|v| v.as_array())?; + // With `response_format` beside it the schema wins — the JSON has to + // come out of a tool call on this route — but only the synthetic + // tool goes on the table. Putting the caller's own tools back under + // no `toolChoice` would hand the model exactly what "none" told it + // not to use. + let tools_json = if caller_forbade_tools { + &[][..] + } else { + req.extra + .get("tools") + .and_then(|v| v.as_array()) + .map(Vec::as_slice) + .unwrap_or_default() + }; let mut tools: Vec = Vec::new(); for entry in tools_json { // OpenAI only defines `type:"function"` tools today; skip any @@ -1863,11 +2090,36 @@ fn build_tool_config(req: &ChatFormat) -> Option { tools.push(Tool::ToolSpec(spec)); } } + if let Some(schema) = json_tool_schema { + if let Ok(spec) = ToolSpecification::builder() + .name(JSON_TOOL_NAME) + .description(JSON_TOOL_DESCRIPTION) + .input_schema(ToolInputSchema::Json(json_to_document(schema))) + .build() + { + tools.push(Tool::ToolSpec(spec)); + } + } if tools.is_empty() { return None; } let mut config = ToolConfiguration::builder().set_tools(Some(tools)); - if let Some(choice) = req.extra.get("tool_choice").and_then(map_tool_choice) { + let client_set_tool_choice = tool_choice_states_a_preference(req.extra.get("tool_choice")); + let thinking_enabled = req + .extra + .get("thinking") + .and_then(|t| t.get("type")) + .and_then(|t| t.as_str()) + .is_some_and(|t| t != "disabled"); + let force_json_tool = json_tool_schema.is_some() + && !client_set_tool_choice + && !thinking_enabled + && converse_honours_forced_tool_choice(publisher); + if force_json_tool { + if let Ok(choice) = SpecificToolChoice::builder().name(JSON_TOOL_NAME).build() { + config = config.tool_choice(ToolChoice::Tool(choice)); + } + } else if let Some(choice) = req.extra.get("tool_choice").and_then(map_tool_choice) { config = config.tool_choice(choice); } config.build().ok() @@ -3657,6 +3909,620 @@ mod tests { ); } + // ─── Structured outputs: chat `response_format` → Bedrock ───────── + + /// The schema shape a caller sends, and the Bedrock spellings of one + /// Claude model id. + fn structured_request(schema: serde_json::Value) -> ChatFormat { + let mut req = ChatFormat::new("my-model", vec![ChatMessage::user("who are you")]); + req.extra.insert( + "response_format".into(), + serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "person", "schema": schema, "strict": true}, + }), + ); + req + } + + fn person_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "nickname": {"type": "string"}, + }, + "required": ["name"], + }) + } + + /// Drive the real dispatch against wiremock and hand back what + /// reached the wire. `route` is the Bedrock operation suffix + /// (`invoke`, `converse`, `converse-stream`); the canned reply is + /// deliberately not a valid one for every route — the request is + /// captured before the response is decoded. + async fn capture_bedrock_body( + model_id: &str, + route: &str, + req: &ChatFormat, + streaming: bool, + ) -> serde_json::Value { + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(format!(r"^/model/.+/{route}$"))) + .respond_with(responder.clone()) + .expect(1) + .mount(&server) + .await; + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with(model_id), + sample_pk_with_secret(valid_secret_json()), + ); + if streaming { + let _ = bridge.chat_stream(req, &ctx).await; + } else { + let _ = bridge.chat(req, &ctx).await; + } + let body = responder.captured_body.lock().unwrap().clone(); + body.unwrap_or_else(|| panic!("no {route} request was captured")) + } + + #[test] + fn native_gate_reads_the_claude_name_out_of_the_bedrock_model_id() { + // Bedrock wraps the Anthropic model name in a publisher tag, an + // optional cross-region prefix and (for profiles) an ARN path. + // The native structured-output gate has to see through all three + // or every Claude on Bedrock silently takes the tool path. + let req = structured_request(person_schema()); + for id in [ + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "eu.anthropic.claude-opus-4-5-20251101-v1:0", + "arn:aws:bedrock:us-east-1:1234:inference-profile/us.anthropic.claude-sonnet-4-5-v1:0", + ] { + assert!( + matches!( + bedrock_structured_output(&req, id), + StructuredOutput::Native(_) + ), + "{id} should take the native path" + ); + } + for id in [ + // Claude, but older than the families that constrain their + // own decoding — the trailing release date is not a minor. + "anthropic.claude-sonnet-4-20250514-v1:0", + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-haiku-20240307-v1:0", + // Other publishers whose Converse supports tool use. + "amazon.nova-pro-v1:0", + "meta.llama3-3-70b-instruct-v1:0", + "mistral.mistral-large-2407-v1:0", + "cohere.command-r-plus-v1:0", + ] { + assert!( + matches!( + bedrock_structured_output(&req, id), + StructuredOutput::Tool(_) + ), + "{id} should take the tool path" + ); + } + for id in [ + // No native structured output AND no Converse tool use, so + // there is no shape to translate into. Dropping the field + // keeps these models answering as they did before rather + // than failing the request on a `toolConfig` they reject. + "amazon.titan-text-express-v1", + "deepseek.r1-v1:0", + "ai21.jamba-1-5-large-v1:0", + // A profile ARN whose last segment is an opaque id naming no + // model this gateway can identify. + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcd1234", + "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/xyz", + ] { + assert!( + matches!(bedrock_structured_output(&req, id), StructuredOutput::None), + "{id} has no structured-output shape and must drop the field" + ); + } + } + + #[tokio::test] + async fn invoke_path_carries_the_schema_in_output_config_for_a_4_5_claude() { + // Claude on Bedrock answers non-streaming over the Anthropic + // Messages wire at /invoke, whose native control is + // `output_config.format` — AWS documents it on that route. + let body = capture_bedrock_body( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "invoke", + &structured_request(person_schema()), + false, + ) + .await; + let format = &body["output_config"]["format"]; + assert_eq!(format["type"], "json_schema"); + assert_eq!(format["schema"]["additionalProperties"], false); + // A caller's optional property stays optional. + assert_eq!(format["schema"]["required"], serde_json::json!(["name"])); + assert!(body.get("response_format").is_none(), "body={body}"); + assert!(body.get("tools").is_none(), "body={body}"); + } + + #[tokio::test] + async fn a_messages_request_to_bedrock_claude_keeps_its_optional_properties() { + // The `/v1/messages` inbound translation used to apply OpenAI + // strict mode's all-required promotion before any bridge saw the + // request, which made a caller's optional property mandatory on + // this edge. The promotion now happens at the OpenAI edge only. + let mut extra = serde_json::json!({ + "output_config": { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "nickname": {"type": "string"}, + }, + "required": ["name"], + }, + } + } + }) + .as_object() + .unwrap() + .clone(); + aisix_provider_anthropic::wire::translate_extras_to_openai_shape( + &mut extra, + aisix_core::MappedEffort::AsWritten, + ); + let mut req = ChatFormat::new("my-model", vec![ChatMessage::user("who are you")]); + req.extra = extra; + + let body = capture_bedrock_body( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "invoke", + &req, + false, + ) + .await; + let schema = &body["output_config"]["format"]["schema"]; + assert_eq!(schema["required"], serde_json::json!(["name"])); + // Sealing still happens — Bedrock rejects an open object. + assert_eq!(schema["additionalProperties"], false); + } + + #[tokio::test] + async fn invoke_path_falls_back_to_the_synthetic_tool_on_an_older_claude() { + let body = capture_bedrock_body( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "invoke", + &structured_request(person_schema()), + false, + ) + .await; + let tools = body["tools"] + .as_array() + .expect("synthetic tool on the wire"); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["name"], JSON_TOOL_NAME); + assert_eq!(tools[0]["input_schema"]["additionalProperties"], false); + assert_eq!( + body["tool_choice"], + serde_json::json!({"type": "tool", "name": JSON_TOOL_NAME}) + ); + assert!(body.get("output_config").is_none(), "body={body}"); + } + + #[tokio::test] + async fn converse_stream_carries_the_schema_as_a_json_string_in_output_config() { + // Streaming goes through Converse for every publisher, and + // Converse takes the schema as a STRING, not as an object. + let body = capture_bedrock_body( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "converse-stream", + &structured_request(person_schema()), + true, + ) + .await; + let text_format = &body["outputConfig"]["textFormat"]; + assert_eq!(text_format["type"], "json_schema"); + let definition = &text_format["structure"]["jsonSchema"]; + // The caller's own schema name, so a compilation error quotes + // something they recognise. + assert_eq!(definition["name"], "person"); + let schema: serde_json::Value = serde_json::from_str( + definition["schema"] + .as_str() + .expect("schema is a JSON string"), + ) + .expect("schema string parses back"); + assert_eq!(schema["additionalProperties"], false); + assert_eq!(schema["required"], serde_json::json!(["name"])); + assert!(body.get("toolConfig").is_none(), "body={body}"); + } + + #[tokio::test] + async fn converse_output_config_carries_the_narrowed_schema() { + // Bedrock documents the same unsupported-keyword set as + // Anthropic for its structured outputs, so what goes into the + // `outputConfig` string has already been narrowed. + let schema = serde_json::json!({ + "type": "object", + "properties": {"name": {"type": "string", "maxLength": 20}}, + }); + let mut req = structured_request(schema); + req.stream = Some(true); + let body = capture_bedrock_body( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "converse-stream", + &req, + true, + ) + .await; + let sent: serde_json::Value = serde_json::from_str( + body["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] + .as_str() + .unwrap(), + ) + .unwrap(); + assert!(sent["properties"]["name"].get("maxLength").is_none()); + assert_eq!(sent["properties"]["name"]["description"], "maxLength: 20"); + } + + #[tokio::test] + async fn converse_tool_path_appends_the_synthetic_tool_and_forces_it_on_nova() { + let body = capture_bedrock_body( + "amazon.nova-pro-v1:0", + "converse", + &structured_request(person_schema()), + false, + ) + .await; + let tools = body["toolConfig"]["tools"].as_array().unwrap(); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["toolSpec"]["name"], JSON_TOOL_NAME); + assert_eq!( + tools[0]["toolSpec"]["inputSchema"]["json"]["additionalProperties"], + false + ); + assert_eq!( + body["toolConfig"]["toolChoice"]["tool"]["name"], + JSON_TOOL_NAME + ); + assert!(body.get("outputConfig").is_none(), "body={body}"); + } + + #[tokio::test] + async fn converse_tool_path_leaves_the_choice_auto_where_the_family_has_none() { + // Converse honours an explicit `toolChoice` on Anthropic and + // Amazon Nova only; sending one anywhere else fails the whole + // request, so the synthetic tool is offered rather than forced. + let body = capture_bedrock_body( + "meta.llama3-3-70b-instruct-v1:0", + "converse", + &structured_request(person_schema()), + false, + ) + .await; + let tools = body["toolConfig"]["tools"].as_array().unwrap(); + assert_eq!(tools[0]["toolSpec"]["name"], JSON_TOOL_NAME); + assert!( + body["toolConfig"].get("toolChoice").is_none(), + "body={body}" + ); + } + + #[tokio::test] + async fn converse_tool_path_keeps_the_callers_tools_and_their_tool_choice() { + let mut req = structured_request(person_schema()); + req.extra.insert( + "tools".into(), + serde_json::json!([{ + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + }]), + ); + req.extra.insert( + "tool_choice".into(), + serde_json::json!({"type": "function", "function": {"name": "get_weather"}}), + ); + let body = capture_bedrock_body("amazon.nova-pro-v1:0", "converse", &req, false).await; + let tools = body["toolConfig"]["tools"].as_array().unwrap(); + assert_eq!(tools.len(), 2); + assert_eq!(tools[0]["toolSpec"]["name"], "get_weather"); + assert_eq!(tools[1]["toolSpec"]["name"], JSON_TOOL_NAME); + // A choice the caller made themselves outranks the forcing. + assert_eq!( + body["toolConfig"]["toolChoice"]["tool"]["name"], + "get_weather" + ); + } + + #[tokio::test] + async fn a_tool_choice_of_auto_still_outranks_forcing_the_synthetic_tool() { + // `auto` is the client saying the model decides. Overriding it + // would mean an agent loop's own tools could never be called + // while `response_format` is set. The synthetic tool is still + // offered, so the model can reach the JSON by itself. + let mut req = structured_request(person_schema()); + req.extra.insert("tool_choice".into(), "auto".into()); + let body = capture_bedrock_body("amazon.nova-pro-v1:0", "converse", &req, false).await; + assert_eq!( + body["toolConfig"]["tools"][0]["toolSpec"]["name"], + JSON_TOOL_NAME + ); + assert!( + body["toolConfig"].get("toolChoice").is_none(), + "body={body}" + ); + } + + #[tokio::test] + async fn converse_tool_choice_none_offers_only_the_synthetic_tool() { + // "none" means the caller wants no tool call this turn. The + // schema still has to come out of one on this route, but their + // own tools must not go back on the table unforced. + let mut req = structured_request(person_schema()); + req.extra.insert( + "tools".into(), + serde_json::json!([{ + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + }]), + ); + req.extra.insert("tool_choice".into(), "none".into()); + let body = capture_bedrock_body("amazon.nova-pro-v1:0", "converse", &req, false).await; + let tools = body["toolConfig"]["tools"].as_array().unwrap(); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["toolSpec"]["name"], JSON_TOOL_NAME); + assert!( + body["toolConfig"].get("toolChoice").is_none(), + "body={body}" + ); + } + + #[tokio::test] + async fn a_publisher_without_tool_use_gets_no_tool_config_from_response_format() { + // Titan Text has no Converse tool support; a `toolConfig` it + // never asked for would fail the whole request. + let body = capture_bedrock_body( + "amazon.titan-text-express-v1", + "converse", + &structured_request(person_schema()), + false, + ) + .await; + assert!(body.get("toolConfig").is_none(), "body={body}"); + assert!(body.get("outputConfig").is_none(), "body={body}"); + } + + #[tokio::test] + async fn converse_tool_path_leaves_the_choice_auto_under_extended_thinking() { + let mut req = structured_request(person_schema()); + req.extra.insert( + "thinking".into(), + serde_json::json!({"type": "enabled", "budget_tokens": 2048}), + ); + let body = capture_bedrock_body("amazon.nova-pro-v1:0", "converse", &req, false).await; + assert_eq!( + body["toolConfig"]["tools"][0]["toolSpec"]["name"], + JSON_TOOL_NAME + ); + assert!( + body["toolConfig"].get("toolChoice").is_none(), + "body={body}" + ); + } + + #[tokio::test] + async fn json_object_without_a_schema_emits_nothing_on_either_bedrock_path() { + let mut req = ChatFormat::new("my-model", vec![ChatMessage::user("hi")]); + req.extra.insert( + "response_format".into(), + serde_json::json!({"type": "json_object"}), + ); + let invoke = capture_bedrock_body( + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "invoke", + &req, + false, + ) + .await; + assert!(invoke.get("output_config").is_none(), "body={invoke}"); + assert!(invoke.get("tools").is_none(), "body={invoke}"); + assert!(invoke.get("response_format").is_none(), "body={invoke}"); + + let converse = capture_bedrock_body("amazon.nova-pro-v1:0", "converse", &req, false).await; + assert!(converse.get("outputConfig").is_none(), "body={converse}"); + assert!(converse.get("toolConfig").is_none(), "body={converse}"); + } + + #[tokio::test] + async fn a_small_stream_budget_does_not_cut_the_fake_stream_leg() { + // On a streaming dispatch `ctx.deadline` is the streaming budget, + // which bounds a chunk gap rather than a whole completion. The + // tool route's upstream leg is not streaming, so it runs under + // the end-to-end budget the context carries alongside it. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(250)) + .set_body_json(serde_json::json!({ + "id": "msg_json", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20240620-v1", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": JSON_TOOL_NAME, + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + })), + ) + .expect(1) + .mount(&server) + .await; + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("anthropic.claude-3-5-sonnet-20240620-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ) + // A chunk-gap budget the completion would blow through, beside + // the end-to-end budget it fits inside. + .with_deadline(Duration::from_millis(50)) + .with_non_streaming_deadline(Some(Duration::from_secs(30))); + + let mut req = structured_request(person_schema()); + req.stream = Some(true); + let stream = bridge + .chat_stream(&req, &ctx) + .await + .expect("the fake-stream leg must not be cut by the chunk-gap budget"); + let chunks: Vec = futures::StreamExt::collect::>(stream) + .await + .into_iter() + .map(Result::unwrap) + .collect(); + assert_eq!( + chunks[1].delta.content.as_deref(), + Some(r#"{"name":"Ada"}"#) + ); + } + + #[tokio::test] + async fn converse_synthetic_tool_reply_comes_back_as_json_content() { + // The reverse translation: the model's call to the synthetic + // tool is the answer, so the client sees plain JSON content and + // no tool call it never asked for. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/converse$")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "output": {"message": {"role": "assistant", "content": [ + {"toolUse": { + "toolUseId": "tooluse_json", + "name": JSON_TOOL_NAME, + "input": {"name": "Ada"} + }} + ]}}, + "stopReason": "tool_use", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1} + }))) + .expect(1) + .mount(&server) + .await; + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("amazon.nova-pro-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let chat = bridge + .chat(&structured_request(person_schema()), &ctx) + .await + .unwrap(); + assert_eq!(chat.message.content.as_deref(), Some(r#"{"name":"Ada"}"#)); + assert!(!chat.message.extra.contains_key("tool_calls")); + assert_eq!(chat.finish_reason, FinishReason::Stop); + } + + #[tokio::test] + async fn invoke_synthetic_tool_reply_comes_back_as_json_content() { + // The /invoke half of the reverse translation: Claude families + // that cannot constrain their own decoding answer by calling the + // synthetic tool, and that call is the answer. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_json", + "model": "claude-3-5-sonnet-20240620-v1", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": JSON_TOOL_NAME, + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + }))) + .expect(1) + .mount(&server) + .await; + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("anthropic.claude-3-5-sonnet-20240620-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let chat = bridge + .chat(&structured_request(person_schema()), &ctx) + .await + .unwrap(); + assert_eq!(chat.message.content.as_deref(), Some(r#"{"name":"Ada"}"#)); + assert!(!chat.message.extra.contains_key("tool_calls")); + assert_eq!(chat.finish_reason, FinishReason::Stop); + } + + #[tokio::test] + async fn streaming_the_tool_path_fake_streams_the_non_streaming_answer() { + // The JSON only exists once the tool call is complete, so the + // stream request never reaches Converse: it runs over the + // publisher's own non-streaming wire and is rendered as chunks. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_json", + "model": "claude-3-5-sonnet-20240620-v1", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": JSON_TOOL_NAME, + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 4}, + }))) + .expect(1) + .mount(&server) + .await; + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("anthropic.claude-3-5-sonnet-20240620-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let mut req = structured_request(person_schema()); + req.stream = Some(true); + let stream = bridge.chat_stream(&req, &ctx).await.unwrap(); + let chunks: Vec = futures::StreamExt::collect::>(stream) + .await + .into_iter() + .map(Result::unwrap) + .collect(); + assert_eq!(chunks.len(), 4); + assert_eq!(chunks[0].delta.role, Some(Role::Assistant)); + assert_eq!( + chunks[1].delta.content.as_deref(), + Some(r#"{"name":"Ada"}"#) + ); + assert_eq!(chunks[2].finish_reason, Some(FinishReason::Stop)); + assert!(chunks[3].usage.is_some()); + } + #[tokio::test] async fn chat_converse_translates_tool_use_response_to_tool_calls() { // The other half of #560: a Converse `toolUse` content block must diff --git a/crates/aisix-provider-openai/Cargo.toml b/crates/aisix-provider-openai/Cargo.toml index 9ebd68a8..2bd66d12 100644 --- a/crates/aisix-provider-openai/Cargo.toml +++ b/crates/aisix-provider-openai/Cargo.toml @@ -23,5 +23,10 @@ bytes.workspace = true http.workspace = true [dev-dependencies] +# Only for the strict-structured-output test: the `/v1/messages` inbound +# translation lives in the Anthropic crate, and what this edge must emit +# for it is exactly what that translation feeds in. No cycle — the +# Anthropic crate depends on aisix-gateway and aisix-core only. +aisix-provider-anthropic = { path = "../aisix-provider-anthropic" } tokio = { workspace = true, features = ["macros", "rt", "time"] } wiremock.workspace = true diff --git a/crates/aisix-provider-openai/src/bridge.rs b/crates/aisix-provider-openai/src/bridge.rs index fc74d82a..4aede960 100644 --- a/crates/aisix-provider-openai/src/bridge.rs +++ b/crates/aisix-provider-openai/src/bridge.rs @@ -20,6 +20,7 @@ //! - elapsed deadline → `BridgeError::Timeout { elapsed_ms }` use aisix_core::{RequestOverrides, ResponseOverrides, StreamDoneMarker}; +use aisix_gateway::structured_output::close_object_schemas; use aisix_gateway::url_cache::cached_endpoint_url; use aisix_gateway::{ apply_request_headers, Bridge, BridgeContext, BridgeError, ChatChunk, ChatChunkStream, @@ -313,6 +314,7 @@ fn prepare_outbound_body( ) -> Result { let mut body = serde_json::to_value(typed) .map_err(|e| BridgeError::Config(format!("serialize request body: {e}")))?; + close_strict_response_format_schema(&mut body); if let Some(r) = request { apply_param_renames(&mut body, &r.param_renames); if let Some(constraints) = &r.param_constraints { @@ -326,6 +328,36 @@ fn prepare_outbound_body( Ok(body) } +/// Apply OpenAI strict mode's schema rule at the edge that declares it: +/// when `response_format.json_schema.strict` is true, every object in +/// the schema must carry `additionalProperties: false` and list each of +/// its declared properties in `required`, or the API rejects it. +/// +/// Public because every OpenAI-wire edge has to apply it — Azure +/// OpenAI builds its own body from the same typed structs, and an +/// edge that skips this rejects a strict schema the caller never had +/// to write out in full. +/// +/// This lives here rather than wherever the `response_format` was +/// assembled because the promotion is only correct for *this* wire. The +/// same normalised request also reaches the Anthropic, Bedrock and +/// Gemini edges, and there `required` is an ordinary keyword whose +/// contents are the caller's own statement — promoting it would silently +/// make their optional fields mandatory. A caller who sent `strict` +/// themselves gets the same treatment they would have got from OpenAI's +/// own validation, so a schema that was already complete is unchanged. +pub fn close_strict_response_format_schema(body: &mut Value) { + let Some(json_schema) = body.pointer_mut("/response_format/json_schema") else { + return; + }; + if json_schema.get("strict").and_then(|s| s.as_bool()) != Some(true) { + return; + } + if let Some(schema) = json_schema.get_mut("schema") { + close_object_schemas(schema); + } +} + /// Build the base outbound `HeaderMap` (Authorization, Content-Type, /// x-aisix-request-id, and optionally Accept: text/event-stream /// for streaming calls), then merge any `default_headers` the PK carries. @@ -777,6 +809,104 @@ fn parse_stream_chunk( #[cfg(test)] mod tests { use super::*; + + // ── strict structured outputs at the OpenAI edge ────────────────── + + /// The exact `response_format` an Anthropic-shaped `/v1/messages` + /// request used to put on the OpenAI wire, when the strict closing + /// still happened during the inbound translation. The closing moved + /// here so the Anthropic, Bedrock and Gemini edges stop inheriting + /// it; this pins that the OpenAI body did not move with it. + #[test] + fn a_translated_messages_request_reaches_openai_byte_for_byte_as_before() { + use aisix_gateway::{ChatFormat, ChatMessage}; + use aisix_provider_anthropic::wire::translate_extras_to_openai_shape; + + // What the `/v1/messages` caller sent. + let mut extra = serde_json::json!({ + "output_config": { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": { + "type": "array", + "items": { + "type": "object", + "properties": {"high": {"type": "number"}}, + }, + }, + }, + }, + } + } + }) + .as_object() + .unwrap() + .clone(); + translate_extras_to_openai_shape(&mut extra, aisix_core::MappedEffort::AsWritten); + + let mut req = ChatFormat::new("m", vec![ChatMessage::user("weather?")]); + req.extra = extra; + let messages = messages_from(&req); + let typed = build_request(&req, "gpt-4o", &messages, false); + let body = prepare_outbound_body(&typed, None, None).unwrap(); + + assert_eq!( + body["response_format"], + serde_json::json!({ + "type": "json_schema", + "json_schema": { + "name": "structured_output", + "strict": true, + "schema": { + "type": "object", + "additionalProperties": false, + "required": ["city", "days"], + "properties": { + "city": {"type": "string"}, + "days": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["high"], + "properties": {"high": {"type": "number"}}, + }, + }, + }, + }, + }, + }) + ); + } + + #[test] + fn a_response_format_that_is_not_strict_reaches_openai_untouched() { + use aisix_gateway::{ChatFormat, ChatMessage}; + + let schema = serde_json::json!({ + "type": "object", + "properties": {"a": {"type": "string"}, "b": {"type": "string"}}, + "required": ["a"], + }); + let mut req = ChatFormat::new("m", vec![ChatMessage::user("hi")]); + req.extra.insert( + "response_format".into(), + serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema, "strict": false}, + }), + ); + let messages = messages_from(&req); + let typed = build_request(&req, "gpt-4o", &messages, false); + let body = prepare_outbound_body(&typed, None, None).unwrap(); + // Not strict: the caller's `required` is theirs, and OpenAI does + // not demand the closing. + assert_eq!(body["response_format"]["json_schema"]["schema"], schema); + } use aisix_core::{Model, ProviderKey}; use aisix_gateway::{ChatMessage, FinishReason, Role}; use std::sync::Arc; diff --git a/crates/aisix-provider-openai/src/lib.rs b/crates/aisix-provider-openai/src/lib.rs index a906e7bd..a494ad02 100644 --- a/crates/aisix-provider-openai/src/lib.rs +++ b/crates/aisix-provider-openai/src/lib.rs @@ -15,4 +15,4 @@ mod bridge; pub mod overrides; pub mod wire; -pub use bridge::{OpenAiBridge, OPENAI_DEFAULT_BASE}; +pub use bridge::{close_strict_response_format_schema, OpenAiBridge, OPENAI_DEFAULT_BASE}; diff --git a/crates/aisix-provider-vertex/src/bridge.rs b/crates/aisix-provider-vertex/src/bridge.rs index d33952e7..d906a46f 100644 --- a/crates/aisix-provider-vertex/src/bridge.rs +++ b/crates/aisix-provider-vertex/src/bridge.rs @@ -24,6 +24,10 @@ use aisix_gateway::{ sse::{SseDecoder, SseEvent}, + structured_output::{ + apply_schema_limits, json_schema_from_response_format, response_into_fake_stream_chunks, + unwrap_json_tool_call, GEMINI_OPENAPI_SCHEMA_LIMITS, + }, Bridge, BridgeContext, BridgeError, ChatChunk, ChatChunkStream, ChatDelta, ChatFormat, ChatMessage, ChatResponse, EmbeddingObject, EmbeddingRequest, EmbeddingResponse, EmbeddingUsage, EmbeddingVector, FinishReason, Role, UsageStats, @@ -52,7 +56,7 @@ use crate::wire; // `StreamState` (the same decoder the direct Anthropic bridge uses). use aisix_provider_anthropic::wire::{ build_request as build_anthropic_request, response_into_chat_response, split_system, - AnthropicResponse, AnthropicStreamEvent, StreamState, + structured_output_for, AnthropicResponse, AnthropicStreamEvent, StreamState, StructuredOutput, }; // Llama + the OpenAI-compatible MaaS family on Vertex use the OpenAI @@ -703,6 +707,27 @@ impl Bridge for VertexBridge { mistral-* / jamba-*" )) })?; + // The Claude tool route's JSON only exists once the synthetic + // tool call has been assembled, so it cannot be streamed as it + // arrives. Run the request non-streaming and fake-stream the + // translated result: the client sees an ordinary chunk sequence, + // and usage rides its own terminal chunk exactly as on a real + // stream. Gemini never takes that route — this bridge sends it + // no tools — and the OpenAI-shim publishers speak + // `response_format` natively. + if publisher == VertexPublisher::Anthropic + && matches!( + structured_output_for(req, upstream_id), + StructuredOutput::Tool(_) + ) + { + // The leg is not streaming, so it runs under the budget a + // non-streaming call would have got — the streaming budget + // this context carries bounds a chunk gap, not a completion. + let chunks = + response_into_fake_stream_chunks(self.chat(req, &ctx.non_streaming_ctx()).await?); + return Ok(Box::pin(futures::stream::iter(chunks.into_iter().map(Ok)))); + } match publisher { VertexPublisher::Google => self.chat_gemini_stream(req, ctx, upstream_id).await, VertexPublisher::OpenAiCompat => { @@ -879,7 +904,7 @@ impl VertexBridge { }, )?; - let typed = build_gemini_request(req); + let typed = build_gemini_request(req, upstream_id); // Audit LOW-4: Gemini requires `contents` to be a non-empty // array. If the caller passed system-only messages (lifted to // `systemInstruction`), `contents` ends up empty and Vertex @@ -991,6 +1016,10 @@ impl VertexBridge { // body shaping, differing only in the version string. let (system, messages) = split_system(req).map_err(|e| BridgeError::InvalidUpstreamConfig(format!("{e}")))?; + let synthetic_json_tool = matches!( + structured_output_for(req, upstream_id), + StructuredOutput::Tool(_) + ); let anthropic_req = build_anthropic_request(req, upstream_id, system, messages, false); let mut body_value = serde_json::to_value(&anthropic_req) .map_err(|e| BridgeError::Config(format!("serialize Anthropic request body: {e}")))?; @@ -1030,7 +1059,11 @@ impl VertexBridge { .json() .await .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; - Ok(response_into_chat_response(parsed)) + let mut chat = response_into_chat_response(parsed); + if synthetic_json_tool { + unwrap_json_tool_call(&mut chat); + } + Ok(chat) }) .await } @@ -1611,7 +1644,7 @@ impl VertexBridge { }, )?; - let typed = build_gemini_request(req); + let typed = build_gemini_request(req, upstream_id); if typed.contents.is_empty() { return Err(BridgeError::Config( "vertex chat: messages must include at least one user / \ @@ -1860,6 +1893,10 @@ fn apply_body_overrides(body: &mut serde_json::Value, ctx: &BridgeContext) { // ─── Gemini wire shapes ──────────────────────────────────────────────── +/// The only `responseMimeType` this bridge sets: Gemini's structured +/// output is gated on it, and both schema fields are inert without it. +const GEMINI_JSON_MIME_TYPE: &str = "application/json"; + /// Gemini's `generateContent` request body per /// . /// @@ -1899,6 +1936,158 @@ struct GeminiGenerationConfig { top_p: Option, #[serde(skip_serializing_if = "Option::is_none", rename = "maxOutputTokens")] max_output_tokens: Option, + /// Gemini's constrained-decoding switch. `"application/json"` is the + /// only value this bridge sets; without it neither schema field has + /// any effect. + #[serde(skip_serializing_if = "Option::is_none", rename = "responseMimeType")] + response_mime_type: Option<&'static str>, + /// Gemini 2 and later: a standard JSON Schema, sent as the caller + /// wrote it. + #[serde(skip_serializing_if = "Option::is_none", rename = "responseJsonSchema")] + response_json_schema: Option, + /// Gemini 1.x: the older OpenAPI-flavoured schema dialect. See + /// [`gemini_openapi_schema`]. + #[serde(skip_serializing_if = "Option::is_none", rename = "responseSchema")] + response_schema: Option, +} + +/// What a caller's OpenAI `response_format` becomes in Gemini's +/// `generationConfig`. +/// +/// Gemini has no tool-call fallback here — this bridge sends no `tools` +/// at all, so there is nothing to force — and no prompt injection: a +/// model that cannot constrain its decoding is left to answer as it +/// would have. +#[derive(Debug, PartialEq, Eq)] +enum GeminiJsonOutput { + /// No `response_format`, or `{"type":"text"}`: nothing to emit. + None, + /// `{"type":"json_object"}` — JSON, but no schema to constrain it to. + MimeOnly, + /// Gemini 2+: `responseJsonSchema`, standard JSON Schema. + JsonSchema(serde_json::Value), + /// Gemini 1.x: `responseSchema`, the OpenAPI-flavoured dialect. + OpenApiSchema(serde_json::Value), +} + +/// Decide what `response_format` becomes for `upstream_model`. +/// +/// Both schema fields carry the same document; which one Gemini reads it +/// out of is a generation thing. `responseJsonSchema` takes ordinary +/// JSON Schema and only exists from Gemini 2 onwards; `responseSchema` +/// is the older OpenAPI-derived subset every generation accepts. An +/// unrecognisable model name therefore falls back to `responseSchema`, +/// the one that works everywhere. +fn gemini_json_output(req: &ChatFormat, upstream_model: &str) -> GeminiJsonOutput { + let Some(response_format) = req.extra.get("response_format") else { + return GeminiJsonOutput::None; + }; + let Some(schema) = json_schema_from_response_format(response_format) else { + return match response_format.get("type").and_then(|t| t.as_str()) { + Some("json_object") => GeminiJsonOutput::MimeOnly, + _ => GeminiJsonOutput::None, + }; + }; + if gemini_major_version(upstream_model).is_some_and(|major| major >= 2) { + GeminiJsonOutput::JsonSchema(schema) + } else { + GeminiJsonOutput::OpenApiSchema(gemini_openapi_schema(&schema)) + } +} + +/// Read the generation off a Gemini model name: `2` from +/// `gemini-2.5-flash`, `3` from `gemini-3-pro-preview`. The gateway +/// holds no capability map, so the name is all there is. +/// +/// `None` for anything that is not a `gemini-` name — including +/// the undated experimental aliases (`gemini-exp-1206`) — which lands on +/// the older schema field, the one every generation accepts. +fn gemini_major_version(model: &str) -> Option { + let lowered = model.trim().to_ascii_lowercase(); + // Vertex accepts both the bare id and the `models/` spelling. + let name = lowered.rsplit('/').next().unwrap_or(&lowered); + let rest = name.strip_prefix("gemini")?.trim_start_matches(['-', '_']); + let digits: String = rest.chars().take_while(char::is_ascii_digit).collect(); + digits.parse().ok() +} + +/// Rewrite a JSON Schema into the OpenAPI-derived dialect Gemini 1.x +/// reads out of `responseSchema`: +/// +/// * `type` is an upper-case OpenAPI type name (`OBJECT`, `STRING`), +/// * members outside the `Schema` type — `additionalProperties` +/// among them — do not exist and are rejected by name, which +/// [`apply_schema_limits`] has already dealt with, +/// * `propertyOrdering` fixes the order the model emits an object's +/// members in — omitted, the order is unspecified. +/// +/// Run after [`apply_schema_limits`], which has already removed the +/// keywords the dialect has no member for and inlined its `$ref`s. +/// +/// The ordering emitted is the order the properties appear in the schema +/// as this gateway serialises it, so the request is self-consistent. +fn gemini_openapi_schema(schema: &serde_json::Value) -> serde_json::Value { + let mut out = schema.clone(); + // First narrow the schema to the dialect's vocabulary — this is + // also what inlines `$ref`, which the dialect has no spelling for — + // then rewrite what survives into the dialect's own shape. + apply_schema_limits(&mut out, &GEMINI_OPENAPI_SCHEMA_LIMITS); + rewrite_gemini_openapi_schema(&mut out); + out +} + +fn rewrite_gemini_openapi_schema(schema: &mut serde_json::Value) { + let Some(obj) = schema.as_object_mut() else { + return; + }; + match obj.get_mut("type") { + Some(serde_json::Value::String(ty)) => *ty = ty.to_ascii_uppercase(), + // A union type (`["string","null"]`, how strict mode spells an + // optional field) would otherwise reach the wire with lower-case + // names Vertex does not recognise. + Some(serde_json::Value::Array(types)) => { + for ty in types.iter_mut() { + if let serde_json::Value::String(ty) = ty { + *ty = ty.to_ascii_uppercase(); + } + } + } + _ => {} + } + if let Some(properties) = obj.get("properties").and_then(|p| p.as_object()) { + let ordering: Vec = + properties.keys().map(|k| k.as_str().into()).collect(); + if !ordering.is_empty() { + obj.insert("propertyOrdering".to_string(), ordering.into()); + } + } + if let Some(properties) = obj.get_mut("properties").and_then(|p| p.as_object_mut()) { + for property in properties.values_mut() { + rewrite_gemini_openapi_schema(property); + } + } + // `items` also has the draft-07 tuple form. Gemini's `Schema.items` + // is a single schema, so a tuple array is a shape this dialect + // cannot express at all and Vertex rejects it — the same standing + // as an unresolvable `$ref`. The elements are still rewritten, so + // the two walkers agree about where schemas live and the request + // that goes up is the caller's own, not a half-converted one. + match obj.get_mut("items") { + Some(serde_json::Value::Array(items)) => { + for item in items { + rewrite_gemini_openapi_schema(item); + } + } + Some(items) => rewrite_gemini_openapi_schema(items), + None => {} + } + for key in ["anyOf", "oneOf", "allOf"] { + if let Some(branches) = obj.get_mut(key).and_then(|b| b.as_array_mut()) { + for branch in branches { + rewrite_gemini_openapi_schema(branch); + } + } + } } /// Translate the gateway's [`ChatFormat`] into Gemini's @@ -1913,7 +2102,12 @@ struct GeminiGenerationConfig { /// - Tool messages: out of scope for D5.2.a; treated as user text /// (preserves conversation history without 400ing the upstream) /// - `temperature`, `top_p`, `max_tokens` → `generationConfig.*` -fn build_gemini_request(req: &ChatFormat) -> GeminiGenerateContentRequest { +/// - `response_format` → `generationConfig.responseMimeType` plus the +/// schema field this model's generation reads — see +/// [`gemini_json_output`]. Every other `extra` field is dropped: this +/// body carries only the fields named on the struct, so an OpenAI-only +/// knob can never ride onto the wire and 400. +fn build_gemini_request(req: &ChatFormat, upstream_model: &str) -> GeminiGenerateContentRequest { let mut system_parts: Vec = Vec::new(); let mut contents: Vec = Vec::new(); for m in &req.messages { @@ -1943,16 +2137,33 @@ fn build_gemini_request(req: &ChatFormat) -> GeminiGenerateContentRequest { }], }) }; - let generation_config = - if req.temperature.is_some() || req.top_p.is_some() || req.max_tokens.is_some() { - Some(GeminiGenerationConfig { - temperature: req.temperature, - top_p: req.top_p, - max_output_tokens: req.max_tokens, - }) - } else { - None + let json_output = gemini_json_output(req, upstream_model); + let generation_config = if req.temperature.is_some() + || req.top_p.is_some() + || req.max_tokens.is_some() + || json_output != GeminiJsonOutput::None + { + let (response_mime_type, response_json_schema, response_schema) = match json_output { + GeminiJsonOutput::None => (None, None, None), + GeminiJsonOutput::MimeOnly => (Some(GEMINI_JSON_MIME_TYPE), None, None), + GeminiJsonOutput::JsonSchema(schema) => { + (Some(GEMINI_JSON_MIME_TYPE), Some(schema), None) + } + GeminiJsonOutput::OpenApiSchema(schema) => { + (Some(GEMINI_JSON_MIME_TYPE), None, Some(schema)) + } }; + Some(GeminiGenerationConfig { + temperature: req.temperature, + top_p: req.top_p, + max_output_tokens: req.max_tokens, + response_mime_type, + response_json_schema, + response_schema, + }) + } else { + None + }; GeminiGenerateContentRequest { contents, system_instruction, @@ -2684,7 +2895,7 @@ mod tests { #[test] fn build_gemini_request_translates_user_turn() { let req = ChatFormat::new("my-gemini", vec![ChatMessage::user("hi")]); - let body = build_gemini_request(&req); + let body = build_gemini_request(&req, "gemini-2.0-flash"); assert_eq!(body.contents.len(), 1); assert_eq!(body.contents[0].role, "user"); assert_eq!(body.contents[0].parts[0].text, "hi"); @@ -2701,7 +2912,7 @@ mod tests { ChatMessage::assistant("hello back"), ], ); - let body = build_gemini_request(&req); + let body = build_gemini_request(&req, "gemini-2.0-flash"); assert_eq!(body.contents.len(), 2); assert_eq!(body.contents[0].role, "user"); // Gemini uses `model`, NOT `assistant`. @@ -2717,7 +2928,7 @@ mod tests { ChatMessage::user("hi"), ], ); - let body = build_gemini_request(&req); + let body = build_gemini_request(&req, "gemini-2.0-flash"); // System NOT in contents[]. assert_eq!(body.contents.len(), 1); assert_eq!(body.contents[0].role, "user"); @@ -2736,7 +2947,7 @@ mod tests { ChatMessage::user("hi"), ], ); - let body = build_gemini_request(&req); + let body = build_gemini_request(&req, "gemini-2.0-flash"); let sys = body.system_instruction.as_ref().unwrap(); assert_eq!(sys.parts[0].text, "rule 1\n\nrule 2"); } @@ -2747,13 +2958,229 @@ mod tests { req.temperature = Some(0.7); req.top_p = Some(0.9); req.max_tokens = Some(100); - let body = build_gemini_request(&req); + let body = build_gemini_request(&req, "gemini-2.0-flash"); let gc = body.generation_config.as_ref().unwrap(); assert_eq!(gc.temperature, Some(0.7)); assert_eq!(gc.top_p, Some(0.9)); assert_eq!(gc.max_output_tokens, Some(100)); } + // ─── Gemini structured outputs ───────────────────────────────────── + + fn gemini_request_with_response_format(response_format: serde_json::Value) -> ChatFormat { + let mut req = ChatFormat::new("my-gemini", vec![ChatMessage::user("who are you")]); + req.extra.insert("response_format".into(), response_format); + req + } + + fn json_schema_format(schema: serde_json::Value) -> serde_json::Value { + serde_json::json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema, "strict": true}, + }) + } + + fn person_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "pet": {"type": "object", "properties": {"kind": {"type": "string"}}}, + }, + "required": ["name"], + "additionalProperties": false, + }) + } + + #[test] + fn gemini_generation_gate_reads_the_major_off_the_model_name() { + for (name, major) in [ + ("gemini-2.0-flash", Some(2)), + ("gemini-2.5-flash", Some(2)), + ("gemini-2.5-pro-preview-05-06", Some(2)), + ("gemini-3-pro-preview", Some(3)), + ("gemini-3", Some(3)), + ("gemini-10-ultra", Some(10)), + ("GEMINI-2.5-FLASH", Some(2)), + ("models/gemini-2.5-flash", Some(2)), + ("gemini-1.5-pro", Some(1)), + ("gemini-1.0-pro-002", Some(1)), + // Not a `gemini-` name: no generation to read. + ("gemini-exp-1206", None), + ("gemini", None), + ("text-embedding-005", None), + ("claude-sonnet-4-5", None), + ] { + assert_eq!(gemini_major_version(name), major, "{name}"); + } + } + + #[test] + fn gemini_2_and_later_carry_the_schema_in_response_json_schema() { + let req = gemini_request_with_response_format(json_schema_format(person_schema())); + let body = build_gemini_request(&req, "gemini-2.5-flash"); + let gc = body.generation_config.as_ref().unwrap(); + assert_eq!(gc.response_mime_type, Some("application/json")); + // Standard JSON Schema, forwarded as the caller wrote it. + assert_eq!(gc.response_json_schema.as_ref(), Some(&person_schema())); + assert!(gc.response_schema.is_none()); + } + + #[test] + fn gemini_1_x_carries_the_schema_in_the_openapi_dialect() { + let req = gemini_request_with_response_format(json_schema_format(person_schema())); + let body = build_gemini_request(&req, "gemini-1.5-pro"); + let gc = body.generation_config.as_ref().unwrap(); + assert_eq!(gc.response_mime_type, Some("application/json")); + assert!(gc.response_json_schema.is_none()); + let schema = gc.response_schema.as_ref().unwrap(); + // Upper-case OpenAPI type names, no `additionalProperties` + // (Gemini rejects it), and an explicit member ordering. + assert_eq!(schema["type"], "OBJECT"); + assert_eq!(schema["properties"]["name"]["type"], "STRING"); + assert_eq!(schema["properties"]["pet"]["type"], "OBJECT"); + assert_eq!( + schema["properties"]["pet"]["properties"]["kind"]["type"], + "STRING" + ); + assert!(schema.get("additionalProperties").is_none()); + assert_eq!( + schema["propertyOrdering"], + serde_json::json!(["name", "pet"]) + ); + assert_eq!( + schema["properties"]["pet"]["propertyOrdering"], + serde_json::json!(["kind"]) + ); + // `required` is the caller's own statement on both dialects. + assert_eq!(schema["required"], serde_json::json!(["name"])); + } + + #[test] + fn gemini_openapi_dialect_reaches_arrays_and_branches() { + let req = gemini_request_with_response_format(json_schema_format(serde_json::json!({ + "type": "array", + "items": { + "type": "object", + "properties": {"id": {"type": "integer"}}, + "additionalProperties": false, + }, + "anyOf": [{"type": "string"}], + }))); + let body = build_gemini_request(&req, "gemini-1.5-flash"); + let schema = body.generation_config.unwrap().response_schema.unwrap(); + assert_eq!(schema["type"], "ARRAY"); + assert_eq!(schema["items"]["type"], "OBJECT"); + assert_eq!(schema["items"]["properties"]["id"]["type"], "INTEGER"); + assert!(schema["items"].get("additionalProperties").is_none()); + assert_eq!(schema["anyOf"][0]["type"], "STRING"); + } + + #[test] + fn gemini_property_ordering_matches_the_schema_as_it_goes_on_the_wire() { + // `propertyOrdering` is only useful if it names the properties + // in the order the request itself presents them; a list that + // disagrees with the accompanying `properties` object would fix + // an order the schema does not show. + let req = gemini_request_with_response_format(json_schema_format(serde_json::json!({ + "type": "object", + "properties": { + "zeta": {"type": "string"}, + "alpha": {"type": "string"}, + "mid": {"type": "string"}, + }, + }))); + let body = serde_json::to_value(build_gemini_request(&req, "gemini-1.5-pro")).unwrap(); + let schema = &body["generationConfig"]["responseSchema"]; + let on_the_wire: Vec<&str> = schema["properties"] + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect(); + let ordering: Vec<&str> = schema["propertyOrdering"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap()) + .collect(); + assert_eq!(ordering, on_the_wire); + } + + #[test] + fn gemini_openapi_dialect_inlines_definitions_and_uppercases_union_types() { + // The dialect is an OpenAPI `Schema` object: it rejects members + // by name, so `$defs` cannot ride along and `$ref` has no + // spelling. The definitions are inlined and the blocks removed + // before the dialect rewrite, so what `$ref` pointed at gets the + // same upper-casing as everything else. + let req = gemini_request_with_response_format(json_schema_format(serde_json::json!({ + "type": "object", + "properties": { + "pet": {"$ref": "#/$defs/Pet"}, + "nickname": {"type": ["string", "null"]}, + }, + "$defs": { + "Pet": { + "type": "object", + "properties": {"kind": {"type": "string"}}, + "additionalProperties": false, + }, + }, + }))); + let schema = build_gemini_request(&req, "gemini-1.5-pro") + .generation_config + .unwrap() + .response_schema + .unwrap(); + assert!(schema.get("$defs").is_none(), "{schema}"); + let pet = &schema["properties"]["pet"]; + assert!(pet.get("$ref").is_none(), "{schema}"); + assert_eq!(pet["type"], "OBJECT"); + assert_eq!(pet["properties"]["kind"]["type"], "STRING"); + assert!(pet.get("additionalProperties").is_none()); + assert_eq!( + schema["properties"]["nickname"]["type"], + serde_json::json!(["STRING", "NULL"]) + ); + } + + #[test] + fn gemini_json_object_asks_for_json_without_a_schema() { + for model in ["gemini-2.5-flash", "gemini-1.5-pro"] { + let req = + gemini_request_with_response_format(serde_json::json!({"type": "json_object"})); + let gc = build_gemini_request(&req, model).generation_config.unwrap(); + assert_eq!(gc.response_mime_type, Some("application/json")); + assert!(gc.response_json_schema.is_none()); + assert!(gc.response_schema.is_none()); + } + } + + #[test] + fn gemini_response_format_text_emits_no_generation_config_at_all() { + let req = gemini_request_with_response_format(serde_json::json!({"type": "text"})); + assert!(build_gemini_request(&req, "gemini-2.5-flash") + .generation_config + .is_none()); + } + + #[test] + fn gemini_response_format_never_reaches_the_wire_verbatim() { + // It has no top-level Gemini counterpart; forwarding it 400s. + let req = gemini_request_with_response_format(json_schema_format(person_schema())); + let body = serde_json::to_value(build_gemini_request(&req, "gemini-2.5-flash")).unwrap(); + assert!(body.get("response_format").is_none()); + assert_eq!( + body["generationConfig"]["responseMimeType"], + "application/json" + ); + assert_eq!( + body["generationConfig"]["responseJsonSchema"], + person_schema() + ); + } + // ─── Gemini response translation ─────────────────────────────────── #[test] @@ -3128,7 +3555,7 @@ mod tests { // ─── Dispatch end-to-end against wiremock via api_base override ── - use wiremock::matchers::{header, method, path}; + use wiremock::matchers::{header, method, path, path_regex}; use wiremock::{Mock, MockServer, Request as MockRequest, Respond, ResponseTemplate}; #[derive(Clone, Default)] @@ -3276,6 +3703,132 @@ mod tests { ); } + #[tokio::test] + async fn vertex_claude_carries_the_schema_in_output_config_on_a_4_5_family() { + // Claude on Vertex speaks the Anthropic Messages wire, so a + // caller's `response_format` becomes whatever the shared + // serializer makes of it — the native control here, since the + // Vertex model id IS the Claude name (`@date`, not `-date`). + let server = MockServer::start().await; + let responder = CapturingAnthropicResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/v1/projects/.+:rawPredict$")) + .respond_with(responder.clone()) + .expect(1) + .mount(&server) + .await; + let bridge = VertexBridge::new().with_api_base_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("claude-sonnet-4-5@20250929"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = gemini_request_with_response_format(json_schema_format(person_schema())); + let _ = bridge.chat(&req, &ctx).await.unwrap(); + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + assert_eq!(body["output_config"]["format"]["type"], "json_schema"); + assert_eq!( + body["output_config"]["format"]["schema"]["additionalProperties"], + false + ); + assert!(body.get("response_format").is_none(), "body={body}"); + assert!(body.get("tools").is_none(), "body={body}"); + } + + #[tokio::test] + async fn a_small_stream_budget_does_not_cut_the_fake_stream_leg() { + // On a streaming dispatch the deadline is the streaming budget, + // which bounds a chunk gap rather than a whole completion. The + // Claude tool route's upstream leg is not streaming, so it runs + // under the end-to-end budget carried beside it. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/v1/projects/.+:rawPredict$")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(std::time::Duration::from_millis(250)) + .set_body_json(serde_json::json!({ + "id": "msg_json", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": "json_tool_call", + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 3, "output_tokens": 5}, + })), + ) + .mount(&server) + .await; + let bridge = VertexBridge::new().with_api_base_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("claude-3-5-sonnet-v2@20241022"), + sample_pk_with_secret(valid_secret_json()), + ) + .with_deadline(std::time::Duration::from_millis(50)) + .with_non_streaming_deadline(Some(std::time::Duration::from_secs(30))); + + let mut req = gemini_request_with_response_format(json_schema_format(person_schema())); + req.stream = Some(true); + let stream = bridge + .chat_stream(&req, &ctx) + .await + .expect("the fake-stream leg must not be cut by the chunk-gap budget"); + let chunks: Vec = futures::StreamExt::collect::>(stream) + .await + .into_iter() + .map(Result::unwrap) + .collect(); + assert_eq!( + chunks[1].delta.content.as_deref(), + Some(r#"{"name":"Ada"}"#) + ); + } + + #[tokio::test] + async fn vertex_claude_synthetic_tool_reply_comes_back_as_json_content() { + // An older Claude family takes the tool route, and the call it + // makes is the answer — the client must not be handed a tool + // call it never offered. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/v1/projects/.+:rawPredict$")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_json", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [{ + "type": "tool_use", + "id": "toolu_json", + "name": "json_tool_call", + "input": {"name": "Ada"}, + }], + "stop_reason": "tool_use", + "usage": {"input_tokens": 3, "output_tokens": 5}, + }))) + .expect(1) + .mount(&server) + .await; + let bridge = VertexBridge::new().with_api_base_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("claude-3-5-sonnet-v2@20241022"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = gemini_request_with_response_format(json_schema_format(person_schema())); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content.as_deref(), Some(r#"{"name":"Ada"}"#)); + assert!(!chat.message.extra.contains_key("tool_calls")); + assert_eq!(chat.finish_reason, FinishReason::Stop); + } + /// Capturing responder that returns a native Anthropic SSE stream /// (the `:streamRawPredict` wire) while recording the inbound request /// body — the streaming analogue of [`CapturingAnthropicResponder`]. diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 2ec7d909..57fa1457 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1679,6 +1679,11 @@ async fn dispatch( if let Some(d) = stream_budget { ctx = ctx.with_deadline(d); } + // A bridge that has to answer this streaming request with a + // non-streaming upstream leg (the structured-output tool + // route) measures that leg against the end-to-end budget, + // not against the per-chunk one set above. + ctx = ctx.with_non_streaming_deadline(timeouts.request); // How many times to re-hit the SAME target (with backoff) on a // retryable failure before failing over to the next one. diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index af162a2e..f8b57f0a 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -2096,6 +2096,12 @@ async fn cross_provider_dispatch( if let Some(d) = connect_deadline { ctx = ctx.with_deadline(d); } + // See chat.rs: the structured-output tool route answers a streaming + // request with a non-streaming upstream leg, which is entitled to + // the end-to-end budget rather than the per-chunk one. + if is_stream { + ctx = ctx.with_non_streaming_deadline(timeouts.request); + } let provider_label = provider.to_ascii_lowercase(); let provider_key_id = model.provider_key_id.as_deref().unwrap_or("unknown"); let upstream_model = model.upstream_model().unwrap_or("unknown").to_string(); diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index e6e14fd4..5dc6c372 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -2224,6 +2224,12 @@ async fn responses_cross_provider_to_target( if let Some(d) = connect_deadline { ctx = ctx.with_deadline(d); } + // See chat.rs: the structured-output tool route answers a streaming + // request with a non-streaming upstream leg, which is entitled to + // the end-to-end budget rather than the per-chunk one. + if is_stream { + ctx = ctx.with_non_streaming_deadline(timeouts.request); + } let provider_label = provider.to_ascii_lowercase(); // least_busy: count this target as in-flight for the upstream call diff --git a/tests/e2e/src/cases/anthropic-structured-output-e2e.test.ts b/tests/e2e/src/cases/anthropic-structured-output-e2e.test.ts new file mode 100644 index 00000000..763b5380 --- /dev/null +++ b/tests/e2e/src/cases/anthropic-structured-output-e2e.test.ts @@ -0,0 +1,286 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + ProxyClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// Structured outputs on Anthropic-protocol upstreams. A chat caller (or a +// `/v1/responses` caller, whose `text.format` the bridge turns into +// `response_format`) that asks for JSON used to get prose: the field was +// consumed and dropped. It now becomes one of two request shapes, picked +// by the target model's family: +// +// * Claude 4.5 and later → `output_config.format`, the model's own +// structured-output control. +// * everything else (older Claude, and non-Claude models behind +// Anthropic-compatible endpoints) → a synthetic `json_tool_call` tool +// under a forced `tool_choice`, whose call the decoder translates back +// into plain JSON content. + +const CALLER_PLAINTEXT = "sk-anthropic-structured-output"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const PERSON_SCHEMA = { + type: "object", + properties: { + name: { type: "string" }, + age: { type: "integer" }, + }, +}; + +const RESPONSE_FORMAT = { + type: "json_schema", + json_schema: { name: "person", schema: PERSON_SCHEMA, strict: true }, +}; + +// What a Claude 4.5+ model returns once `output_config.format` constrains +// it: ordinary text that happens to be the JSON document. +const NATIVE_REPLY = { + id: "msg_native_json", + type: "message", + role: "assistant", + model: "claude-sonnet-4-5-20250929", + content: [{ type: "text", text: '{"name":"Ada","age":36}' }], + stop_reason: "end_turn", + usage: { input_tokens: 14, output_tokens: 9 }, +}; + +// What the tool path gets back: the model calls the synthetic tool, and +// the call's input is the answer. +const TOOL_REPLY = { + id: "msg_tool_json", + type: "message", + role: "assistant", + model: "claude-3-5-haiku-20241022", + content: [ + { + type: "tool_use", + id: "toolu_json", + name: "json_tool_call", + input: { name: "Ada", age: 36 }, + }, + ], + stop_reason: "tool_use", + usage: { input_tokens: 21, output_tokens: 12 }, +}; + +function lastBody(upstream: OpenAiUpstream): Record { + const last = upstream.receivedRequests.at(-1); + expect(last?.path).toBe("/v1/messages"); + return JSON.parse(last!.body); +} + +function chat(app: SpawnedApp, body: unknown): Promise { + return fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify(body), + }); +} + +describe("chat response_format → Anthropic structured outputs", () => { + let app: SpawnedApp | undefined; + let nativeUpstream: OpenAiUpstream | undefined; + let toolUpstream: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + nativeUpstream = await startOpenAiUpstream({ nonStreamBody: NATIVE_REPLY }); + toolUpstream = await startOpenAiUpstream({ nonStreamBody: TOOL_REPLY }); + app = await spawnApp(); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const nativePk = await seed.createProviderKey({ + display_name: "structured-native-pk", + provider: "anthropic", + adapter: "anthropic", + secret: "sk-ant-mock", + api_base: nativeUpstream.baseUrl, + }); + // The alias is deliberately not the upstream name: the family gate + // reads the model the gateway dispatches to, not what the caller + // typed. + await seed.createModel({ + display_name: "json-native", + provider: "anthropic", + model_name: "claude-sonnet-4-5", + provider_key_id: nativePk.id, + }); + const toolPk = await seed.createProviderKey({ + display_name: "structured-tool-pk", + provider: "anthropic", + adapter: "anthropic", + secret: "sk-ant-mock", + api_base: toolUpstream.baseUrl, + }); + await seed.createModel({ + display_name: "json-legacy", + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: toolPk.id, + }); + // Seeded last, so this key authenticating implies the whole seed + // set has reached the gateway's snapshot. + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["json-native", "json-legacy"], + }); + const proxy = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation( + async () => (await proxy.listModels()).status === 200, + ); + }); + + afterAll(async () => { + await app?.exit(); + await nativeUpstream?.close(); + await toolUpstream?.close(); + }); + + test("claude 4.5+ takes the native path: output_config.format, no response_format", async (ctx) => { + if (!etcdReachable || !app || !nativeUpstream) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-native", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.choices[0].message.content).toBe('{"name":"Ada","age":36}'); + + const sent = lastBody(nativeUpstream); + const format = (sent.output_config as Record)?.format; + expect(format?.type).toBe("json_schema"); + expect(format?.schema?.properties?.name?.type).toBe("string"); + // Anthropic rejects an open object schema. + expect(format?.schema?.additionalProperties).toBe(false); + // The OpenAI spelling never reaches the Anthropic body, and the + // native path adds no tool. + expect(sent.response_format).toBeUndefined(); + expect(sent.tools).toBeUndefined(); + expect(sent.tool_choice).toBeUndefined(); + }); + + test("older claude takes the tool path and the call comes back as JSON content", async (ctx) => { + if (!etcdReachable || !app || !toolUpstream) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-legacy", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + + const sent = lastBody(toolUpstream); + expect(sent.response_format).toBeUndefined(); + expect(sent.output_config).toBeUndefined(); + const tools = sent.tools as Array>; + expect(tools).toHaveLength(1); + expect(tools[0].name).toBe("json_tool_call"); + expect(tools[0].input_schema.additionalProperties).toBe(false); + expect(sent.tool_choice).toEqual({ type: "tool", name: "json_tool_call" }); + + // The caller never offered a tool, so it must not be told the model + // stopped to call one. + const body = await res.json(); + const message = body.choices[0].message; + expect(JSON.parse(message.content)).toEqual({ name: "Ada", age: 36 }); + expect(message.tool_calls).toBeUndefined(); + expect(body.choices[0].finish_reason).toBe("stop"); + }); + + test("streaming on the tool path fake-streams the JSON and still reports usage", async (ctx) => { + if (!etcdReachable || !app || !toolUpstream) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-legacy", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + stream: true, + stream_options: { include_usage: true }, + }); + expect(res.status).toBe(200); + expect(res.headers.get("content-type")).toContain("text/event-stream"); + const raw = await res.text(); + + // The upstream leg ran non-streaming — the JSON only exists once the + // tool call is complete. + const sent = lastBody(toolUpstream); + expect(sent.stream).toBe(false); + expect(sent.tool_choice).toEqual({ type: "tool", name: "json_tool_call" }); + + const frames = raw + .split("\n") + .filter((l) => l.startsWith("data: ") && !l.includes("[DONE]")) + .map((l) => JSON.parse(l.slice(6))); + expect(frames[0].choices[0].delta.role).toBe("assistant"); + const text = frames + .map((f) => f.choices?.[0]?.delta?.content ?? "") + .join(""); + expect(JSON.parse(text)).toEqual({ name: "Ada", age: 36 }); + expect(frames.some((f) => f.choices?.[0]?.finish_reason === "stop")).toBe( + true, + ); + expect( + frames.some((f) => f.usage?.completion_tokens === 12), + ).toBe(true); + }); + + test("/v1/responses text.format reaches Anthropic as output_config.format", async (ctx) => { + if (!etcdReachable || !app || !nativeUpstream) { + ctx.skip(); + return; + } + const res = await fetch(`${app.proxyUrl}/v1/responses`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "json-native", + input: "who is Ada", + text: { + format: { + type: "json_schema", + name: "person", + schema: PERSON_SCHEMA, + strict: true, + }, + }, + }), + }); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.output[0].content[0].text).toBe('{"name":"Ada","age":36}'); + + const sent = lastBody(nativeUpstream); + const format = (sent.output_config as Record)?.format; + expect(format?.type).toBe("json_schema"); + expect(format?.schema?.additionalProperties).toBe(false); + expect(sent.response_format).toBeUndefined(); + }); +}); diff --git a/tests/e2e/src/cases/provider-structured-output-e2e.test.ts b/tests/e2e/src/cases/provider-structured-output-e2e.test.ts new file mode 100644 index 00000000..2a9ac4ce --- /dev/null +++ b/tests/e2e/src/cases/provider-structured-output-e2e.test.ts @@ -0,0 +1,460 @@ +import { createHash } from "node:crypto"; +import { createServer, type Server } from "node:http"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + waitConfigPropagation, + type SpawnedApp, +} from "../harness/index.js"; + +// Structured outputs on the two providers whose chat bridge builds its +// own request shape rather than forwarding an OpenAI body. Both used to +// drop a caller's `response_format` on the floor and answer in prose. +// +// * Gemini takes `generationConfig.responseMimeType` plus the schema, +// in `responseJsonSchema` from Gemini 2 onwards and in the older +// OpenAPI-flavoured `responseSchema` before that. +// * Bedrock picks by whether the model constrains its own decoding. +// A Claude 4.5 answers non-streaming over the Anthropic Messages +// `/invoke` wire, whose control is `output_config.format`. Every +// other model takes the synthetic `json_tool_call` tool on Converse, +// and the call it makes is translated back into JSON content. + +const CALLER_PLAINTEXT = "sk-provider-structured-output"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const PERSON_SCHEMA = { + type: "object", + properties: { + name: { type: "string" }, + nickname: { type: "string" }, + }, + required: ["name"], +}; + +const RESPONSE_FORMAT = { + type: "json_schema", + json_schema: { name: "person", schema: PERSON_SCHEMA, strict: true }, +}; + +const ANSWER = '{"name":"Ada","nickname":"Countess"}'; + +// The streaming budget bounds the gap between chunks; the request +// budget bounds the whole call. The tool route answers a streaming +// request with ONE non-streaming upstream call, so it has to be +// measured against the second — these three values are what tells the +// two apart. +const STREAM_BUDGET_MS = 400; +const REQUEST_BUDGET_MS = 30_000; +const SLOW_UPSTREAM_MS = 1_200; + +interface RecordedRequest { + path: string; + body: string; +} + +interface RecordingUpstream { + baseUrl: string; + received: RecordedRequest[]; + close(): Promise; +} + +/** + * A JSON upstream that answers every route from one reply function, + * optionally after a delay — which is how a completion slower than one + * streaming chunk-gap budget is reproduced. + */ +async function startJsonUpstream( + reply: (path: string) => unknown, + delayMs = 0, +): Promise { + const received: RecordedRequest[] = []; + const server: Server = createServer((req, res) => { + res.on("error", () => {}); + const chunks: Buffer[] = []; + req.on("data", (c: Buffer) => chunks.push(c)); + req.on("end", () => { + const path = (req.url ?? "/").split("?")[0]; + received.push({ path, body: Buffer.concat(chunks).toString("utf8") }); + const send = () => { + res.statusCode = 200; + res.setHeader("content-type", "application/json"); + res.end(JSON.stringify(reply(path))); + }; + if (delayMs > 0) setTimeout(send, delayMs); + else send(); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const addr = server.address(); + if (addr === null || typeof addr === "string") throw new Error("no port"); + return { + baseUrl: `http://127.0.0.1:${addr.port}`, + received, + close: () => + new Promise((resolve, reject) => + server.close((e) => (e ? reject(e) : resolve())), + ), + }; +} + +function lastRequest(upstream: RecordingUpstream): { + path: string; + body: Record; +} { + const last = upstream.received.at(-1); + expect(last, "upstream received no request").toBeDefined(); + return { path: last!.path, body: JSON.parse(last!.body) }; +} + +function chat(app: SpawnedApp, body: unknown): Promise { + return fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify(body), + }); +} + +describe("chat response_format → Gemini and Bedrock", () => { + let app: SpawnedApp | undefined; + let gemini: RecordingUpstream | undefined; + let bedrock: RecordingUpstream | undefined; + let slowBedrock: RecordingUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + gemini = await startJsonUpstream(() => ({ + candidates: [ + { + content: { role: "model", parts: [{ text: ANSWER }] }, + finishReason: "STOP", + }, + ], + usageMetadata: { + promptTokenCount: 7, + candidatesTokenCount: 11, + totalTokenCount: 18, + }, + })); + // One Bedrock endpoint serves both routes: the Anthropic Messages + // envelope on `/invoke`, the Converse envelope on `/converse`. + bedrock = await startJsonUpstream((path) => + path.endsWith("/invoke") + ? { + id: "msg_bedrock_json", + type: "message", + role: "assistant", + model: "claude-sonnet-4-5-20250929", + content: [{ type: "text", text: ANSWER }], + stop_reason: "end_turn", + usage: { input_tokens: 7, output_tokens: 11 }, + } + : { + output: { + message: { + role: "assistant", + content: [ + { + toolUse: { + toolUseId: "tooluse_json", + name: "json_tool_call", + input: { name: "Ada", nickname: "Countess" }, + }, + }, + ], + }, + }, + stopReason: "tool_use", + usage: { inputTokens: 7, outputTokens: 11, totalTokens: 18 }, + metrics: { latencyMs: 1 }, + }, + ); + + // Answers the Converse route with the synthetic tool call, but only + // after longer than the streaming chunk-gap budget seeded below. + slowBedrock = await startJsonUpstream( + () => ({ + output: { + message: { + role: "assistant", + content: [ + { + toolUse: { + toolUseId: "tooluse_slow", + name: "json_tool_call", + input: { name: "Ada", nickname: "Countess" }, + }, + }, + ], + }, + }, + stopReason: "tool_use", + usage: { inputTokens: 7, outputTokens: 11, totalTokens: 18 }, + metrics: { latencyMs: 1 }, + }), + SLOW_UPSTREAM_MS, + ); + + app = await spawnApp(); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const vertexPk = await seed.createProviderKey({ + display_name: "structured-vertex-pk", + provider: "google", + adapter: "vertex", + secret: JSON.stringify({ + access_token: "ya29.structured-e2e", + project: "proj-e2e", + region: "us-central1", + }), + api_base: gemini.baseUrl, + }); + // The aliases are deliberately not the upstream names: every gate + // here reads the model the gateway dispatches to, never what the + // caller typed. + await seed.createModel({ + display_name: "json-gemini", + provider: "google", + model_name: "gemini-2.5-flash", + provider_key_id: vertexPk.id, + }); + await seed.createModel({ + display_name: "json-gemini-legacy", + provider: "google", + model_name: "gemini-1.5-pro", + provider_key_id: vertexPk.id, + }); + + const bedrockPk = await seed.createProviderKey({ + display_name: "structured-bedrock-pk", + provider: "bedrock", + adapter: "bedrock", + secret: JSON.stringify({ + access_key_id: "AKIA-structured-e2e", + secret_access_key: "sk-structured-e2e", + region: "us-west-2", + }), + api_base: bedrock.baseUrl, + }); + await seed.createModel({ + display_name: "json-claude-bedrock", + provider: "bedrock", + model_name: "anthropic.claude-sonnet-4-5-20250929-v1:0", + provider_key_id: bedrockPk.id, + }); + await seed.createModel({ + display_name: "json-nova", + provider: "bedrock", + model_name: "amazon.nova-pro-v1:0", + provider_key_id: bedrockPk.id, + }); + + const slowPk = await seed.createProviderKey({ + display_name: "structured-slow-bedrock-pk", + provider: "bedrock", + adapter: "bedrock", + secret: JSON.stringify({ + access_key_id: "AKIA-structured-slow", + secret_access_key: "sk-structured-slow", + region: "us-west-2", + }), + api_base: slowBedrock.baseUrl, + }); + // A chunk-gap budget the completion blows through, beside an + // end-to-end budget it fits inside — the shape an operator sets when + // they want slow-first-token failover but long completions. + await seed.createModel({ + display_name: "json-nova-slow", + provider: "bedrock", + model_name: "amazon.nova-pro-v1:0", + provider_key_id: slowPk.id, + stream_timeout: STREAM_BUDGET_MS, + timeout: REQUEST_BUDGET_MS, + }); + + // Seeded last, so this key authenticating implies the whole seed set + // has reached the gateway's snapshot. + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, + }); + if (res.status !== 200) { + await res.text(); + return false; + } + const body = (await res.json()) as { data?: Array<{ id?: string }> }; + return (body.data ?? []).some((m) => m.id === "json-nova"); + }); + }); + + afterAll(async () => { + await app?.exit(); + await gemini?.close(); + await bedrock?.close(); + await slowBedrock?.close(); + }); + + test("a streaming tool-route request is not cut by the chunk-gap budget", async (ctx) => { + if (!etcdReachable || !app || !slowBedrock) { + ctx.skip(); + return; + } + // The model carries stream_timeout=400ms and timeout=30s, and the + // upstream takes 1.2s. Measured against the streaming budget — which + // is what the bridge's deadline is on a streaming dispatch — this + // call is cut off; measured against the request budget it is fine. + const res = await chat(app, { + model: "json-nova-slow", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + stream: true, + }); + expect(res.status).toBe(200); + expect(res.headers.get("content-type")).toContain("text/event-stream"); + const raw = await res.text(); + const frames = raw + .split("\n") + .filter((l) => l.startsWith("data: ") && !l.includes("[DONE]")) + .map((l) => JSON.parse(l.slice(6))); + const text = frames + .map((f) => f.choices?.[0]?.delta?.content ?? "") + .join(""); + expect(JSON.parse(text)).toEqual({ name: "Ada", nickname: "Countess" }); + // The upstream leg really did run non-streaming on the Converse route. + expect(slowBedrock.received.at(-1)?.path).toMatch(/\/converse$/); + }); + + test("gemini 2+ gets responseMimeType and responseJsonSchema", async (ctx) => { + if (!etcdReachable || !app || !gemini) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-gemini", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + const body = await res.json(); + expect(JSON.parse(body.choices[0].message.content)).toEqual({ + name: "Ada", + nickname: "Countess", + }); + + const sent = lastRequest(gemini); + expect(sent.path).toContain("gemini-2.5-flash:generateContent"); + expect(sent.body.generationConfig.responseMimeType).toBe( + "application/json", + ); + // Ordinary JSON Schema, forwarded as the caller wrote it — the + // caller's `required` included. + expect(sent.body.generationConfig.responseJsonSchema).toEqual( + PERSON_SCHEMA, + ); + expect(sent.body.generationConfig.responseSchema).toBeUndefined(); + // The OpenAI spelling has no Gemini counterpart; forwarding it 400s. + expect(sent.body.response_format).toBeUndefined(); + }); + + test("gemini 1.x gets the OpenAPI-flavoured responseSchema", async (ctx) => { + if (!etcdReachable || !app || !gemini) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-gemini-legacy", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + + const sent = lastRequest(gemini); + expect(sent.path).toContain("gemini-1.5-pro:generateContent"); + const gc = sent.body.generationConfig; + expect(gc.responseMimeType).toBe("application/json"); + expect(gc.responseJsonSchema).toBeUndefined(); + expect(gc.responseSchema.type).toBe("OBJECT"); + expect(gc.responseSchema.properties.name.type).toBe("STRING"); + expect(gc.responseSchema.propertyOrdering).toEqual(["name", "nickname"]); + expect(gc.responseSchema.required).toEqual(["name"]); + }); + + test("a bedrock claude 4.5 gets output_config.format on the messages wire", async (ctx) => { + if (!etcdReachable || !app || !bedrock) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-claude-bedrock", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.choices[0].message.content).toBe(ANSWER); + + const sent = lastRequest(bedrock); + expect(sent.path).toMatch(/\/invoke$/); + expect(sent.body.output_config.format.type).toBe("json_schema"); + // Sealed, because Bedrock rejects an open object — but the caller's + // optional `nickname` is still optional. + expect(sent.body.output_config.format.schema.additionalProperties).toBe( + false, + ); + expect(sent.body.output_config.format.schema.required).toEqual(["name"]); + expect(sent.body.response_format).toBeUndefined(); + expect(sent.body.tools).toBeUndefined(); + }); + + test("a bedrock nova gets the synthetic tool and its call comes back as JSON", async (ctx) => { + if (!etcdReachable || !app || !bedrock) { + ctx.skip(); + return; + } + const res = await chat(app, { + model: "json-nova", + messages: [{ role: "user", content: "who is Ada" }], + response_format: RESPONSE_FORMAT, + }); + expect(res.status).toBe(200); + + const sent = lastRequest(bedrock); + expect(sent.path).toMatch(/\/converse$/); + expect(sent.body.outputConfig).toBeUndefined(); + const tools = sent.body.toolConfig.tools; + expect(tools).toHaveLength(1); + expect(tools[0].toolSpec.name).toBe("json_tool_call"); + expect(tools[0].toolSpec.inputSchema.json.additionalProperties).toBe(false); + // Nova is one of the two families whose Converse honours a + // `toolChoice`, so the tool is forced rather than merely offered. + expect(sent.body.toolConfig.toolChoice).toEqual({ + tool: { name: "json_tool_call" }, + }); + + // The caller never offered a tool, so it must not be told the model + // stopped to call one. + const body = await res.json(); + const message = body.choices[0].message; + expect(JSON.parse(message.content)).toEqual({ + name: "Ada", + nickname: "Countess", + }); + expect(message.tool_calls).toBeUndefined(); + expect(body.choices[0].finish_reason).toBe("stop"); + }); +});