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 85bbb543 [integrations][java][python] Apply Ollama native structured
output (#981)
85bbb543 is described below
commit 85bbb5433669f72e6492592f4399ca5a15892fb9
Author: Weiqing Yang <[email protected]>
AuthorDate: Mon Aug 31 00:49:30 2026 -0700
[integrations][java][python] Apply Ollama native structured output (#981)
Generated-by: Claude Code 2.1.226
---
integrations/chat-models/ollama/pom.xml | 11 ++
.../ollama/OllamaChatModelConnection.java | 183 ++++++++++++++++++---
.../ollama/OllamaChatModelConnectionTest.java | 168 ++++++++++++++++++-
pom.xml | 16 ++
.../integrations/chat_models/ollama_chat_model.py | 89 +++++++++-
.../chat_models/tests/test_ollama_chat_model.py | 42 +++++
.../tests/test_ollama_native_structured_output.py | 151 +++++++++++++++++
7 files changed, 628 insertions(+), 32 deletions(-)
diff --git a/integrations/chat-models/ollama/pom.xml
b/integrations/chat-models/ollama/pom.xml
index 2ca08957..49638efd 100644
--- a/integrations/chat-models/ollama/pom.xml
+++ b/integrations/chat-models/ollama/pom.xml
@@ -43,6 +43,17 @@ under the License.
<artifactId>ollama4j</artifactId>
<version>${ollama4j.version}</version>
</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>
</dependencies>
</project>
\ No newline at end of file
diff --git
a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
index e86f1277..be42f834 100644
---
a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
+++
b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
@@ -20,6 +20,14 @@ package
org.apache.flink.agents.integrations.chatmodels.ollama;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
+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 io.github.ollama4j.exceptions.RoleNotFoundException;
import io.github.ollama4j.models.chat.*;
import io.github.ollama4j.models.request.OllamaChatEndpointCaller;
@@ -177,39 +185,62 @@ public class OllamaChatModelConnection extends
BaseChatModelConnection {
}
}
+ /**
+ * Whether Ollama can constrain generation to a schema for {@code
effectiveModel}.
+ *
+ * <p>Always {@code true}, and deliberately independent of the argument:
schema-constrained
+ * decoding is applied by the Ollama server's sampler rather than by the
model, so it holds for
+ * every model served by a server at or above v0.5.0. There is also no
model-level signal to key
+ * on. Ollama's model capability set — completion, tools, insert, vision,
embedding, thinking,
+ * image, audio — carries nothing schema-related, {@code /api/show}
reports exactly that set,
+ * and {@code /api/version} reports only a version string. Since a server
runs arbitrary local
+ * models, any allowlist would be invented, and would report not-capable
for models that do
+ * work.
+ *
+ * <p>Three deployments break the guarantee, none of them distinguishable
from a model name: a
+ * server below v0.5.0 rejects the {@code format} field with HTTP 400;
Ollama Cloud accepts the
+ * request but does not enforce the schema; and the MLX runner accepts the
field and drops it.
+ *
+ * <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) {
- try {
- // convert think to think mode.
- final Object think = modelParams.getOrDefault("think", true);
- ThinkMode thinkMode = ThinkMode.ENABLED;
- for (ThinkMode mode : ThinkMode.values()) {
- if (mode.getValue().equals(think)) {
- thinkMode = mode;
- break;
- }
- }
+ return doChat(messages, tools, modelParams, null);
+ }
+
+ /**
+ * Translates {@code outputSchema} into Ollama's native {@code 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 boolean extractReasoning =
(boolean) modelParams.getOrDefault("extract_reasoning",
true);
- final List<Tools.Tool> ollamaTools =
this.convertToOllamaTools(tools);
- final List<OllamaChatMessage> ollamaChatMessages =
- messages.stream()
- .map(this::convertToOllamaChatMessages)
- .collect(Collectors.toList());
-
- final String modelName = (String) modelParams.get("model");
final OllamaChatRequest chatRequest =
- OllamaChatRequest.builder()
- .withMessages(ollamaChatMessages)
- .withModel(modelName)
- .withThinking(thinkMode)
- .withUseTools(false)
- .build();
-
- chatRequest.setTools(ollamaTools);
+ buildRequest(messages, tools, modelParams, outputSchema);
final OllamaChatResult ollamaChatResult =
this.caller.callSync(chatRequest);
final OllamaChatResponseModel ollamaChatResponse =
ollamaChatResult.getResponseModel();
final OllamaChatMessage ollamaChatMessage =
ollamaChatResponse.getMessage();
@@ -229,6 +260,7 @@ public class OllamaChatModelConnection extends
BaseChatModelConnection {
}
// Stash token usage if model name is available
+ final String modelName = (String) modelParams.get("model");
if (modelName != null && !modelName.isBlank()) {
Integer promptTokens = ollamaChatResponse.getPromptEvalCount();
Integer completionTokens = ollamaChatResponse.getEvalCount();
@@ -245,6 +277,107 @@ public class OllamaChatModelConnection extends
BaseChatModelConnection {
}
}
+ // Package-private so the request body (including the native format) can
be asserted without
+ // issuing a live call through the Ollama endpoint caller.
+ OllamaChatRequest buildRequest(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
+ // convert think to think mode.
+ final Object think = modelParams.getOrDefault("think", true);
+ ThinkMode thinkMode = ThinkMode.ENABLED;
+ for (ThinkMode mode : ThinkMode.values()) {
+ if (mode.getValue().equals(think)) {
+ thinkMode = mode;
+ break;
+ }
+ }
+
+ final List<Tools.Tool> ollamaTools = this.convertToOllamaTools(tools);
+ final List<OllamaChatMessage> ollamaChatMessages =
+ messages.stream()
+ .map(this::convertToOllamaChatMessages)
+ .collect(Collectors.toList());
+
+ final String modelName = (String) modelParams.get("model");
+ final OllamaChatRequest chatRequest =
+ OllamaChatRequest.builder()
+ .withMessages(ollamaChatMessages)
+ .withModel(modelName)
+ .withThinking(thinkMode)
+ .withUseTools(false)
+ .build();
+
+ chatRequest.setTools(ollamaTools);
+
+ // 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 schema is a request field of its own rather than a sampling
option, so it is set as
+ // the request's format, which is left unset when no native
translation applies and is then
+ // omitted from the serialized body rather than serialized as null.
+ //
+ // TODO(#912): the requested strategy is not visible here, so this
re-check cannot tell an
+ // explicit NATIVE request apart from one that merely resolved to
native. A caller asking
+ // for NATIVE on a schema form this branch skips therefore gets an
unconstrained response
+ // instead of an error. Once strategy resolution is wired up, NATIVE
must either bypass
+ // this capability re-check or fail explicitly.
+ if (outputSchema instanceof Class &&
supportsNativeStructuredOutput(modelName)) {
+ chatRequest.setFormat(toNativeFormat((Class<?>) outputSchema));
+ }
+
+ return chatRequest;
+ }
+
+ // Derives the JSON schema Ollama's format field expects from a POJO
class. Every setting below
+ // addresses a concrete way the generated schema otherwise fails to
constrain generation:
+ //
+ // - 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.
"getSummary()".
+ // - MAP_VALUES_AS_ADDITIONAL_PROPERTIES gives a Map its value schema.
Without it the map
+ // admits any value, and a model does emit values that the declared
value type then fails
+ // to deserialize.
+ // - Sorting fields before methods and applying no further comparison
leaves properties in
+ // declaration order. Ollama's grammar fixes generation order to the
order the schema
+ // declares its properties, so the default alphabetical order would
condition generation on
+ // an order the class does not read in.
+ // - 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. The
+ // response is read back into the same class with an ObjectMapper, so
a property that
+ // @JsonProperty renames or @JsonIgnore drops has to be stated in the
schema under the name
+ // the mapper reads, or a response that satisfies the schema still
fails to deserialize.
+ // 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 gains nothing: Ollama's
grammar already
+ // refuses a key the schema does not declare, even one a prompt
explicitly asks for, and
+ // only an explicit additionalProperties: true admits one.
+ // - DEFINITION_FOR_MAIN_SCHEMA lets a recursive type generate a schema,
but when the
+ // document root is a $ref and one $defs entry references another, the
server drops the
+ // grammar and returns a free-form object. Any nested type used twice
is extracted into
+ // $defs, so enabling it would silently unconstrain a common shape to
rescue a rare one. A
+ // recursive type instead fails loudly, with HTTP 400 from the server.
+ private static ObjectNode toNativeFormat(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()));
+ return new
SchemaGenerator(configBuilder.build()).generateSchema(schemaClass);
+ }
+
/**
* Converts Ollama tool calls to the format expected by the Flink Agents
framework.
*
diff --git
a/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
index 06570c29..f8b70633 100644
---
a/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
+++
b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
@@ -17,7 +17,13 @@
*/
package org.apache.flink.agents.integrations.chatmodels.ollama;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.JsonNode;
+import io.github.ollama4j.models.chat.OllamaChatRequest;
import io.github.ollama4j.tools.Tools;
+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;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.tools.Tool;
@@ -27,18 +33,62 @@ import org.apache.flink.agents.api.tools.ToolResponse;
import org.apache.flink.agents.api.tools.ToolType;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.NullAndEmptySource;
+import org.junit.jupiter.params.provider.ValueSource;
+import java.util.ArrayList;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
+import java.util.Optional;
import static org.assertj.core.api.Assertions.assertThat;
/**
- * Unit tests for {@link OllamaChatModelConnection}'s tool-schema conversion —
no network access.
+ * Unit tests for {@link OllamaChatModelConnection}'s tool-schema conversion
and native
+ * structured-output behavior — no network access. The structured-output
assertions inspect the body
+ * built by {@code buildRequest}, and exercise the capability predicate
directly.
*/
class OllamaChatModelConnectionTest {
private static final ResourceContext NOOP =
ResourceContext.fromGetResource((a, b) -> null);
+ private static final String DRAFT_2020_12 =
"https://json-schema.org/draft/2020-12/schema";;
+
+ /**
+ * 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 OllamaChatModelConnection connection() {
ResourceDescriptor desc =
ResourceDescriptor.Builder.newBuilder(OllamaChatModelConnection.class.getName())
@@ -64,6 +114,16 @@ class OllamaChatModelConnectionTest {
}
}
+ private static Map<String, Object> params(String model) {
+ Map<String, Object> params = new HashMap<>();
+ params.put("model", model);
+ return params;
+ }
+
+ private static List<ChatMessage> userMessage() {
+ return List.of(new ChatMessage(MessageRole.USER, "hi"));
+ }
+
@Test
@DisplayName("A schema without a 'required' key converts with every
property optional")
void testSchemaWithoutRequiredKey() {
@@ -101,4 +161,110 @@ class OllamaChatModelConnectionTest {
assertThat(tool.getToolSpec().getParameters().getProperties().get("b").isRequired())
.isFalse();
}
+
+ @Test
+ @DisplayName("A POJO output schema is sent as the native format")
+ void buildRequestSetsFormatForPojoSchema() {
+ OllamaChatRequest request =
+ connection()
+ .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Report.class);
+
+ assertThat(request.getFormat()).isInstanceOf(JsonNode.class);
+ JsonNode schema = (JsonNode) request.getFormat();
+ assertThat(schema.path("type").asText()).isEqualTo("object");
+ assertThat(schema.path("properties").has("summary")).isTrue();
+ }
+
+ @Test
+ @DisplayName("No output schema leaves the request without a format")
+ void buildRequestOmitsFormatWithoutSchema() {
+ OllamaChatRequest request =
+ connection().buildRequest(userMessage(), List.of(),
params("qwen3:4b"), null);
+
+ assertThat(request.getFormat()).isNull();
+ }
+
+ @Test
+ @DisplayName("A RowTypeInfo-shaped schema stays on the prompt fallback")
+ void buildRequestLeavesFormatUnsetForRowTypeInfo() {
+ // A RowTypeInfo schema arrives wrapped in OutputSchema rather than as
a bare POJO Class, so
+ // it must not activate native structured output. OutputSchema cannot
be instantiated here
+ // because RowTypeInfo is not on this module's classpath; any
non-Class schema object
+ // exercises the same gate.
+ Object nonClassSchema = "row<name STRING>";
+
+ OllamaChatRequest request =
+ connection()
+ .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), nonClassSchema);
+
+ assertThat(request.getFormat()).isNull();
+ }
+
+ @Test
+ @DisplayName("The generated schema constrains draft, property order, map
values and required")
+ void generatedSchemaShapeIsConstraining() {
+ OllamaChatRequest request =
+ connection()
+ .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Report.class);
+ JsonNode schema = (JsonNode) request.getFormat();
+
+ // Ollama fixes generation order to the order the schema declares its
properties, so the
+ // emitted order has to follow the class rather than the alphabet. A
getter backed by no
+ // field must not surface as a property of its own.
+ assertThat(schema.path("$schema").asText()).isEqualTo(DRAFT_2020_12);
+ assertThat(schema.path("properties").fieldNames())
+ .toIterable()
+ .containsExactly("summary", "counts", "note", "total");
+
+ // A map without a value schema admits any value, which the model does
take up and which
+ // then fails to deserialize into the declared type.
+
assertThat(schema.path("properties").path("counts").path("additionalProperties").isObject())
+ .isTrue();
+ assertThat(
+ schema.path("properties")
+ .path("counts")
+ .path("additionalProperties")
+ .path("type")
+ .asText())
+ .isEqualTo("integer");
+
+ // Every field is required except the one the caller declared
omissible.
+ assertThat(textValues(schema.path("required")))
+ .containsExactlyInAnyOrder("summary", "counts", "total");
+ }
+
+ @Test
+ @DisplayName("The generated schema names properties the way Jackson
deserializes them")
+ void generatedSchemaFollowsJacksonPropertyNames() {
+ OllamaChatRequest request =
+ connection()
+ .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Profile.class);
+ JsonNode schema = (JsonNode) request.getFormat();
+
+ // The response is read back with an ObjectMapper, which accepts the
renamed property and
+ // rejects the Java field name, and which discards an ignored property
the schema would
+ // otherwise force the model to fabricate.
+ assertThat(schema.path("properties").fieldNames())
+ .toIterable()
+ .containsExactly("full_name", "age");
+ assertThat(textValues(schema.path("required")))
+ .containsExactlyInAnyOrder("full_name", "age");
+ }
+
+ @ParameterizedTest
+ @NullAndEmptySource
+ @ValueSource(strings = {"qwen3:4b", "llama3.2", "gpt-oss:20b",
"some-private-local-model"})
+ @DisplayName("Capability is reported for any model, since the server
provides it")
+ void supportsNativeStructuredOutputIsServerNotModelGated(String model) {
+ // Null and empty are included because the capability does not depend
on the argument at
+ // all, so the guard the sibling connections need for their allowlists
would be a silent
+ // behavior change here.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isTrue();
+ }
+
+ private static List<String> textValues(JsonNode arrayNode) {
+ List<String> values = new ArrayList<>();
+ arrayNode.forEach(element -> values.add(element.asText()));
+ return values;
+ }
}
diff --git a/pom.xml b/pom.xml
index fb78d863..2e6ce6fc 100644
--- a/pom.xml
+++ b/pom.xml
@@ -54,6 +54,7 @@ under the License.
<bytebuddy.version>1.15.4</bytebuddy.version>
<cel.version>0.12.0</cel.version>
<protobuf.version>4.33.5</protobuf.version>
+ <victools.jsonschema.version>4.38.0</victools.jsonschema.version>
<gpg.useagent>true</gpg.useagent>
<arguments />
</properties>
@@ -85,6 +86,21 @@ under the License.
<artifactId>protobuf-java</artifactId>
<version>${protobuf.version}</version>
</dependency>
+ <!--
+ victools JSON Schema BOM. The generator also reaches the
classpath transitively
+ through provider SDKs, together with its sibling modules
jsonschema-module-jackson
+ and jsonschema-module-swagger-2. Managing the family through the
BOM keeps all three
+ on one version; pinning only the directly declared artifact
would let the siblings
+ move independently with an SDK upgrade, and the enforcer
configures no convergence
+ rule that would report the split.
+ -->
+ <dependency>
+ <groupId>com.github.victools</groupId>
+ <artifactId>jsonschema-generator-bom</artifactId>
+ <version>${victools.jsonschema.version}</version>
+ <type>pom</type>
+ <scope>import</scope>
+ </dependency>
</dependencies>
</dependencyManagement>
diff --git a/python/flink_agents/integrations/chat_models/ollama_chat_model.py
b/python/flink_agents/integrations/chat_models/ollama_chat_model.py
index bd16905a..a17f5047 100644
--- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py
@@ -19,7 +19,8 @@ import uuid
from typing import Any, Dict, List, Literal, Sequence
from ollama import Client, Message
-from pydantic import Field
+from pydantic import BaseModel, Field
+from typing_extensions import override
from flink_agents.api.agents.types import OutputSchema
from flink_agents.api.chat_message import ChatMessage, MessageRole
@@ -34,6 +35,23 @@ DEFAULT_CONTEXT_WINDOW = 2048
DEFAULT_REQUEST_TIMEOUT = 30.0
+def _native_format(output_schema: Any) -> Dict[str, Any] | None:
+ """Build the Ollama ``format`` payload for a native structured-output
request.
+
+ Returns ``None`` (leaving the request unconstrained) unless the schema is a
+ ``BaseModel`` subclass. A ``RowTypeInfo`` schema is skipped so it keeps the
+ prompt-engineering fallback.
+ """
+ if output_schema is None:
+ return None
+ 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.model_json_schema()
+
+
class OllamaChatModelConnection(BaseChatModelConnection):
"""Ollama ChatModelServer which manage the connection to the Ollama server.
@@ -82,6 +100,30 @@ class OllamaChatModelConnection(BaseChatModelConnection):
self.__client = Client(host=self.base_url,
timeout=self.request_timeout)
return self.__client
+ @override
+ def supports_native_structured_output(self, effective_model: str | None)
-> bool:
+ """Whether Ollama can constrain generation to a schema for
``effective_model``.
+
+ Always ``True``, and deliberately independent of the argument:
+ schema-constrained decoding is applied by the Ollama server's sampler
rather
+ than by the model, so it holds for every model served by a server at
or above
+ v0.5.0. There is also no model-level signal to key on. Ollama's model
capability
+ set -- completion, tools, insert, vision, embedding, thinking, image,
audio --
+ carries nothing schema-related, ``/api/show`` reports exactly that
set, and
+ ``/api/version`` reports only a version string. Since a server runs
arbitrary
+ local models, any allowlist would be invented, and would report
not-capable for
+ models that do work.
+
+ Three deployments break the guarantee, none of them distinguishable
from a model
+ name: a server below v0.5.0 rejects the ``format`` field with HTTP
400; Ollama
+ Cloud accepts the request but does not enforce the schema; and the MLX
runner
+ accepts the field and drops it.
+
+ Reads no instance state, so capability stays answerable independently
of how the
+ connection was configured.
+ """
+ return True
+
def chat(
self,
messages: Sequence[ChatMessage],
@@ -91,12 +133,26 @@ class OllamaChatModelConnection(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.
+ Parameters
+ ----------
+ messages : Sequence[ChatMessage]
+ Input message sequence.
+ tools : Optional[List[Tool]]
+ List of tools that can be called by the model.
+ output_schema : OutputSchema | None
+ The schema the response should conform to, or ``None`` for an
unconstrained
+ response. A ``BaseModel`` schema is sent as Ollama's native
``format``
+ argument so the server constrains decoding to it; any other schema
form,
+ notably a ``RowTypeInfo``, keeps the prompt-engineering fallback.
+ **kwargs : Any
+ Additional parameters passed to the model service (e.g.,
temperature,
+ num_ctx, etc.)
+
+ Returns:
+ -------
+ ChatMessage
+ Model response message
"""
- self._reject_unsupported_output_schema(output_schema)
ollama_messages = self.__convert_to_ollama_messages(messages)
# Convert tool format
@@ -105,6 +161,26 @@ class OllamaChatModelConnection(BaseChatModelConnection):
ollama_tools = [to_openai_tool(metadata=tool.metadata) for tool in
tools]
model_name = kwargs.pop("model")
+
+ # 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. The schema is a request field of its
own rather
+ # than a sampling option, so it is passed as the format argument,
which is
+ # omitted altogether when no native translation applies.
+ #
+ # TODO(#912): the requested strategy is not visible here, so this
re-check
+ # cannot tell an explicit NATIVE request apart from one that merely
resolved to
+ # native. A caller asking for NATIVE on a schema form this branch skips
+ # therefore gets an unconstrained response instead of an error. Once
strategy
+ # resolution is wired up, NATIVE must either bypass this capability
re-check or
+ # fail explicitly.
+ native_format = None
+ if output_schema is not None and
self.supports_native_structured_output(
+ model_name
+ ):
+ native_format = _native_format(output_schema)
+ format_kwargs = {} if native_format is None else {"format":
native_format}
+
response = self.client.chat(
model=model_name,
messages=ollama_messages,
@@ -113,6 +189,7 @@ class OllamaChatModelConnection(BaseChatModelConnection):
options=kwargs,
keep_alive=kwargs.get("keep_alive", False),
think=kwargs.get("think", True),
+ **format_kwargs,
)
ollama_tool_calls = response.message.tool_calls
diff --git
a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
index cceedcd2..e868dd39 100644
---
a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
@@ -15,11 +15,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+import json
import os
from unittest.mock import MagicMock
import pytest
+from pydantic import BaseModel
+from flink_agents.api.agents.types import OutputSchema
from flink_agents.api.chat_message import ChatMessage, MessageRole
from flink_agents.api.resource import Resource, ResourceType
from flink_agents.api.resource_context import ResourceContext
@@ -111,6 +114,45 @@ def test_ollama_chat_with_tools() -> None:
assert add(**tool_call["function"]["arguments"]) == 3
+class Person(BaseModel):
+ """A flat output schema with one string field and one integer field."""
+
+ name: str
+ age: int
+
+
[email protected](
+ client is None, reason="Ollama client is not available or test model is
missing"
+)
+def test_ollama_chat_with_output_schema() -> None:
+ """Verify a BaseModel schema constrains the server response to that schema.
+
+ Asserts schema conformance only. Ollama constrains decoding to the schema
on
+ the server side, so the shape of the response is guaranteed while its
values
+ are not; asserting the values would depend on the model's accuracy.
+ """
+ server = OllamaChatModelConnection(request_timeout=120.0)
+
+ response = server.chat(
+ [
+ ChatMessage(
+ role=MessageRole.USER,
+ content="Ada Lovelace is 36 years old. Extract the person.",
+ )
+ ],
+ model=test_model,
+ output_schema=OutputSchema(output_schema=Person),
+ think=False,
+ )
+
+ payload = json.loads(response.content)
+ assert {"name", "age"}.issubset(payload)
+
+ person = Person.model_validate_json(response.content)
+ assert isinstance(person.name, str)
+ assert isinstance(person.age, int)
+
+
def test_model_field_roundtrip() -> None:
"""Verify `model` is preserved through pydantic dump/validate
round-trip."""
setup = OllamaChatModelSetup(connection="conn", model="test-model")
diff --git
a/python/flink_agents/integrations/chat_models/tests/test_ollama_native_structured_output.py
b/python/flink_agents/integrations/chat_models/tests/test_ollama_native_structured_output.py
new file mode 100644
index 00000000..0d1d2934
--- /dev/null
+++
b/python/flink_agents/integrations/chat_models/tests/test_ollama_native_structured_output.py
@@ -0,0 +1,151 @@
+################################################################################
+# 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.
+#################################################################################
+from typing import Any, Dict
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic import BaseModel
+from pyflink.common.typeinfo import Types
+
+from flink_agents.api.agents.types import OutputSchema
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.integrations.chat_models.ollama_chat_model import (
+ OllamaChatModelConnection,
+)
+
+
+class Person(BaseModel):
+ """A representative flat BaseModel output schema."""
+
+ name: str
+ age: int
+
+
+class Company(BaseModel):
+ """A nested output schema, whose JSON schema carries a ``$defs`` entry."""
+
+ name: str
+ owner: Person
+
+
+def _connection() -> OllamaChatModelConnection:
+ """A connection whose Ollama client is a mock, so no server is
contacted."""
+ conn = OllamaChatModelConnection()
+ response = MagicMock()
+ response.message.role = "assistant"
+ response.message.content = "ok"
+ response.message.tool_calls = None
+ response.prompt_eval_count = 1
+ response.eval_count = 2
+ mock_client = MagicMock()
+ mock_client.chat.return_value = response
+ conn._OllamaChatModelConnection__client = mock_client
+ return conn
+
+
+def _chat_call_kwargs(conn: OllamaChatModelConnection) -> Dict[str, Any]:
+ return conn.client.chat.call_args.kwargs
+
+
+def _messages() -> list[ChatMessage]:
+ return [ChatMessage(role=MessageRole.USER, content="hi")]
+
+
+def test_native_applied_for_base_model() -> None:
+ """A BaseModel schema reaches the request as the native ``format``
argument."""
+ conn = _connection()
+ conn.chat(
+ _messages(), model="qwen3",
output_schema=OutputSchema(output_schema=Person)
+ )
+ native_format = _chat_call_kwargs(conn)["format"]
+ assert native_format["properties"].keys() == {"name", "age"}
+
+
+def test_format_absent_without_schema() -> None:
+ """A call without a schema carries no ``format``, leaving generation
unconstrained."""
+ conn = _connection()
+ conn.chat(_messages(), model="qwen3")
+ assert "format" not in _chat_call_kwargs(conn)
+
+
+def test_native_not_applied_for_row_type_info() -> None:
+ """A RowTypeInfo schema has no native translation and keeps the prompt
fallback."""
+ conn = _connection()
+ row_type = Types.ROW_NAMED(["name"], [Types.STRING()])
+ conn.chat(
+ _messages(), model="qwen3",
output_schema=OutputSchema(output_schema=row_type)
+ )
+ assert "format" not in _chat_call_kwargs(conn)
+
+
+def test_schema_is_model_json_schema() -> None:
+ """The payload is pydantic's schema verbatim, ``$defs`` and all.
+
+ A nested schema is used because it is the shape a hand-rolled translation
+ diverges on: inlining or renaming a ``$defs`` entry breaks the ``$ref``
targets
+ the server resolves when it builds the grammar.
+ """
+ conn = _connection()
+ conn.chat(
+ _messages(), model="qwen3",
output_schema=OutputSchema(output_schema=Company)
+ )
+ assert _chat_call_kwargs(conn)["format"] == Company.model_json_schema()
+
+
+def test_schema_not_passed_as_sampling_option() -> None:
+ """The schema never reaches ``options``, which the server reads as
sampling options.
+
+ A schema written into the forwarded kwargs would arrive there instead of in
+ ``format``, so the server would apply no grammar and report no error.
+ """
+ conn = _connection()
+ conn.chat(
+ _messages(), model="qwen3",
output_schema=OutputSchema(output_schema=Person)
+ )
+ options = _chat_call_kwargs(conn)["options"]
+ assert "format" not in options
+ assert Person.model_json_schema() not in options.values()
+
+
[email protected](
+ "model",
+ ["qwen3", "llama3.2", "gemma3:270m", "mistral", "an-unknown-model", "",
None],
+)
+def test_supports_native_structured_output(model: str | None) -> None:
+ """Capability is reported for every model, including an absent model name.
+
+ The capability is the server's, not the model's, so there is no model name
it
+ can be keyed on and none it should report not-capable for.
+ """
+ conn = OllamaChatModelConnection()
+ assert conn.supports_native_structured_output(model) is True
+
+
+def test_schema_accepted_not_rejected() -> None:
+ """A schema with no native translation is answered, not refused.
+
+ Rejecting is what a connection without native structured output does. This
one
+ has it, so a schema form it cannot translate natively falls back to the
prompt
+ engineering the caller already applied rather than raising.
+ """
+ conn = _connection()
+ row_type = Types.ROW_NAMED(["name"], [Types.STRING()])
+ response = conn.chat(
+ _messages(), model="qwen3",
output_schema=OutputSchema(output_schema=row_type)
+ )
+ assert response.content == "ok"
