This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 3b8f6acb [integrations][watsonx] Support watsonx.ai native structured
output (#1117)
3b8f6acb is described below
commit 3b8f6acbfa6b9d61ad7f22027c4abaa9ea7481ea
Author: Weiqing Yang <[email protected]>
AuthorDate: Mon Sep 14 20:03:44 2026 -0700
[integrations][watsonx] Support watsonx.ai native structured output (#1117)
Generated-by: Claude Code 2.1.260 (Claude Opus 5)
---
integrations/chat-models/watsonx/pom.xml | 11 +
.../watsonx/WatsonxChatModelConnection.java | 165 +++++++++-
.../watsonx/WatsonxChatModelConnectionTest.java | 317 ++++++++++++++++--
.../watsonx/WatsonxChatModelLiveTest.java | 32 ++
.../tests/test_watsonx_native_structured_output.py | 358 +++++++++++++++++++++
.../chat_models/watsonx/watsonx_chat_model.py | 129 +++++++-
6 files changed, 981 insertions(+), 31 deletions(-)
diff --git a/integrations/chat-models/watsonx/pom.xml
b/integrations/chat-models/watsonx/pom.xml
index 864426ff..c71baae3 100644
--- a/integrations/chat-models/watsonx/pom.xml
+++ b/integrations/chat-models/watsonx/pom.xml
@@ -46,6 +46,17 @@ under the License.
<artifactId>jackson-databind</artifactId>
</dependency>
+ <!-- Versions managed by the victools BOM imported in the root pom. -->
+ <dependency>
+ <groupId>com.github.victools</groupId>
+ <artifactId>jsonschema-generator</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>com.github.victools</groupId>
+ <artifactId>jsonschema-module-jackson</artifactId>
+ </dependency>
+
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-annotations</artifactId>
diff --git
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
index 1c1606c5..131d94cb 100644
---
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
+++
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
@@ -24,6 +24,13 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.github.victools.jsonschema.generator.Option;
+import com.github.victools.jsonschema.generator.OptionPreset;
+import com.github.victools.jsonschema.generator.SchemaGenerator;
+import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder;
+import com.github.victools.jsonschema.generator.SchemaVersion;
+import com.github.victools.jsonschema.generator.impl.PropertySortUtils;
+import com.github.victools.jsonschema.module.jackson.JacksonModule;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
@@ -46,6 +53,7 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.regex.Matcher;
@@ -236,14 +244,65 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
return value;
}
+ /**
+ * Whether watsonx.ai can constrain generation to a schema for {@code
effectiveModel}.
+ *
+ * <p>Always {@code true}, and deliberately independent of the argument.
The constraint is
+ * applied by the serving runtime, through vLLM guided decoding, rather
than by a per-model
+ * capability. IBM states that chat API support requires the vLLM runtime,
on its pages about
+ * inferencing <i>custom</i> foundation models; it does not state in so
many words that its own
+ * provided models are served the same way. That last step rests instead
on the chat request
+ * body exposing vLLM's guided-decoding parameters verbatim.
+ *
+ * <p>There is no allowlist because IBM publishes no per-model
structured-output signal to key
+ * on: the foundation-model comparison table has columns for chat and tool
interaction and none
+ * for structured output, and the model-specs API exposes no such flag. A
list written here
+ * would encode a gate nobody documents and would report not-capable for
models that do work.
+ *
+ * <p>This diverges from the base contract, which describes capability as
model-dependent and
+ * requires an unrecognized model to report {@code false}. That rule
guards against failing at
+ * the provider for a model whose capability is unknown; here capability
is a property of the
+ * endpoint rather than of the model, so there is no unknown to guard
against.
+ *
+ * <p>Reads no instance state, so capability stays answerable
independently of how the
+ * connection was configured.
+ */
+ @Override
+ protected boolean supportsNativeStructuredOutput(String effectiveModel) {
+ return true;
+ }
+
@Override
public ChatMessage chat(
List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
modelParams) {
+ return doChat(messages, tools, modelParams, null);
+ }
+
+ /**
+ * Translates {@code outputSchema} into watsonx.ai's native {@code
response_format} field when
+ * it is a POJO {@link Class}. Any other schema form — notably a {@code
RowTypeInfo} wrapped in
+ * {@code OutputSchema} — has no native translation here and leaves the
request unconstrained,
+ * so that the prompt-engineering fallback still governs the response.
+ */
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
+ return doChat(messages, tools, modelParams, outputSchema);
+ }
+
+ private ChatMessage doChat(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
try {
final String modelName = (String) modelParams.get("model");
final boolean extractReasoning =
Boolean.TRUE.equals(modelParams.get("extract_reasoning"));
- final ObjectNode payload = buildPayload(messages, tools,
modelParams);
+ final ObjectNode payload = buildPayload(messages, tools,
modelParams, outputSchema);
if (projectId != null && !projectId.isEmpty()) {
payload.put("project_id", projectId);
} else {
@@ -397,10 +456,14 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
}
@VisibleForTesting
- static ObjectNode buildPayload(
- List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
modelParams) {
+ ObjectNode buildPayload(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
final ObjectNode payload = MAPPER.createObjectNode();
- payload.put("model_id", (String) modelParams.get("model"));
+ final String modelName = (String) modelParams.get("model");
+ payload.put("model_id", modelName);
payload.set("messages", convertMessages(messages));
if (tools != null && !tools.isEmpty()) {
@@ -418,6 +481,40 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
@SuppressWarnings("unchecked")
final Map<String, Object> additionalKwargs =
(Map<String, Object>) modelParams.get("additional_kwargs");
+
+ // Native structured output applies only for a POJO Class schema; any
other schema form,
+ // such as a RowTypeInfo wrapped in OutputSchema, keeps the
prompt-engineering fallback.
+ // The derived schema is a request field of its own rather than a
sampling option, so it is
+ // written at the payload root. When no native translation applies the
key stays absent
+ // rather than being written as a null, which would still be a present
field on the wire.
+ //
+ // TODO(#912): the requested strategy is not visible here, so a
request that explicitly
+ // asked for NATIVE cannot be told apart from one that merely resolved
to it. Capability is
+ // unconditional on this connection, so the schema form is the only
way through to an
+ // unconstrained response: a caller who asked for NATIVE and passed a
schema this branch
+ // cannot translate gets one silently. Once strategy resolution is
wired up, NATIVE must
+ // either bypass this re-check or fail explicitly.
+ if (outputSchema instanceof Class &&
supportsNativeStructuredOutput(modelName)) {
+ // A caller reaches the same payload field through either channel.
Only the branch that
+ // actually sends a derived schema may reject the caller's value;
every path that skips
+ // it leaves that value untouched. A null is not a conflict,
because both write loops
+ // below skip null values and it therefore never reaches the
payload.
+ final boolean callerSuppliedResponseFormat =
+ modelParams.get("response_format") != null
+ || (additionalKwargs != null
+ && additionalKwargs.get("response_format")
!= null);
+ if (callerSuppliedResponseFormat) {
+ throw new IllegalArgumentException(
+ "The "
+ + ((Class<?>) outputSchema).getSimpleName()
+ + " output schema is sent as response_format,
so response_format"
+ + " must not also be set in the model
parameters or in"
+ + " additional_kwargs. Remove that value, or
omit the output"
+ + " schema to set response_format directly.");
+ }
+ payload.set("response_format", toNativeResponseFormat((Class<?>)
outputSchema));
+ }
+
if (additionalKwargs != null) {
final Set<String> collisions = new
java.util.HashSet<>(additionalKwargs.keySet());
collisions.retainAll(RESERVED_ADDITIONAL_KWARGS);
@@ -449,6 +546,66 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
return payload;
}
+ // Derives the response_format value watsonx.ai expects from a POJO class.
Every setting below
+ // addresses a concrete way the generated schema otherwise misstates the
contract:
+ //
+ // - DRAFT_2020_12 is the draft pydantic generates on the Python side,
so a schema derived
+ // from a Java class states the same contract in the same dialect.
+ // - The PLAIN_JSON preset keeps generation to fields. Without a preset,
getters surface as
+ // properties of their own, named after the accessor call, e.g.
"getDerived()".
+ // - MAP_VALUES_AS_ADDITIONAL_PROPERTIES gives a Map its value schema.
Without it the map
+ // admits any value, and a response that satisfies the schema can
still fail to deserialize
+ // into the declared value type at the caller.
+ // - Sorting fields before methods and applying no further comparison
leaves properties in
+ // declaration order, which is the order pydantic emits, so the
documents the two languages
+ // derive from the same shape stay aligned.
+ // - The required check marks every field required except an Optional
one. The default marks
+ // nothing required, which lets a model omit fields at will, while
marking everything
+ // required would force the fields a caller declared omissible.
+ // - The Jackson module makes the schema name properties the way Jackson
names them, because
+ // a caller deserializing the response into this class honors
@JsonProperty and skips
+ // @JsonIgnore. This connection returns the content as a string and
never deserializes into
+ // the schema class, so a property stated under the wrong name
produces a response that
+ // satisfies the schema and still fails to read back, at the caller
rather than here. It is
+ // applied with no JacksonOption, so it contributes property naming
and visibility only:
+ // the required set and the property order stay the ones configured
below.
+ //
+ // Two settings are deliberately absent:
+ //
+ // - FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT is omitted for
cross-language parity:
+ // pydantic's plain schema does not emit additionalProperties: false,
so emitting it on the
+ // Java side alone would make the two documents disagree. The endpoint
does not require it,
+ // whether or not strict is set.
+ // - DEFINITION_FOR_MAIN_SCHEMA is left off, which is victools' default.
Turning it on only
+ // moves the document root into $defs and makes the root a $ref, and
nothing establishes
+ // that watsonx benefits from either shape, so there is no reason to
deviate. It does not
+ // gate recursion: a self-referential type generates either way, as
{"child": {"$ref": "#"}}
+ // with the option off, and any type used twice is extracted into
$defs regardless. A
+ // recursive POJO reaches the provider unguarded either way, and this
setting cannot
+ // change that.
+ private static ObjectNode toNativeResponseFormat(Class<?> schemaClass) {
+ SchemaGeneratorConfigBuilder configBuilder =
+ new SchemaGeneratorConfigBuilder(
+ SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
+ .with(Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES)
+ .with(new JacksonModule());
+ configBuilder
+ .forTypesInGeneral()
+
.withPropertySorter(PropertySortUtils.SORT_PROPERTIES_FIELDS_BEFORE_METHODS);
+ configBuilder
+ .forFields()
+ .withRequiredCheck(field ->
!Optional.class.equals(field.getRawMember().getType()));
+
+ final ObjectNode responseFormat = MAPPER.createObjectNode();
+ responseFormat.put("type", "json_schema");
+ final ObjectNode jsonSchema = responseFormat.putObject("json_schema");
+ jsonSchema.put("name", schemaClass.getSimpleName());
+ jsonSchema.set(
+ "schema", new
SchemaGenerator(configBuilder.build()).generateSchema(schemaClass));
+ jsonSchema.put("strict", true);
+ return responseFormat;
+ }
+
/**
* Converts framework chat messages to the watsonx.ai (OpenAI-compatible)
message format.
*
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
index f594e932..5dda6f7b 100644
---
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
@@ -17,6 +17,8 @@
*/
package org.apache.flink.agents.integrations.chatmodels.watsonx;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -33,14 +35,18 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.NullAndEmptySource;
import org.junit.jupiter.params.provider.ValueSource;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Stream;
@@ -59,6 +65,42 @@ class WatsonxChatModelConnectionTest {
private static final String CHAT_RESPONSE =
"{\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\","
+ "\"content\":\"Hello!\"},\"finish_reason\":\"stop\"}]}";
+ private static final String DRAFT_2020_12 =
"https://json-schema.org/draft/2020-12/schema";;
+ private static final String MODEL = "ibm/granite-3-3-8b-instruct";
+ private static final Map<String, Object> CALLER_FORMAT = Map.of("type",
"json_object");
+
+ /**
+ * Output schema fixture shaped to expose the schema-generation settings.
+ *
+ * <p>Fields are declared out of alphabetical order, {@code counts} is a
map whose values carry
+ * a type, {@code note} is the only optional field, and {@code getDerived}
is a getter backed by
+ * no field.
+ */
+ public static class Report {
+ public String summary;
+ public Map<String, Integer> counts;
+ public Optional<String> note;
+ public int total;
+
+ public String getDerived() {
+ return summary + total;
+ }
+ }
+
+ /**
+ * Output schema fixture shaped to expose Jackson's property model.
+ *
+ * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
+ * and {@code secret} is not deserialized at all.
+ */
+ public static class Profile {
+ @JsonProperty("full_name")
+ public String name;
+
+ @JsonIgnore public String secret;
+
+ public int age;
+ }
private static ResourceDescriptor descriptor(String url, String apiKey,
String projectId) {
ResourceDescriptor.Builder b =
@@ -115,7 +157,44 @@ class WatsonxChatModelConnectionTest {
return connection.chat(
List.of(new ChatMessage(MessageRole.USER, "Hello!")),
List.of(),
- Map.of("model", "ibm/granite-3-3-8b-instruct"));
+ Map.of("model", MODEL));
+ }
+
+ /** A connection that is never sent a request, for the payload-building
tests. */
+ private static WatsonxChatModelConnection connection() {
+ return new WatsonxChatModelConnection(
+ descriptor("https://us-south.ml.cloud.ibm.com";, "test-key",
"test-project"),
+ NOOP,
+ NO_ENVIRONMENT);
+ }
+
+ private static ObjectNode payloadFor(Object outputSchema) {
+ return payloadFor(Map.of("model", MODEL), outputSchema);
+ }
+
+ private static ObjectNode payloadFor(Map<String, Object> modelParams,
Object outputSchema) {
+ return connection()
+ .buildPayload(
+ List.of(new ChatMessage(MessageRole.USER, "Hello!")),
+ List.of(),
+ modelParams,
+ outputSchema);
+ }
+
+ private static JsonNode derivedSchema(Object outputSchema) {
+ return
payloadFor(outputSchema).path("response_format").path("json_schema").path("schema");
+ }
+
+ private static List<String> fieldNames(JsonNode objectNode) {
+ List<String> names = new ArrayList<>();
+ objectNode.fieldNames().forEachRemaining(names::add);
+ return names;
+ }
+
+ private static List<String> textValues(JsonNode arrayNode) {
+ List<String> values = new ArrayList<>();
+ arrayNode.forEach(element -> values.add(element.asText()));
+ return values;
}
@ParameterizedTest(name = "{0}")
@@ -303,12 +382,10 @@ class WatsonxChatModelConnectionTest {
@DisplayName("Model params are copied top-level into the payload")
void testBuildPayload() {
ObjectNode payload =
- WatsonxChatModelConnection.buildPayload(
- List.of(new ChatMessage(MessageRole.USER, "Hello!")),
- List.of(),
+ payloadFor(
Map.of(
"model",
- "ibm/granite-3-3-8b-instruct",
+ MODEL,
"temperature",
0.5,
"max_tokens",
@@ -318,9 +395,10 @@ class WatsonxChatModelConnectionTest {
"extract_reasoning",
true,
"additional_kwargs",
- Map.of("top_p", 0.9)));
+ Map.of("top_p", 0.9)),
+ null);
-
assertThat(payload.get("model_id").asText()).isEqualTo("ibm/granite-3-3-8b-instruct");
+ assertThat(payload.get("model_id").asText()).isEqualTo(MODEL);
assertThat(payload.get("temperature").asDouble()).isEqualTo(0.5);
assertThat(payload.get("max_tokens").asInt()).isEqualTo(256);
assertThat(payload.get("top_p").asDouble()).isEqualTo(0.5);
@@ -332,14 +410,13 @@ class WatsonxChatModelConnectionTest {
assertThatThrownBy(
() ->
- WatsonxChatModelConnection.buildPayload(
- List.of(new
ChatMessage(MessageRole.USER, "Hello!")),
- List.of(),
+ payloadFor(
Map.of(
"model",
- "ibm/granite-3-3-8b-instruct",
+ MODEL,
"additional_kwargs",
- Map.of("temperature", 5.0))))
+ Map.of("temperature", 5.0)),
+ null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("additional_kwargs")
.hasMessageContaining("temperature");
@@ -348,14 +425,9 @@ class WatsonxChatModelConnectionTest {
List.of("model_id", "messages", "tools", "project_id",
"space_id")) {
assertThatThrownBy(
() ->
- WatsonxChatModelConnection.buildPayload(
- List.of(new
ChatMessage(MessageRole.USER, "Hello!")),
- List.of(),
- Map.of(
- "model",
-
"ibm/granite-3-3-8b-instruct",
- requestOwnedField,
- "override")))
+ payloadFor(
+ Map.of("model", MODEL,
requestOwnedField, "override"),
+ null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining(requestOwnedField);
}
@@ -672,4 +744,209 @@ class WatsonxChatModelConnectionTest {
assertThat(function.get("name")).isEqualTo("add");
assertThat(function.get("arguments")).isEqualTo(Map.of("a", 1, "b",
2));
}
+
+ @ParameterizedTest
+ @NullAndEmptySource
+ @ValueSource(
+ strings = {
+ "ibm/granite-3-3-8b-instruct",
+ "meta-llama/llama-3-3-70b-instruct",
+ "mistralai/mistral-large",
+ " "
+ })
+ @DisplayName("Capability is reported for any model, since the endpoint
provides it")
+ void supportsNativeStructuredOutputIsUnconditional(String model) {
+ // Capability comes from the serving runtime the chat API requires
rather than from the
+ // model, so a model allowlist here would silently drop back to the
prompt fallback for
+ // anything the list had not caught up with. Null and blank are
included because the
+ // answer does not depend on the argument at all.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isTrue();
+ }
+
+ @Test
+ @DisplayName("A POJO output schema is sent as a native json_schema
response format")
+ void buildPayloadWritesResponseFormatForPojoSchema() {
+ ObjectNode payload = payloadFor(Report.class);
+
+ JsonNode responseFormat = payload.path("response_format");
+
assertThat(responseFormat.path("type").asText()).isEqualTo("json_schema");
+ JsonNode jsonSchema = responseFormat.path("json_schema");
+ assertThat(jsonSchema.path("name").asText()).isEqualTo("Report");
+ assertThat(jsonSchema.path("strict").booleanValue()).isTrue();
+ assertThat(fieldNames(jsonSchema.path("schema").path("properties")))
+ .containsExactlyInAnyOrder("summary", "counts", "note",
"total");
+ // The draft is the one pydantic emits on the Python side, so both
languages state the
+ // same contract in the same dialect. Nothing else in this suite pins
the draft.
+
assertThat(jsonSchema.path("schema").path("$schema").asText()).isEqualTo(DRAFT_2020_12);
+ }
+
+ @Test
+ @DisplayName("No output schema leaves the payload without a response
format")
+ void buildPayloadOmitsResponseFormatWithoutSchema() {
+ ObjectNode payload = payloadFor(null);
+
+ // A null is still a present response_format field in the serialized
body, so the key has
+ // to be absent rather than written as null.
+ assertThat(payload.has("response_format")).isFalse();
+ }
+
+ @Test
+ @DisplayName("A RowTypeInfo-shaped schema stays on the prompt fallback")
+ void buildPayloadOmitsResponseFormatForRowTypeInfo() {
+ // A RowTypeInfo schema arrives wrapped rather than as a bare POJO
Class, so it must not
+ // activate native structured output. The wrapper cannot be built here
because RowTypeInfo
+ // is not on this module's classpath; any non-Class schema object
exercises the same gate.
+ // Deriving a schema from it must degrade to the fallback instead of
failing the request.
+ Object nonClassSchema = "row<name STRING>";
+
+ ObjectNode payload = payloadFor(nonClassSchema);
+
+ assertThat(payload.has("response_format")).isFalse();
+ }
+
+ @Test
+ @DisplayName("The derived schema names properties the way Jackson
deserializes them")
+ void derivedSchemaHonorsJacksonAnnotations() {
+ JsonNode schema = derivedSchema(Profile.class);
+
+ // A caller deserializing the response into this class accepts the
renamed property and
+ // rejects the Java field name, and discards an ignored property that
the schema would
+ // otherwise state as required and so force a value for.
+ assertThat(fieldNames(schema.path("properties")))
+ .containsExactlyInAnyOrder("full_name", "age");
+ assertThat(textValues(schema.path("required")))
+ .containsExactlyInAnyOrder("full_name", "age");
+ }
+
+ @Test
+ @DisplayName("The derived schema requires every field the caller did not
make optional")
+ void derivedSchemaMarksNonOptionalFieldsRequired() {
+ JsonNode schema = derivedSchema(Report.class);
+
+ // Without a required set the model may omit fields the caller
declared, and with an
+ // all-inclusive one it must invent a value for the field the caller
made omissible.
+ assertThat(textValues(schema.path("required")))
+ .containsExactlyInAnyOrder("summary", "counts", "total");
+ }
+
+ @Test
+ @DisplayName("The derived schema gives map values their own schema")
+ void derivedSchemaGivesMapValuesTheirSchema() {
+ JsonNode counts =
derivedSchema(Report.class).path("properties").path("counts");
+
+ // A map without a value schema admits any value, so a response can
satisfy the schema and
+ // still fail to deserialize into the declared value type at the
caller.
+ assertThat(counts.path("additionalProperties").isObject()).isTrue();
+
assertThat(counts.path("additionalProperties").path("type").asText()).isEqualTo("integer");
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("callerSuppliedResponseFormats")
+ @DisplayName("A caller response format alongside an output schema is
rejected")
+ void chatWithSchemaAndCallerResponseFormatRaises(
+ String ignoredCaseName, Map<String, Object> modelParams) throws
Exception {
+ HttpServer server = startServer();
+ server.createContext(
+ "/ml/v1/text/chat", exchange -> sendJson(exchange, 200,
CHAT_RESPONSE));
+
+ try {
+ WatsonxChatModelConnection connection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), false, 0), NOOP,
NO_ENVIRONMENT);
+
+ // Both channels land in the same payload field as the derived
schema, so letting
+ // either side win silently sends a format the caller never asked
for, or drops the
+ // schema the agent depends on to parse the reply. The stub
answers the request, so a
+ // guard that stops rejecting the combination fails here rather
than hanging on a call
+ // to the real endpoint.
+ assertThatThrownBy(
+ () ->
+ connection.chat(
+ List.of(new
ChatMessage(MessageRole.USER, "Hello!")),
+ List.of(),
+ modelParams,
+ Report.class))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("response_format");
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ private static Stream<Arguments> callerSuppliedResponseFormats() {
+ return Stream.of(
+ Arguments.of("model params", callerFormatInModelParams()),
+ Arguments.of("additional_kwargs",
callerFormatInAdditionalKwargs()));
+ }
+
+ private static Map<String, Object> callerFormatInModelParams() {
+ return Map.of("model", MODEL, "response_format", CALLER_FORMAT);
+ }
+
+ private static Map<String, Object> callerFormatInAdditionalKwargs() {
+ return Map.of(
+ "model", MODEL, "additional_kwargs", Map.of("response_format",
CALLER_FORMAT));
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("skippedNativePathsWithCallerResponseFormat")
+ @DisplayName("A caller response format is forwarded when no schema is
translated")
+ void callerResponseFormatSurvivesWhenNoSchemaIsSent(
+ String ignoredCaseName, Map<String, Object> modelParams, Object
outputSchema) {
+ // response_format is on none of the reserved sets, so callers set it
directly today. The
+ // rejection belongs to the branch that actually derives a schema:
applied any wider it
+ // turns every one of those existing calls into an error.
+ ObjectNode payload = payloadFor(modelParams, outputSchema);
+
+
assertThat(payload.path("response_format")).isEqualTo(MAPPER.valueToTree(CALLER_FORMAT));
+ }
+
+ private static Stream<Arguments>
skippedNativePathsWithCallerResponseFormat() {
+ return Stream.of(
+ Arguments.of("no schema, model params",
callerFormatInModelParams(), null),
+ Arguments.of(
+ "no schema, additional_kwargs",
callerFormatInAdditionalKwargs(), null),
+ // A schema the branch cannot translate skips it for the other
reason, so the
+ // caller's value has to survive that arm too.
+ Arguments.of(
+ "untranslatable schema, model params",
+ callerFormatInModelParams(),
+ "row<name STRING>"));
+ }
+
+ @Test
+ @DisplayName("The serialized request body carries response_format at the
document root")
+ void serializedRequestBodyCarriesResponseFormatAtRoot() throws Exception {
+ HttpServer server = startServer();
+ AtomicReference<JsonNode> captured = new AtomicReference<>();
+ server.createContext(
+ "/ml/v1/text/chat",
+ exchange -> {
+ captured.set(MAPPER.readTree(exchange.getRequestBody()));
+ sendJson(exchange, 200, CHAT_RESPONSE);
+ });
+
+ try {
+ WatsonxChatModelConnection connection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), false, 0), NOOP,
NO_ENVIRONMENT);
+
+ connection.chat(
+ List.of(new ChatMessage(MessageRole.USER, "Hello!")),
+ List.of(),
+ Map.of("model", MODEL),
+ Report.class);
+
+ // The scope id is injected after the payload is built, so only
the wire body proves
+ // that a schema handed to chat reaches the request at all, and
that it sits at the
+ // root beside the other request fields rather than nested inside
one of them.
+ JsonNode body = captured.get();
+ assertThat(body).isNotNull();
+ assertThat(body.has("response_format")).isTrue();
+
assertThat(body.path("response_format").path("json_schema").path("name").asText())
+ .isEqualTo("Report");
+ } finally {
+ server.stop(0);
+ }
+ }
}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
index 50197e18..5baaabc6 100644
---
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.flink.agents.integrations.chatmodels.watsonx;
+import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.resource.ResourceContext;
@@ -42,6 +43,12 @@ class WatsonxChatModelLiveTest {
private static final ResourceContext NOOP =
ResourceContext.fromGetResource((a, b) -> null);
+ /** Flat output schema for the live structured-output call. */
+ public static class Answer {
+ public String verdict;
+ public int score;
+ }
+
static boolean credentialsAvailable() {
return isSet("WATSONX_URL")
&& (isSet("WATSONX_API_KEY") || isSet("WATSONX_TOKEN"))
@@ -83,4 +90,29 @@ class WatsonxChatModelLiveTest {
assertThat(response.getExtraArgs().get("promptTokens")).isNotNull();
assertThat(response.getExtraArgs().get("completionTokens")).isNotNull();
}
+
+ @Test
+ @DisplayName("A POJO output schema comes back as content that deserializes
into it")
+ void testChatWithOutputSchema() throws Exception {
+ ChatMessage response =
+ connection()
+ .chat(
+ List.of(
+ new ChatMessage(
+ MessageRole.USER,
+ "Rate the sentence \"the build
is green\" and"
+ + " report a verdict
and a score.")),
+ List.of(),
+ Map.of("model", model(), "max_tokens", 200),
+ Answer.class);
+
+ // This is the only check that the schema was applied rather than
answered with prose. A
+ // reply that does not deserialize is a result to investigate rather
than a diagnosis on
+ // its own: the schema may not have been enforced, or may not have
described the type the
+ // caller passed, or the reply may carry a property the schema permits
— no
+ // additionalProperties: false is sent — but this reader rejects.
+ assertThat(response.getContent()).isNotBlank();
+ Answer answer = new ObjectMapper().readValue(response.getContent(),
Answer.class);
+ assertThat(answer.verdict).isNotNull();
+ }
}
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_native_structured_output.py
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_native_structured_output.py
new file mode 100644
index 00000000..92a8ca18
--- /dev/null
+++
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_native_structured_output.py
@@ -0,0 +1,358 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import json
+import os
+from typing import Any, Callable, Dict, Tuple
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic import BaseModel
+from pyflink.common.typeinfo import BasicTypeInfo, RowTypeInfo
+
+from flink_agents.api.agents.types import OutputSchema
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.integrations.chat_models.watsonx.watsonx_chat_model import (
+ DEFAULT_MODEL,
+ WatsonxChatModelConnection,
+)
+
+test_model = os.environ.get("WATSONX_CHAT_MODEL", DEFAULT_MODEL)
+credentials_available = (
+ "WATSONX_URL" in os.environ
+ and ("WATSONX_API_KEY" in os.environ or "WATSONX_TOKEN" in os.environ)
+) and ("WATSONX_PROJECT_ID" in os.environ or "WATSONX_SPACE_ID" in os.environ)
+
+
+class Report(BaseModel):
+ """Output schema fixture, shaped like the Java suite's fixture of the same
name.
+
+ The shapes are held in step so the two languages can be compared: a plain
field,
+ a map with a typed value, an omissible field, and a second plain field.
+ """
+
+ summary: str
+ counts: Dict[str, int]
+ note: str | None = None
+ total: int
+
+
+class Unrenderable(BaseModel):
+ """A schema carrying a member that no JSON Schema can express."""
+
+ cb: Callable[[int], int]
+
+
+class Answer(BaseModel):
+ """Flat output schema for the live structured-output call."""
+
+ verdict: str
+ score: int
+
+
+# What the Java suite pins for the Report shape. The two documents are
compared on
+# property names, the required set and value types rather than byte for byte:
+# victools always emits $schema and never title, and pydantic is the reverse.
+JAVA_PROPERTY_NAMES = {"summary", "counts", "note", "total"}
+JAVA_REQUIRED = {"summary", "counts", "total"}
+
+CALLER_FORMAT = {"type": "json_object"}
+
+CHAT_RESPONSE = {
+ "id": "chatcmpl-1",
+ "model_id": test_model,
+ "choices": [
+ {
+ "index": 0,
+ "message": {"role": "assistant", "content": "{}"},
+ "finish_reason": "stop",
+ }
+ ],
+}
+
+
+def _connection() -> WatsonxChatModelConnection:
+ """A connection built from fake credentials, contacting nothing on its
own."""
+ return WatsonxChatModelConnection(
+ url="https://us-south.ml.cloud.ibm.com";,
+ api_key="fake-key",
+ project_id="fake-project",
+ )
+
+
+def _mocked_connection(
+ monkeypatch: pytest.MonkeyPatch,
+) -> Tuple[WatsonxChatModelConnection, MagicMock]:
+ """A connection whose provider call is a mock, returned alongside it.
+
+ Assigning the private client is what keeps the call offline: the connection
+ builds a real one lazily on first use otherwise.
+ """
+ provider_model = MagicMock()
+ provider_model.chat.return_value = CHAT_RESPONSE
+ monkeypatch.setattr(
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.ModelInference",
+ MagicMock(return_value=provider_model),
+ )
+ connection = _connection()
+ connection._client = MagicMock()
+ return connection, provider_model
+
+
+def _sent_params(
+ monkeypatch: pytest.MonkeyPatch,
+ output_schema: OutputSchema | None = None,
+ **kwargs: Any,
+) -> Dict[str, Any]:
+ """The params dict one chat call hands to the provider."""
+ connection, provider_model = _mocked_connection(monkeypatch)
+ connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
+ output_schema=output_schema,
+ **kwargs,
+ )
+ return provider_model.chat.call_args.kwargs["params"] or {}
+
+
[email protected]("effective_model", [DEFAULT_MODEL, "no-such-model",
None])
+def test_supports_native_structured_output_is_unconditional(
+ effective_model: str | None,
+) -> None:
+ """Capability is reported for any model, including one nothing recognizes.
+
+ The constraint is applied by the serving runtime rather than by the model,
so a
+ gate on the model string would turn callers away from a path that works.
+ """
+ assert _connection().supports_native_structured_output(effective_model) is
True
+
+
+def test_params_carry_response_format_for_basemodel_schema(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A BaseModel schema reaches the provider as the endpoint's json_schema
form.
+
+ The schema slot is the model's plain rendering. A stricter rendering would
force
+ every field required and forbid additional properties, neither of which
the Java
+ side emits.
+ """
+ params = _sent_params(monkeypatch,
output_schema=OutputSchema(output_schema=Report))
+
+ assert params["response_format"] == {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "Report",
+ "schema": Report.model_json_schema(),
+ "strict": True,
+ },
+ }
+
+
+def test_params_omit_response_format_without_schema(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """No schema leaves the key absent rather than present and empty."""
+ params = _sent_params(monkeypatch, temperature=0.5)
+
+ assert "response_format" not in params
+
+
+def test_row_type_info_schema_returns_and_sends_no_response_format(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A RowTypeInfo schema is accepted and left to the prompt-engineering
fallback.
+
+ Both halves are asserted. Reporting native capability exempts this
connection
+ from the tree-wide rule that a connection rejects a schema it cannot
translate,
+ so nothing else pins that a schema form with no native translation still
returns
+ a response, and that no format is derived from it.
+ """
+ schema = OutputSchema(
+ output_schema=RowTypeInfo([BasicTypeInfo.INT_TYPE_INFO()], ["total"])
+ )
+ connection, provider_model = _mocked_connection(monkeypatch)
+
+ response = connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
+ output_schema=schema,
+ temperature=0.5,
+ )
+
+ assert response.role == MessageRole.ASSISTANT
+ assert "response_format" not in
provider_model.chat.call_args.kwargs["params"]
+
+
+def test_unrenderable_model_reports_render_failure(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A model with no JSON Schema is reported as that, naming the model.
+
+ Rendering through the shared helper is what produces this wording. Calling
the
+ model's own rendering directly would surface an error naming neither the
schema
+ nor a way out of it.
+ """
+ with pytest.raises(
+ TypeError, match="Unrenderable cannot be rendered as a JSON Schema"
+ ):
+ _sent_params(
+ monkeypatch, output_schema=OutputSchema(output_schema=Unrenderable)
+ )
+
+
[email protected](
+ "caller_kwargs",
+ [
+ pytest.param({"response_format": CALLER_FORMAT}, id="kwarg"),
+ pytest.param(
+ {"additional_kwargs": {"response_format": CALLER_FORMAT}},
+ id="additional_kwargs",
+ ),
+ ],
+)
+def test_schema_and_caller_response_format_raises(
+ monkeypatch: pytest.MonkeyPatch, caller_kwargs: Dict[str, Any]
+) -> None:
+ """A schema and a caller response format with a value cannot both be
honored.
+
+ Both channels reach the same request field, so both are guarded. Choosing
the
+ derived value would drop a setting the caller made deliberately, and
choosing the
+ caller's would report the schema as applied when it was not.
+ """
+ with pytest.raises(
+ ValueError, match="Report output schema is sent as response_format"
+ ):
+ _sent_params(
+ monkeypatch,
+ output_schema=OutputSchema(output_schema=Report),
+ **caller_kwargs,
+ )
+
+
[email protected](
+ ("output_schema", "caller_kwargs"),
+ [
+ pytest.param(None, {"response_format": CALLER_FORMAT}, id="no schema,
kwarg"),
+ pytest.param(
+ None,
+ {"additional_kwargs": {"response_format": CALLER_FORMAT}},
+ id="no schema, additional_kwargs",
+ ),
+ pytest.param(
+ OutputSchema(
+ output_schema=RowTypeInfo([BasicTypeInfo.INT_TYPE_INFO()],
["total"])
+ ),
+ {"response_format": CALLER_FORMAT},
+ id="untranslatable schema, kwarg",
+ ),
+ ],
+)
+def test_caller_response_format_survives_when_no_schema_is_sent(
+ monkeypatch: pytest.MonkeyPatch,
+ output_schema: OutputSchema | None,
+ caller_kwargs: Dict[str, Any],
+) -> None:
+ """A caller response format is forwarded whenever no schema is derived.
+
+ response_format is on none of the reserved sets, so callers set it directly
+ today and the rejection belongs to the branch that actually derives a
schema.
+ Applied any wider it turns those calls into errors. The branch is skipped
for two
+ distinct reasons, no schema at all and a schema with no native
translation, and
+ the caller's value has to survive both.
+ """
+ params = _sent_params(monkeypatch, output_schema=output_schema,
**caller_kwargs)
+
+ assert params["response_format"] == CALLER_FORMAT
+
+
+def test_null_caller_response_format_is_not_a_conflict(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A caller response format of None is replaced by the derived one, not
refused.
+
+ Nothing the caller set survives to compete with the schema, so there is no
+ conflict to report and the request still carries the derived envelope.
+ """
+ params = _sent_params(
+ monkeypatch,
+ output_schema=OutputSchema(output_schema=Report),
+ response_format=None,
+ )
+
+ assert params["response_format"]["json_schema"]["name"] == "Report"
+
+
+def test_derived_schema_matches_java_semantically(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """The document sent from Python states the contract the Java one states.
+
+ A caller may reach the same model through either language, so the two must
agree
+ on which properties exist, which of them the response must carry, and what
a map
+ value has to be.
+ """
+ params = _sent_params(monkeypatch,
output_schema=OutputSchema(output_schema=Report))
+ schema = params["response_format"]["json_schema"]["schema"]
+
+ assert set(schema["properties"]) == JAVA_PROPERTY_NAMES
+ assert set(schema["required"]) == JAVA_REQUIRED
+ assert schema["properties"]["summary"]["type"] == "string"
+ assert schema["properties"]["total"]["type"] == "integer"
+ assert schema["properties"]["counts"]["additionalProperties"]["type"] ==
"integer"
+
+
+def test_live_fixture_document_omits_additional_properties() -> None:
+ """The live call's schema states the same contract the Java recipe states.
+
+ A model config forbidding extras would render additionalProperties into the
+ document, so the two languages would send different contracts for one
fixture.
+ Nothing else offline reads what the live call sends.
+ """
+ assert "additionalProperties" not in Answer.model_json_schema()
+
+
[email protected]
[email protected](
+ not credentials_available, reason="watsonx.ai credentials are not set"
+)
+def test_chat_with_output_schema() -> None:
+ """A schema comes back as content that parses into the model it described.
+
+ The reply's own keys are checked, not just that it parses, because nothing
else
+ here would notice an extra one: no additionalProperties is sent, so the
schema
+ permits it, and the reader drops a property it does not declare.
+
+ A reply that fails this is a result to investigate rather than a diagnosis
on its
+ own: the schema may not have been applied, or may not have described the
type the
+ caller passed.
+ """
+ response = WatsonxChatModelConnection().chat(
+ [
+ ChatMessage(
+ role=MessageRole.USER,
+ content='Rate the sentence "the build is green" and report a
verdict'
+ " and a score.",
+ )
+ ],
+ output_schema=OutputSchema(output_schema=Answer),
+ model=test_model,
+ max_tokens=200,
+ )
+
+ assert response.content is not None
+ assert response.content.strip() != ""
+ parsed = json.loads(response.content)
+ assert set(parsed) == {"verdict", "score"}
+ assert Answer(**parsed).verdict is not None
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
index b553e3f9..414ff6f9 100644
--- a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
@@ -28,10 +28,10 @@ import httpx
from ibm_watsonx_ai import APIClient, Credentials
from ibm_watsonx_ai.foundation_models import ModelInference
from ibm_watsonx_ai.wml_client_error import ApiRequestFailure
-from pydantic import Field, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr
from typing_extensions import override
-from flink_agents.api.agents.types import OutputSchema
+from flink_agents.api.agents.types import OutputSchema, render_output_schema
from flink_agents.api.chat_message import ChatMessage, MessageRole
from flink_agents.api.chat_models.chat_model import (
BaseChatModelConnection,
@@ -161,6 +161,49 @@ def _convert_to_watsonx_tool_call(tool_call: Dict[str,
Any]) -> Dict[str, Any]:
}
+def _native_output_model(
+ output_schema: OutputSchema | None,
+) -> type[BaseModel] | None:
+ """The model a schema translates natively to, or ``None`` where none
applies.
+
+ ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
+ translation here and keeps the prompt-engineering fallback.
+
+ Separate from the render below because the caller-conflict check needs to
know
+ whether a schema will be sent, and under what name, before anything is
rendered.
+ """
+ model = (
+ output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
+ )
+ if not (isinstance(model, type) and issubclass(model, BaseModel)):
+ return None
+ return model
+
+
+def _native_response_format(
+ output_schema: OutputSchema | None,
+) -> Dict[str, Any] | None:
+ """Build the ``response_format`` for a native structured-output request.
+
+ Returns ``None``, leaving the request unchanged, unless the schema is a
+ ``BaseModel`` subclass.
+
+ Rendering goes through the shared helper, which raises ``TypeError``
naming the
+ schema class and the remedy when the model has no JSON Schema.
+ """
+ model = _native_output_model(output_schema)
+ if model is None:
+ return None
+ return {
+ "type": "json_schema",
+ "json_schema": {
+ "name": model.__name__,
+ "schema": render_output_schema(model, lambda m:
m.model_json_schema()),
+ "strict": True,
+ },
+ }
+
+
class WatsonxChatModelConnection(BaseChatModelConnection):
"""Connection to the IBM watsonx.ai chat API."""
@@ -332,6 +375,35 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
time.sleep(delay)
attempt += 1
+ @override
+ def supports_native_structured_output(self, effective_model: str | None)
-> bool:
+ """Whether watsonx.ai can constrain generation to a schema for the
given model.
+
+ Always ``True``, and deliberately independent of the argument. The
constraint
+ is applied by the serving runtime, through vLLM guided decoding,
rather than by
+ a per-model capability. IBM states that chat API support requires the
vLLM
+ runtime, on its pages about inferencing *custom* foundation models; it
does not
+ state in so many words that its own provided models are served the
same way.
+ That last step rests instead on the chat request body exposing vLLM's
+ guided-decoding parameters verbatim.
+
+ There is no allowlist because IBM publishes no per-model
structured-output
+ signal to key on: the foundation-model comparison table has columns
for chat
+ and tool interaction and none for structured output, and the
model-specs API
+ exposes no such flag. A list written here would encode a gate nobody
documents
+ and would report not-capable for models that do work.
+
+ This diverges from the base contract, which describes capability as
+ model-dependent and requires an unrecognized model to report
``False``. That
+ rule guards against failing at the provider for a model whose
capability is
+ unknown; here capability is a property of the endpoint rather than of
the
+ model, so there is no unknown to guard against.
+
+ Reads no instance state, so capability stays answerable independently
of how
+ the connection was configured.
+ """
+ return True
+
def chat(
self,
messages: Sequence[ChatMessage],
@@ -341,16 +413,18 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
) -> ChatMessage:
"""Process a sequence of messages, and return a response.
- A non-``None`` ``output_schema`` is rejected: this connection has no
native
- structured-output translation, so callers stay on the
prompt-engineering
- fallback. Declaring the parameter keeps a caller-supplied schema out of
- ``**kwargs``, which is forwarded to the provider SDK.
+ A ``BaseModel`` ``output_schema`` is sent as the ``response_format``
request
+ parameter. Any other schema form, notably a ``RowTypeInfo``, has no
native
+ translation here and keeps the prompt-engineering fallback. Where the
schema
+ is sent natively, a caller-supplied ``response_format`` carrying a
value
+ conflicts with it and raises ``ValueError``. Declaring the parameter
keeps a
+ caller-supplied schema out of ``**kwargs``, which is forwarded to the
+ provider SDK.
When the response carries a finish reason, it is available verbatim as
``extra_args["finish_reason"]``; the key is absent when the provider
reports none.
"""
- self._reject_unsupported_output_schema(output_schema)
model_name = kwargs.pop("model", DEFAULT_MODEL)
extract_reasoning = bool(kwargs.pop("extract_reasoning", False))
tool_choice = kwargs.pop("tool_choice", None)
@@ -373,6 +447,47 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
)
raise ValueError(msg)
+ # Native structured output applies only for a BaseModel schema; any
other
+ # schema form, such as a RowTypeInfo wrapped in OutputSchema, keeps the
+ # prompt-engineering fallback.
+ #
+ # TODO(#912): the requested strategy is not visible here, so a request
that
+ # explicitly asked for NATIVE cannot be told apart from one that merely
+ # resolved to it. Capability is unconditional on this connection, so
the schema
+ # form is the only way through to an unconstrained response: a caller
who asked
+ # for NATIVE and passed a schema this branch cannot translate gets one
+ # silently. Once strategy resolution is wired up, NATIVE must either
bypass
+ # this re-check or fail explicitly.
+ if output_schema is not None and
self.supports_native_structured_output(
+ model_name
+ ):
+ native_model = _native_output_model(output_schema)
+ # A caller reaches the same request field through either channel,
and both
+ # have already merged into request_params. Only the branch that
sends a
+ # derived schema may reject the caller's value; every path that
skips it
+ # leaves that value untouched. A None is not a conflict, because
the
+ # assignment below replaces it before the request is built. Tested
before
+ # the schema is rendered, so a caller who supplied both is told
about the
+ # conflict rather than about a render failure; the name is read
off the
+ # model class and needs no rendered document.
+ if (
+ native_model is not None
+ and request_params.get("response_format") is not None
+ ):
+ msg = (
+ f"The {native_model.__name__} output schema is sent as"
+ " response_format, so response_format must not also be
passed as"
+ " a kwarg or in additional_kwargs. Remove that value, or
omit"
+ " output_schema to set response_format directly."
+ )
+ raise ValueError(msg)
+ response_format = _native_response_format(output_schema)
+ if response_format is not None:
+ # The SDK merges params into the request body at its root, so
the
+ # derived schema travels as the request field it is rather
than as a
+ # sampling option.
+ request_params["response_format"] = response_format
+
tool_specs: List[Dict[str, Any]] | None = (
[to_openai_tool(metadata=tool.metadata) for tool in tools]
if tools
