From 679078a2b5f568658c703dd53c225c7759099930 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 18 Jul 2026 22:00:25 -0700 Subject: [PATCH 1/2] [api][python] Add explicit output schema parameter to the chat path Carry the output schema to chat model connections as an explicit parameter and split the structured output strategy into a user policy and a per-model capability, so a later change can apply a provider's native structured output without threading the schema through the provider-facing model parameters. Java adds a concrete four argument chat overload whose default body ignores the schema and delegates to the existing abstract three argument chat, so connections with no native mechanism need no change. Python adds an explicit output_schema parameter to every connection, including the Java bridge connection, so the schema binds to a named parameter and cannot be forwarded into a provider request through kwargs. StructuredOutputStrategy expresses user intent (AUTO, NATIVE or PROMPT) on the setup and defaults to AUTO. Capability is separate: a connection side predicate over the effective model, since a provider supports its native mechanism only on some models. The base predicate returns false; a connection that adds a native translation overrides it. This is the mechanism only. No connection applies native structured output and no path passes a schema on the chat call yet, so behavior is unchanged. --- .../chat/model/BaseChatModelConnection.java | 51 ++++++++ .../api/chat/model/BaseChatModelSetup.java | 16 +++ .../chat/model/StructuredOutputStrategy.java | 112 +++++++++++++++++ .../api/chat/model/BaseChatModelTest.java | 84 +++++++++++++ .../api/chat_models/chat_model.py | 114 +++++++++++++++++- .../chat_models/tests/test_chat_model_base.py | 84 ++++++++++++- .../mock_chat_model_agent.py | 2 + .../tool_parameter_injection_agent.py | 2 + .../anthropic/anthropic_chat_model.py | 10 +- .../azure/azure_openai_chat_model.py | 7 ++ .../chat_models/ollama_chat_model.py | 10 +- .../chat_models/openai/openai_chat_model.py | 7 ++ .../test_output_schema_param_declared.py | 87 +++++++++++++ .../chat_models/tongyi_chat_model.py | 10 +- .../runtime/java/java_chat_model.py | 6 + 15 files changed, 594 insertions(+), 8 deletions(-) create mode 100644 api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java create mode 100644 python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java index 5181073a7..9cbad62bf 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.tools.Tool; +import javax.annotation.Nullable; + import java.util.List; import java.util.Map; @@ -45,6 +47,26 @@ public ResourceType getResourceType() { return ResourceType.CHAT_MODEL_CONNECTION; } + /** + * Whether this connection can apply the provider's native structured-output API for the given + * model. + * + *

Capability is model-dependent, not connection-wide: a single provider connection + * commonly serves both models that accept a native schema parameter and models that do not. It + * is therefore evaluated against the effective model at request-build time — the model + * actually being called, which per-request parameters may override. + * + *

The default {@code false} keeps a connection on the prompt-engineering fallback. An + * unrecognized model must report {@code false} so that it degrades to the fallback rather than + * failing at the provider. + * + * @param effectiveModel the model the request will be issued against, may be null + * @return true if a schema can be applied natively for {@code effectiveModel} + */ + protected boolean supportsNativeStructuredOutput(String effectiveModel) { + return false; + } + /** * Process a chat request and return a chat response. * @@ -55,4 +77,33 @@ public ResourceType getResourceType() { */ public abstract ChatMessage chat( List messages, List tools, Map modelParams); + + /** + * Process a chat request that carries an output schema, and return a chat response. + * + *

{@code outputSchema} is framework-level execution metadata, kept off {@code modelParams} + * so that it can never reach a provider SDK request as a generation parameter. It is either a + * POJO {@link Class} or an {@link org.apache.flink.agents.api.agents.OutputSchema} (a {@code + * RowTypeInfo} wrapper); the two cases are distinguished by the connection that consumes it. + * + *

This default implementation ignores {@code outputSchema} and delegates to {@link + * #chat(List, List, Map)}. Connections without a native structured-output translation inherit + * it unchanged and need no edits. A connection that does translate a schema into a native + * provider parameter overrides this overload, and reports its capability via {@link + * #supportsNativeStructuredOutput(String)}. + * + * @param messages the input chat messages + * @param tools the tools can be called by the model + * @param modelParams the additional arguments passed to the model + * @param outputSchema the schema the response should conform to, or null for an unconstrained + * response + * @return the chat response containing model outputs + */ + public ChatMessage chat( + List messages, + List tools, + Map modelParams, + @Nullable Object outputSchema) { + return chat(messages, tools, modelParams); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java index 34b2b2994..6c59f6ef7 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java @@ -48,6 +48,7 @@ public abstract class BaseChatModelSetup extends Resource { @Nullable protected String skillDiscoveryPrompt; protected List allowedCommands; protected List allowedScriptDirs; + protected StructuredOutputStrategy structuredOutputStrategy; @Nullable protected BaseChatModelConnection connection; protected final List tools = new ArrayList<>(); @@ -67,6 +68,10 @@ public BaseChatModelSetup(ResourceDescriptor descriptor, ResourceContext resourc declaredScriptDirs == null ? new ArrayList<>() : new ArrayList<>(declaredScriptDirs); + this.structuredOutputStrategy = + StructuredOutputStrategy.fromArgument( + descriptor.getArgument("structured_output_strategy"), + StructuredOutputStrategy.AUTO); } /** @@ -219,4 +224,15 @@ public List getAllowedCommands() { public List getAllowedScriptDirs() { return allowedScriptDirs; } + + /** + * The configured intent about how an output schema should be applied, defaulting to {@link + * StructuredOutputStrategy#AUTO}. Whether native structured output is actually applied combines + * this policy with the connection's model-dependent capability. + * + * @return the structured output strategy + */ + public StructuredOutputStrategy getStructuredOutputStrategy() { + return structuredOutputStrategy; + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java new file mode 100644 index 000000000..4026d1992 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java @@ -0,0 +1,112 @@ +/* + * 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. + */ + +package org.apache.flink.agents.api.chat.model; + +import java.util.Locale; + +/** + * User intent about how an output schema should be applied to a chat request. + * + *

This expresses policy only. Whether a connection can apply the provider's native + * structured-output API is a separate, model-dependent capability question answered by + * {@link BaseChatModelConnection#supportsNativeStructuredOutput(String)}. Policy and capability are + * combined at request-build time. + */ +public enum StructuredOutputStrategy { + /** + * Use the provider's native structured-output API when the effective model is capable of it, + * and fall back to prompt engineering otherwise. This is the default. + */ + AUTO, + + /** + * Always use the provider's native structured-output API, without consulting the capability + * predicate. + */ + NATIVE, + + /** + * Never use the provider's native structured-output API; rely on prompt engineering alone. This + * matches the behavior of connections that have no native translation. + */ + PROMPT; + + /** + * Resolves this policy against a connection's model-dependent capability into whether the + * provider's native structured-output API should be used. + * + *

+ * + * @param modelCapable whether the connection reports the effective model as natively capable + * @return true if native structured output should be applied + */ + public boolean resolvesToNative(boolean modelCapable) { + switch (this) { + case NATIVE: + return true; + case PROMPT: + return false; + case AUTO: + default: + return modelCapable; + } + } + + /** + * Resolves a strategy from a descriptor argument, which may arrive either as a {@code + * StructuredOutputStrategy} or — across the Python bridge, where arguments are carried as JSON + * — as its case-insensitive name. + * + * @param value the raw descriptor argument, may be null + * @param defaultValue the strategy to use when {@code value} is null + * @return the resolved strategy + * @throws IllegalArgumentException if {@code value} is neither null, a {@code + * StructuredOutputStrategy}, nor the name of one + */ + public static StructuredOutputStrategy fromArgument( + Object value, StructuredOutputStrategy defaultValue) { + if (value == null) { + return defaultValue; + } + if (value instanceof StructuredOutputStrategy) { + return (StructuredOutputStrategy) value; + } + if (value instanceof String) { + try { + return valueOf(((String) value).toUpperCase(Locale.ROOT)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + String.format( + "Unknown structured output strategy '%s'. Expected one of: AUTO, NATIVE, PROMPT.", + value), + e); + } + } + throw new IllegalArgumentException( + String.format( + "Unsupported structured output strategy type '%s'. Expected a StructuredOutputStrategy or its name.", + value.getClass().getName())); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java index 43f8c8b05..117a70174 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java @@ -237,6 +237,7 @@ void testChatResponseFormat() { /** Connection that captures the messages passed to it for assertions. */ private static class RecordingConnection extends BaseChatModelConnection { List capturedMessages; + Map capturedModelParams; RecordingConnection() { super( @@ -249,6 +250,7 @@ private static class RecordingConnection extends BaseChatModelConnection { public ChatMessage chat( List messages, List tools, Map modelParams) { this.capturedMessages = new ArrayList<>(messages); + this.capturedModelParams = new HashMap<>(modelParams); return new ChatMessage(MessageRole.ASSISTANT, "ok"); } } @@ -320,6 +322,88 @@ void testChatRefillsTemplateOnSubsequentInvocations() { assertEquals("tool result", connection.capturedMessages.get(1).getContent()); } + @Test + @DisplayName("Default chat() overload ignores outputSchema and delegates to the 3-arg chat()") + void testDefaultChatOverloadIgnoresOutputSchema() { + RecordingConnection connection = new RecordingConnection(); + Map modelParams = new HashMap<>(); + modelParams.put("temperature", 0.5); + + ChatMessage response = + connection.chat( + List.of(new ChatMessage(MessageRole.USER, "hi")), + List.of(), + modelParams, + new Object()); + + // The 3-arg chat() ran (it is what produces "ok") and the schema never reached + // modelParams, so it cannot travel on to a provider SDK request. + assertEquals("ok", response.getContent()); + assertEquals(Map.of("temperature", 0.5), connection.capturedModelParams); + } + + @Test + @DisplayName("Default capability predicate reports no native structured output for any model") + void testDefaultCapabilityPredicateIsFalse() { + RecordingConnection connection = new RecordingConnection(); + + assertFalse(connection.supportsNativeStructuredOutput("gpt-4o")); + assertFalse(connection.supportsNativeStructuredOutput("gpt-3.5-turbo")); + assertFalse(connection.supportsNativeStructuredOutput(null)); + } + + @Test + @DisplayName("Structured-output strategy defaults to AUTO when the descriptor omits it") + void testStructuredOutputStrategyDefaultsToAuto() { + RecordingChatModelSetup setup = + new RecordingChatModelSetup(new RecordingConnection(), null); + + assertEquals(StructuredOutputStrategy.AUTO, setup.getStructuredOutputStrategy()); + } + + @Test + @DisplayName("Structured-output strategy is read from the descriptor argument") + void testStructuredOutputStrategyReadFromDescriptor() { + TestChatModel model = + new TestChatModel( + new ResourceDescriptor( + TestChatModel.class.getName(), + Map.of("structured_output_strategy", "native")), + null); + + assertEquals(StructuredOutputStrategy.NATIVE, model.getStructuredOutputStrategy()); + } + + @Test + @DisplayName("An unrecognized structured-output strategy is rejected instead of defaulting") + void testUnknownStructuredOutputStrategyRejected() { + ResourceDescriptor descriptor = + new ResourceDescriptor( + TestChatModel.class.getName(), + Map.of("structured_output_strategy", "bogus")); + + assertThrows(IllegalArgumentException.class, () -> new TestChatModel(descriptor, null)); + } + + @Test + @DisplayName("AUTO resolves to native only when the effective model is capable") + void testAutoStrategyResolvesToNativeOnlyWhenCapable() { + assertTrue(StructuredOutputStrategy.AUTO.resolvesToNative(true)); + assertFalse(StructuredOutputStrategy.AUTO.resolvesToNative(false)); + } + + @Test + @DisplayName("NATIVE forces native even when the model is not capable") + void testNativeStrategyForcesNativeRegardlessOfCapability() { + assertTrue(StructuredOutputStrategy.NATIVE.resolvesToNative(false)); + } + + @Test + @DisplayName("PROMPT never resolves to native even when the model is capable") + void testPromptStrategyNeverResolvesToNative() { + assertFalse(StructuredOutputStrategy.PROMPT.resolvesToNative(true)); + } + @Test @DisplayName("Test chat with long input") void testChatWithLongInput() { diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index ac4a814aa..a3443cf3d 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -17,11 +17,13 @@ ################################################################################# import re from abc import ABC, abstractmethod +from enum import Enum from typing import Any, ClassVar, Dict, List, Mapping, Sequence, Tuple, cast from pydantic import Field, PrivateAttr from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ( ChatMessage, MessageRole, @@ -33,6 +35,70 @@ from flink_agents.api.tools.tool import Tool +class StructuredOutputStrategy(str, Enum): + """User intent about how an output schema should be applied to a chat request. + + This expresses *policy* only. Whether a connection *can* apply the provider's + native structured-output API is a separate, model-dependent *capability* + question. Policy and capability are combined at request-build time. + + Inherits from ``str`` so the value survives the JSON-carried bridge to Java. + Java serializes this enum as its *name* ("NATIVE") while the value here is + lowercase, so ``_missing_`` accepts either form in any case — matching the + case-insensitive resolver on the Java side. + + Attributes: + ---------- + AUTO : str + Use the provider's native structured-output API when the effective model is + capable of it, and fall back to prompt engineering otherwise. The default. + NATIVE : str + Always use the provider's native structured-output API, without consulting + the capability predicate. + PROMPT : str + Never use the provider's native structured-output API; rely on prompt + engineering alone. Matches the behavior of connections that have no native + translation. + """ + + AUTO = "auto" + NATIVE = "native" + PROMPT = "prompt" + + def resolves_to_native(self, model_capable: bool) -> bool: # noqa: FBT001 + """Resolve this policy against a model's capability into whether to go native. + + ``AUTO`` defers to ``model_capable`` (native when the effective model can, else + the prompt-engineering fallback); ``NATIVE`` always resolves to native, ignoring + ``model_capable``, so an explicit user intent surfaces a provider error rather + than silently degrading; ``PROMPT`` never resolves to native. + + Parameters + ---------- + model_capable : bool + Whether the connection reports the effective model as natively capable. + + Returns: + ------- + bool + ``True`` if native structured output should be applied. + """ + if self is StructuredOutputStrategy.NATIVE: + return True + if self is StructuredOutputStrategy.PROMPT: + return False + return model_capable + + @classmethod + def _missing_(cls, value: object) -> "StructuredOutputStrategy | None": + if isinstance(value, str): + normalized = value.lower() + for member in cls: + if normalized in (member.value, member.name.lower()): + return member + return None + + class BaseChatModelConnection(Resource, ABC): """Base abstract class for chat model connection. @@ -54,6 +120,32 @@ def resource_type(cls) -> ResourceType: """Return resource type of class.""" return ResourceType.CHAT_MODEL_CONNECTION + def supports_native_structured_output(self, effective_model: str | None) -> bool: + """Whether this connection can natively structure output for a given model. + + Capability is *model-dependent*, not connection-wide: a single provider + connection commonly serves both models that accept a native schema parameter and + models that do not, so it is evaluated against the *effective* model at + request-build time — the model actually being called, which per-request + parameters may override. + + The default ``False`` keeps a connection on the prompt-engineering fallback. A + connection that translates a schema into a native provider parameter overrides + this; an unrecognized model must report ``False`` so it degrades to the fallback + rather than failing at the provider. + + Parameters + ---------- + effective_model : str | None + The model the request will be issued against, may be ``None``. + + Returns: + ------- + bool + ``True`` if a schema can be applied natively for ``effective_model``. + """ + return False + DEFAULT_REASONING_PATTERNS: ClassVar[Tuple[re.Pattern[str], ...]] = ( re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), @@ -106,6 +198,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -116,6 +209,14 @@ def chat( Input message sequence tools : Optional[List] 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. This is framework-level execution metadata, and + every implementation must declare it as a named parameter rather than let + it fall into ``**kwargs``: ``**kwargs`` is forwarded to the provider SDK, + so a schema landing there would reach the request body. Implementations + without a native structured-output translation accept and ignore it, + leaving the caller on the prompt-engineering fallback. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -153,6 +254,14 @@ class BaseChatModelSetup(Resource): skill_discovery_prompt: str | None = None allowed_commands: List[str] = Field(default_factory=list) allowed_script_dirs: List[str] = Field(default_factory=list) + structured_output_strategy: StructuredOutputStrategy = Field( + default=StructuredOutputStrategy.AUTO, + description=( + "Intent about how an output schema should be applied. Whether native " + "structured output is actually used combines this policy with the " + "connection's model-dependent capability." + ), + ) @property @abstractmethod @@ -256,9 +365,8 @@ def chat( # Call chat model connection to execute chat merged_kwargs = self.model_kwargs.copy() merged_kwargs.update(kwargs) - return self._get_connection().chat( - messages, tools=self._get_tools(), **merged_kwargs - ) + connection = self._get_connection() + return connection.chat(messages, tools=self._get_tools(), **merged_kwargs) def _record_token_metrics( self, model_name: str, prompt_tokens: int, completion_tokens: int diff --git a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py index 921662360..d6ec5fe5f 100644 --- a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py +++ b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py @@ -18,12 +18,14 @@ from typing import Any, Dict, List, Sequence import pytest -from pydantic import Field, ValidationError +from pydantic import BaseModel, Field, ValidationError +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, BaseChatModelSetup, + StructuredOutputStrategy, ) from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.tools.tool import Tool @@ -41,18 +43,29 @@ def model_kwargs(self) -> Dict[str, Any]: return {"model": self.model} +class _Answer(BaseModel): + """A representative BaseModel output schema.""" + + text: str + + class _RecordingConnection(BaseChatModelConnection): - """Connection that captures the messages it receives for inspection.""" + """Connection that captures the messages and kwargs it receives for inspection.""" captured_messages: List[ChatMessage] = Field(default_factory=list) + captured_kwargs: Dict[str, Any] = Field(default_factory=dict) + captured_output_schema: OutputSchema | None = None def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: self.captured_messages = list(messages) + self.captured_kwargs = dict(kwargs) + self.captured_output_schema = output_schema return ChatMessage(role=MessageRole.ASSISTANT, content="ok") @@ -125,3 +138,70 @@ def test_chat_refills_template_on_subsequent_invocations() -> None: assert len(connection.captured_messages) == 2 assert connection.captured_messages[0].content == "Task: v1" assert connection.captured_messages[1].content == "tool result" + + +def test_default_capability_predicate_is_false() -> None: + """A connection reports no native structured output for any model by default.""" + connection = _RecordingConnection() + + assert connection.supports_native_structured_output("gpt-4o") is False + assert connection.supports_native_structured_output("gpt-3.5-turbo") is False + assert connection.supports_native_structured_output(None) is False + + +def test_setup_routes_output_schema_through_to_connection() -> None: + """chat() forwards a caller's ``output_schema`` on to the connection intact. + + The setup filters what reaches the connection, so a schema it dropped or consumed + would leave the connection unable to apply one at all. That the schema cannot land + in ``**kwargs`` is a separate, tree-wide invariant covered by the connection + signature guard. + """ + setup = _RecordingChatModelSetup(connection="c", model="m") + connection = _RecordingConnection() + setup._resolved_connection = connection + + schema = OutputSchema(output_schema=_Answer) + setup.chat([], output_schema=schema) + + assert connection.captured_output_schema is schema + + +def test_structured_output_strategy_defaults_to_auto() -> None: + """The setup policy defaults to AUTO when unset.""" + setup = _RecordingChatModelSetup(connection="c", model="m") + + assert setup.structured_output_strategy is StructuredOutputStrategy.AUTO + + +@pytest.mark.parametrize("raw", ["NATIVE", "native", "Native"]) +def test_structured_output_strategy_coerces_name_and_value_case_insensitively( + raw: str, +) -> None: + """The policy coerces from either its name or its value, in any case. + + Java serializes this enum as its name ("NATIVE") and its own resolver accepts any + case, so a Python side that only accepted the lowercase value would reject what + Java sends. + """ + setup = _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=raw + ) + + assert setup.structured_output_strategy is StructuredOutputStrategy.NATIVE + + +def test_auto_strategy_resolves_to_native_only_when_capable() -> None: + """AUTO defers to the model's capability.""" + assert StructuredOutputStrategy.AUTO.resolves_to_native(True) is True + assert StructuredOutputStrategy.AUTO.resolves_to_native(False) is False + + +def test_native_strategy_forces_native_regardless_of_capability() -> None: + """NATIVE resolves to native even when the model is not capable.""" + assert StructuredOutputStrategy.NATIVE.resolves_to_native(False) is True + + +def test_prompt_strategy_never_resolves_to_native() -> None: + """PROMPT never resolves to native even when the model is capable.""" + assert StructuredOutputStrategy.PROMPT.resolves_to_native(True) is False diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py index 625bb69f5..0055996f2 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py @@ -29,6 +29,7 @@ from pyflink.datastream import KeySelector from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -95,6 +96,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Generate a tool call or a response according to input.""" diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py index 0b30f9afb..c046da9cb 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py @@ -18,6 +18,7 @@ from pyflink.datastream import KeySelector from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -113,6 +114,7 @@ def chat( self, messages: list[ChatMessage], tools: list[BaseTool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: object, ) -> ChatMessage: """Return a tool call for user input, or echo the tool response.""" diff --git a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py index c077c6c8e..ed9573bb7 100644 --- a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py +++ b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py @@ -24,6 +24,7 @@ from pydantic import Field, PrivateAttr from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -168,9 +169,16 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Direct communication with Anthropic model service for chat conversation.""" + """Direct communication with Anthropic model service for chat conversation. + + ``output_schema`` is accepted and ignored: 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. + """ anthropic_tools = None if tools is not None: anthropic_tools = [ diff --git a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py index 18a092128..0bc5f9f40 100644 --- a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py @@ -21,6 +21,7 @@ from openai import NOT_GIVEN, AzureOpenAI from pydantic import Field, PrivateAttr +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -117,6 +118,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -127,6 +129,11 @@ def chat( Input message sequence tools : Optional[List] List of tools that can be called by the model + output_schema : OutputSchema | None + Accepted and ignored: 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) 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 7c36ec38e..1adf535ee 100644 --- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py +++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py @@ -21,6 +21,7 @@ from ollama import Client, Message from pydantic import Field +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -85,9 +86,16 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Process a sequence of messages, and return a response.""" + """Process a sequence of messages, and return a response. + + ``output_schema`` is accepted and ignored: 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. + """ ollama_messages = self.__convert_to_ollama_messages(messages) # Convert tool format diff --git a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py index 0d487772b..f5aa7c2c3 100644 --- a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py @@ -22,6 +22,7 @@ from pydantic import Field, PrivateAttr from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -138,6 +139,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -148,6 +150,11 @@ def chat( Input message sequence tools : Optional[List] List of tools that can be called by the model + output_schema : OutputSchema | None + Accepted and ignored: 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) diff --git a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py new file mode 100644 index 000000000..f526eddb0 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py @@ -0,0 +1,87 @@ +################################################################################ +# 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. +################################################################################# +"""Guards the ``output_schema`` parameter contract across every chat connection.""" + +import inspect +from typing import Iterator, List, Type + +from flink_agents.api.chat_models import java_chat_model as api_java_chat_model +from flink_agents.api.chat_models.chat_model import BaseChatModelConnection +from flink_agents.e2e_tests.e2e_tests_integration import ( + mock_chat_model_agent, + tool_parameter_injection_agent, +) +from flink_agents.integrations.chat_models import ollama_chat_model, tongyi_chat_model +from flink_agents.integrations.chat_models.anthropic import anthropic_chat_model +from flink_agents.integrations.chat_models.azure import azure_openai_chat_model +from flink_agents.integrations.chat_models.openai import openai_chat_model +from flink_agents.runtime.java import java_chat_model as runtime_java_chat_model + +# A class is only discoverable through __subclasses__() once it has been imported. +# Importing every module that defines a connection is what gives the walk below its +# reach — including the cross-language bridge and the e2e test doubles. +_MODULES_DEFINING_CONNECTIONS = ( + anthropic_chat_model, + api_java_chat_model, + azure_openai_chat_model, + mock_chat_model_agent, + ollama_chat_model, + openai_chat_model, + runtime_java_chat_model, + tongyi_chat_model, + tool_parameter_injection_agent, +) + + +def _subclasses_recursive(cls: Type[object]) -> Iterator[Type[object]]: + for subclass in cls.__subclasses__(): + yield subclass + yield from _subclasses_recursive(subclass) + + +def test_every_connection_declares_output_schema_param() -> None: + """Every connection in the tree must declare ``output_schema`` defaulting to None. + + A connection that omits the parameter absorbs a caller's ``output_schema`` into + ``**kwargs``. Connections forward ``**kwargs`` to their provider SDK — and the + cross-language bridge forwards it to Java as ``modelParams`` — so the schema would + reach the request body as an unknown field. Declaring the parameter is what makes + the leak impossible rather than merely unlikely. + + The tree is walked rather than hand-listed: a hand-kept list silently stops + guarding whatever it was never updated to mention, including the abstract base + itself. + """ + offenders: List[str] = [] + connections = [ + BaseChatModelConnection, + *_subclasses_recursive(BaseChatModelConnection), + ] + for cls in connections: + param = inspect.signature(cls.chat).parameters.get("output_schema") + if param is None: + offenders.append(f"{cls.__module__}.{cls.__qualname__}: parameter missing") + elif param.default is not None: + offenders.append( + f"{cls.__module__}.{cls.__qualname__}: default is " + f"{param.default!r}, expected None" + ) + + assert not offenders, ( + "connections violating the output_schema contract:\n" + "\n".join(offenders) + ) diff --git a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py index 6587a8cba..5dba9edb3 100644 --- a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py @@ -24,6 +24,7 @@ from dashscope import Generation from pydantic import Field +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -101,9 +102,16 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Process a sequence of messages, and return a response.""" + """Process a sequence of messages, and return a response. + + ``output_schema`` is accepted and ignored: 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. + """ tongyi_messages = self.__convert_to_tongyi_messages(messages) tongyi_tools: List[Dict[str, Any]] | None = ( diff --git a/python/flink_agents/runtime/java/java_chat_model.py b/python/flink_agents/runtime/java/java_chat_model.py index 10bc169c1..49d0c04b4 100644 --- a/python/flink_agents/runtime/java/java_chat_model.py +++ b/python/flink_agents/runtime/java/java_chat_model.py @@ -19,6 +19,7 @@ from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.java_chat_model import ( JavaChatModelConnection, @@ -56,6 +57,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Chat method that throws UnsupportedOperationException. @@ -63,6 +65,10 @@ def chat( This connection serves as a Java resource wrapper only. Chat operations should be performed on the Java side using the underlying Java chat model object. + + ``output_schema`` is accepted and ignored. Declaring it keeps a caller-supplied + schema out of ``**kwargs``, which crosses to Java as ``modelParams`` — a + provider-facing map, not a channel for framework execution metadata. """ java_messages = [ self._j_resource_adapter.fromPythonChatMessage(message) From da21d0816e884e56ab7c9bfb64e3f5308d2bc403 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 18 Jul 2026 22:01:44 -0700 Subject: [PATCH 2/2] [integrations][openai] Apply OpenAI native structured output Override the schema-carrying chat on the OpenAI completions connection to request the provider's native structured output when the caller supplies a schema the provider can honor natively. Native applies only when the schema is a POJO class (Java) or a BaseModel subclass (Python) and the effective model supports it. Capability is a connection-side predicate over the effective model, since the OpenAI json schema mode is model dependent: per the provider's structured output guide it is supported on gpt-4o-mini and the gpt-4o snapshots from 2024-08-06 onward, while gpt-4-turbo and gpt-3.5-turbo have json mode only. The predicate matches gpt-4o-mini by prefix (the whole family postdates the cutoff) and the capable gpt-4o snapshots exactly, so a pre-cutoff snapshot such as gpt-4o-2024-05-13 correctly falls back to the prompt path rather than failing at the provider. A RowTypeInfo schema stays on the prompt path. No path passes a schema on the chat call yet, so behavior is unchanged; native output is exercised by direct callers and the unit tests. --- .../openai/OpenAICompletionsConnection.java | 83 +++++++- .../OpenAICompletionsConnectionTest.java | 191 ++++++++++++++++++ .../chat_models/openai/openai_chat_model.py | 80 +++++++- .../test_openai_native_structured_output.py | 183 +++++++++++++++++ 4 files changed, 529 insertions(+), 8 deletions(-) create mode 100644 integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java create mode 100644 python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py diff --git a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java index 29d0dcf78..02a398282 100644 --- a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java +++ b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java @@ -22,11 +22,13 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.openai.client.OpenAIClient; import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.core.JsonSchemaLocalValidation; import com.openai.core.JsonValue; import com.openai.models.ChatModel; import com.openai.models.FunctionDefinition; import com.openai.models.FunctionParameters; import com.openai.models.ReasoningEffort; +import com.openai.models.ResponseFormatJsonSchema; import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionCreateParams; import com.openai.models.chat.completions.ChatCompletionFunctionTool; @@ -43,6 +45,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; /** * A chat model integration for the OpenAI Chat Completions service using the official Java SDK. @@ -119,11 +122,53 @@ public OpenAICompletionsConnection( this.client = builder.build(); } + // Models for which OpenAI documents json_schema strict Structured Outputs support. + // Source of truth: https://platform.openai.com/docs/guides/structured-outputs — json_schema is + // supported on the gpt-4o-mini and gpt-4o-2024-08-06 snapshots "and later"; gpt-4-turbo, + // earlier models, and gpt-3.5-turbo get JSON mode only. + // + // "and later" is temporal, not a name prefix: gpt-4o-2024-05-13 predates the cutoff and does + // NOT support Structured Outputs, so matching a bare "gpt-4o" prefix would misclassify it as + // capable and fail silently at the provider. Prefix matching is therefore used only for the + // gpt-4o-mini family, whose entire lifetime post-dates the cutoff; every other capable model is + // matched exactly. An unrecognized model reports not-capable and degrades to the prompt + // fallback rather than failing at the provider. + private static final String NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX = "gpt-4o-mini"; + private static final Set NATIVE_STRUCTURED_OUTPUT_MODELS = + Set.of("gpt-4o", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"); + + @Override + protected boolean supportsNativeStructuredOutput(String effectiveModel) { + if (effectiveModel == null) { + return false; + } + return effectiveModel.startsWith(NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX) + || NATIVE_STRUCTURED_OUTPUT_MODELS.contains(effectiveModel); + } + @Override public ChatMessage chat( List messages, List tools, Map modelParams) { + return doChat(messages, tools, modelParams, null); + } + + @Override + public ChatMessage chat( + List messages, + List tools, + Map modelParams, + Object outputSchema) { + return doChat(messages, tools, modelParams, outputSchema); + } + + private ChatMessage doChat( + List messages, + List tools, + Map modelParams, + Object outputSchema) { try { - ChatCompletionCreateParams params = buildRequest(messages, tools, modelParams); + ChatCompletionCreateParams params = + buildRequest(messages, tools, modelParams, outputSchema); ChatCompletion completion = client.chat().completions().create(params); ChatMessage response = OpenAIChatCompletionsUtils.convertFromOpenAIMessage( @@ -150,8 +195,13 @@ public ChatMessage chat( } } - private ChatCompletionCreateParams buildRequest( - List messages, List tools, Map rawModelParams) { + // Package-private so the request body (including the native response_format) can be asserted + // without issuing a live API call through the final OpenAI client. + ChatCompletionCreateParams buildRequest( + List messages, + List tools, + Map rawModelParams, + Object outputSchema) { Map modelParams = rawModelParams != null ? new HashMap<>(rawModelParams) : new HashMap<>(); @@ -170,6 +220,13 @@ private ChatCompletionCreateParams buildRequest( builder.tools(convertTools(tools, strictMode)); } + // Native structured output applies only for a POJO Class schema on a model the provider + // documents as capable; a RowTypeInfo (wrapped in OutputSchema) or an incapable model keeps + // the prompt-engineering fallback. + if (outputSchema instanceof Class && supportsNativeStructuredOutput(modelName)) { + builder.responseFormat(toNativeResponseFormat((Class) outputSchema)); + } + Object temperature = modelParams.remove("temperature"); if (temperature instanceof Number) { builder.temperature(((Number) temperature).doubleValue()); @@ -208,6 +265,26 @@ private ChatCompletionCreateParams buildRequest( return builder.build(); } + // Derives the strict json_schema response format from a POJO class via the SDK's typed + // structured-output builder. The Kotlin-facade StructuredOutputsKt.responseFormatFromClass is + // not callable from Java, so the response format is extracted through the typed builder, which + // generates the same strict draft-2020-12 schema, and then reattached to the standard builder. + private static ResponseFormatJsonSchema toNativeResponseFormat(Class schemaClass) { + return ChatCompletionCreateParams.builder() + .model(ChatModel.of("")) + .addUserMessage("") + .responseFormat(schemaClass, JsonSchemaLocalValidation.NO) + .build() + .rawParams() + .responseFormat() + .orElseThrow( + () -> + new IllegalStateException( + "OpenAI SDK did not produce a response_format for schema " + + schemaClass.getName())) + .asJsonSchema(); + } + private List convertTools(List tools, boolean strictMode) { List openaiTools = new ArrayList<>(tools.size()); for (Tool tool : tools) { diff --git a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java new file mode 100644 index 000000000..d95739095 --- /dev/null +++ b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java @@ -0,0 +1,191 @@ +/* + * 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. + */ + +package org.apache.flink.agents.integrations.chatmodels.openai; + +import com.openai.models.ResponseFormatJsonSchema; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +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; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +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 java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link OpenAICompletionsConnection}'s native structured-output behavior. These + * assert the built request body without a live API call by inspecting {@code buildRequest}, and + * exercise the model-dependent capability predicate directly. + */ +class OpenAICompletionsConnectionTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + + /** A representative POJO output schema. */ + public static class Person { + public String name; + public int age; + } + + private static OpenAICompletionsConnection connection() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder(OpenAICompletionsConnection.class.getName()) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("model", "gpt-4o") + .build(); + return new OpenAICompletionsConnection(desc, NOOP); + } + + private static Map params(String model) { + Map params = new HashMap<>(); + params.put("model", model); + return params; + } + + private static List userMessage() { + return List.of(new ChatMessage(MessageRole.USER, "hi")); + } + + @Test + @DisplayName("Native response_format json_schema strict applied for a POJO on a capable model") + void testNativeAppliedForPojoCapableModel() { + ChatCompletionCreateParams params = + connection().buildRequest(userMessage(), List.of(), params("gpt-4o"), Person.class); + + assertThat(params.responseFormat()).isPresent(); + ResponseFormatJsonSchema jsonSchema = params.responseFormat().get().asJsonSchema(); + assertThat(jsonSchema.jsonSchema().strict()).contains(true); + } + + @Test + @DisplayName("Native NOT applied for a POJO on an incapable model (prompt fallback)") + void testNativeNotAppliedForIncapableModel() { + ChatCompletionCreateParams params = + connection() + .buildRequest( + userMessage(), List.of(), params("gpt-3.5-turbo"), Person.class); + + assertThat(params.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied for a pre-cutoff same-family gpt-4o snapshot") + void testNativeNotAppliedForPreCutoffSnapshot() { + // gpt-4o-2024-05-13 predates the Structured Outputs cutoff even though it shares the gpt-4o + // prefix; treating it as capable would fail silently at the provider. + ChatCompletionCreateParams params = + connection() + .buildRequest( + userMessage(), + List.of(), + params("gpt-4o-2024-05-13"), + Person.class); + + assertThat(params.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied when no output schema is supplied") + void testNativeNotAppliedWhenSchemaNull() { + ChatCompletionCreateParams params = + connection().buildRequest(userMessage(), List.of(), params("gpt-4o"), null); + + assertThat(params.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied for a non-POJO schema form (POJO-only scope)") + void testNativeNotAppliedForNonPojoSchema() { + // A RowTypeInfo schema arrives wrapped in OutputSchema (not a bare POJO Class), so it must + // not activate native structured output; any non-Class schema object exercises the same + // instanceof gate. + Object nonClassSchema = "row"; + + ChatCompletionCreateParams params = + connection() + .buildRequest(userMessage(), List.of(), params("gpt-4o"), nonClassSchema); + + assertThat(params.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native applied for a POJO even when tools are bound (no empty-tools gate)") + void testNativeAppliedEvenWhenToolsBound() { + ChatCompletionCreateParams params = + connection() + .buildRequest( + userMessage(), + List.of(new StubTool()), + params("gpt-4o"), + Person.class); + + assertThat(params.responseFormat()).isPresent(); + } + + @Test + @DisplayName("Capability predicate accepts the documented capable models") + void testCapabilityPredicateAcceptsCapableModels() { + OpenAICompletionsConnection connection = connection(); + + assertThat(connection.supportsNativeStructuredOutput("gpt-4o")).isTrue(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-08-06")).isTrue(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-11-20")).isTrue(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4o-mini")).isTrue(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4o-mini-2024-07-18")).isTrue(); + } + + @Test + @DisplayName("Capability predicate rejects incapable, pre-cutoff, unknown, and null models") + void testCapabilityPredicateRejectsIncapableModels() { + OpenAICompletionsConnection connection = connection(); + + assertThat(connection.supportsNativeStructuredOutput("gpt-3.5-turbo")).isFalse(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4-turbo")).isFalse(); + assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-05-13")).isFalse(); + assertThat(connection.supportsNativeStructuredOutput("some-unknown-model")).isFalse(); + assertThat(connection.supportsNativeStructuredOutput(null)).isFalse(); + } + + /** Minimal tool stub; only its presence in the tools list matters. */ + private static class StubTool extends Tool { + StubTool() { + super(new ToolMetadata("add", "adds", "{\"type\":\"object\"}")); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return ToolResponse.success(null); + } + } +} diff --git a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py index f5aa7c2c3..fe5b1905a 100644 --- a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py @@ -19,7 +19,13 @@ import httpx from openai import NOT_GIVEN, OpenAI -from pydantic import Field, PrivateAttr + +# Private SDK module (leading underscore): the openai client itself uses this helper to +# build the strict json_schema for response_format, and there is no public re-export. It +# has existed at this path since the structured-output support in openai 1.66.3 (the +# pinned minimum). A future openai bump that moves it will fail loudly on import here. +from openai.lib._pydantic import to_strict_json_schema +from pydantic import BaseModel, Field, PrivateAttr from typing_extensions import override from flink_agents.api.agents.types import OutputSchema @@ -38,6 +44,47 @@ DEFAULT_OPENAI_MODEL = "gpt-4o-mini" +# Models with documented json_schema strict Structured Outputs support. Source of +# truth: https://platform.openai.com/docs/guides/structured-outputs +# +# json_schema is supported on the gpt-4o-mini and gpt-4o-2024-08-06 snapshots "and +# later"; gpt-4-turbo, earlier models, and gpt-3.5-turbo get JSON mode only. +# +# "and later" is temporal, not a name prefix: gpt-4o-2024-05-13 predates the cutoff +# and does NOT support Structured Outputs, so a bare "gpt-4o" prefix would misclassify +# it as capable and fail silently. Prefix matching is therefore used only for the +# gpt-4o-mini family, whose entire lifetime post-dates the cutoff; every other capable +# model is matched exactly. An unrecognized model reports not-capable and degrades to +# the prompt fallback rather than failing at the provider. +_NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX = "gpt-4o-mini" +_NATIVE_STRUCTURED_OUTPUT_MODELS = frozenset( + {"gpt-4o", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"} +) + + +def _native_response_format(output_schema: Any) -> Dict[str, Any] | None: + """Build the OpenAI ``response_format`` for a native structured-output request. + + Returns ``None`` (leaving behavior unchanged) 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 { + "type": "json_schema", + "json_schema": { + "name": model.__name__, + "schema": to_strict_json_schema(model), + "strict": True, + }, + } + class OpenAIChatModelConnection(BaseChatModelConnection): """The connection to the OpenAI LLM. @@ -135,6 +182,22 @@ def __get_client_kwargs(self) -> Dict[str, Any]: "http_client": self._http_client, } + @override + def supports_native_structured_output(self, effective_model: str | None) -> bool: + """Whether OpenAI documents json_schema strict support for ``effective_model``. + + See the module-level allowlist for the source of truth and the rationale for + matching the gpt-4o-mini family by prefix while matching other capable models + exactly. An unrecognized model reports ``False`` so it degrades to the + prompt-engineering fallback rather than failing at the provider. + """ + if not effective_model: + return False + return ( + effective_model.startswith(_NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX) + or effective_model in _NATIVE_STRUCTURED_OUTPUT_MODELS + ) + def chat( self, messages: Sequence[ChatMessage], @@ -151,10 +214,10 @@ def chat( tools : Optional[List] List of tools that can be called by the model output_schema : OutputSchema | None - Accepted and ignored: 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. + The schema the response should conform to, or ``None`` for an unconstrained + response. Native structured output is applied only for a ``BaseModel`` + schema on a model the provider documents as capable; a ``RowTypeInfo`` + schema or an incapable model keeps the prompt-engineering fallback. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -173,6 +236,13 @@ def chat( tool_spec["function"]["strict"] = strict tool_spec["function"]["parameters"]["additionalProperties"] = False + if output_schema is not None and self.supports_native_structured_output( + kwargs.get("model") + ): + response_format = _native_response_format(output_schema) + if response_format is not None: + kwargs["response_format"] = response_format + response = self.client.chat.completions.create( messages=convert_to_openai_messages(messages), tools=tool_specs or NOT_GIVEN, diff --git a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py new file mode 100644 index 000000000..c50d03080 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py @@ -0,0 +1,183 @@ +################################################################################ +# 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 +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.openai.openai_chat_model import ( + OpenAIChatModelConnection, +) +from flink_agents.plan.function import PythonFunction +from flink_agents.plan.tools.function_tool import FunctionTool + + +class Person(BaseModel): + """A representative BaseModel output schema.""" + + name: str + age: int + + +def _connection() -> OpenAIChatModelConnection: + conn = OpenAIChatModelConnection( + name="openai", api_key="test-key", api_base_url="http://localhost" + ) + mock_client = MagicMock() + mock_message = MagicMock() + mock_message.role = "assistant" + mock_message.content = "ok" + mock_message.tool_calls = None + mock_client.chat.completions.create.return_value.choices = [ + MagicMock(message=mock_message) + ] + mock_client.chat.completions.create.return_value.usage = None + conn._client = mock_client + return conn + + +def _create_call_kwargs(conn: OpenAIChatModelConnection) -> dict[str, Any]: + return conn.client.chat.completions.create.call_args.kwargs + + +def _add(a: int, b: int) -> int: + """Add two integers. + + Parameters + ---------- + a : int + first + b : int + second + + Returns: + ------- + int + sum + """ + return a + b + + +def test_native_applied_for_basemodel_capable_model() -> None: + """response_format json_schema strict applied for a BaseModel on a capable model.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model="gpt-4o", + output_schema=OutputSchema(output_schema=Person), + ) + response_format = _create_call_kwargs(conn)["response_format"] + assert response_format["type"] == "json_schema" + assert response_format["json_schema"]["strict"] is True + assert response_format["json_schema"]["schema"]["additionalProperties"] is False + + +def test_native_not_applied_for_incapable_model() -> None: + """Native NOT applied for a BaseModel on an incapable model (prompt fallback).""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model="gpt-3.5-turbo", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_for_pre_cutoff_snapshot() -> None: + """Native NOT applied for a pre-cutoff same-family gpt-4o snapshot. + + gpt-4o-2024-05-13 predates the Structured Outputs cutoff even though it shares the + gpt-4o prefix; treating it as capable would fail silently at the provider. + """ + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model="gpt-4o-2024-05-13", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_when_schema_none() -> None: + """Native NOT applied when no output schema is supplied.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model="gpt-4o", + output_schema=None, + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_for_row_type_info() -> None: + """Native NOT applied for a RowTypeInfo schema (BaseModel-only scope).""" + conn = _connection() + row_type = Types.ROW_NAMED(["name"], [Types.STRING()]) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model="gpt-4o", + output_schema=OutputSchema(output_schema=row_type), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_applied_even_when_tools_bound() -> None: + """Native applied for a BaseModel even when tools are bound (no empty-tools gate).""" + conn = _connection() + tool = FunctionTool(func=PythonFunction.from_callable(_add)) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + tools=[tool], + model="gpt-4o", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" in _create_call_kwargs(conn) + + +@pytest.mark.parametrize( + "model", + [ + "gpt-4o", + "gpt-4o-2024-08-06", + "gpt-4o-2024-11-20", + "gpt-4o-mini", + "gpt-4o-mini-2024-07-18", + ], +) +def test_capability_predicate_accepts_capable_models(model: str) -> None: + """The capability predicate accepts the documented capable models.""" + assert _connection().supports_native_structured_output(model) is True + + +@pytest.mark.parametrize( + "model", + [ + "gpt-3.5-turbo", + "gpt-4-turbo", + "gpt-4o-2024-05-13", + "some-unknown-model", + None, + ], +) +def test_capability_predicate_rejects_incapable_models(model: str | None) -> None: + """The capability predicate rejects incapable, pre-cutoff, unknown, and None models.""" + assert _connection().supports_native_structured_output(model) is False