From 679078a2b5f568658c703dd53c225c7759099930 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 18 Jul 2026 22:00:25 -0700 Subject: [PATCH 1/3] [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 aead976409bffad31ebcb50a354db3e1f36ca2b4 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Fri, 24 Jul 2026 00:23:55 -0700 Subject: [PATCH 2/3] [api][python] Reject an output schema a chat connection cannot translate The schema-carrying chat path silently ignored an output schema when the connection had no native structured-output translation, so a caller could receive an unconstrained response with no way to tell it apart from a schema-conforming one. Java's default four-argument chat now throws UnsupportedOperationException for a non-null schema and still delegates to the abstract three-argument overload for null, so a connection that translates a schema overrides it and every other connection keeps its current behavior for schemaless calls. Python's chat is abstract, leaving no inherited body that could refuse a schema, so a shared guard on the connection base raises NotImplementedError and each connection without a native translation calls it as the first statement of chat, before any client access. The walk over connection subclasses now also asserts the rejection, holding a connection added to an imported module to the same contract. --- .../chat/model/BaseChatModelConnection.java | 21 +++-- .../api/chat/model/BaseChatModelTest.java | 31 +++++-- .../api/chat_models/chat_model.py | 34 +++++++- .../chat_models/tests/test_chat_model_base.py | 20 +++++ .../mock_chat_model_agent.py | 8 +- .../tool_parameter_injection_agent.py | 8 +- .../anthropic/anthropic_chat_model.py | 3 +- .../azure/azure_openai_chat_model.py | 3 +- .../chat_models/ollama_chat_model.py | 3 +- .../chat_models/openai/openai_chat_model.py | 3 +- .../test_output_schema_param_declared.py | 81 +++++++++++++++++++ .../chat_models/tongyi_chat_model.py | 3 +- .../runtime/java/java_chat_model.py | 11 ++- 13 files changed, 206 insertions(+), 23 deletions(-) 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 9cbad62bf..b4200958a 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 @@ -86,11 +86,12 @@ public abstract ChatMessage chat( * 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)}. + *

A schema must not be handed to a connection that has no native translation for it: this + * default implementation rejects a non-null {@code outputSchema} rather than dropping it, so an + * unconstrained response can never be mistaken for a schema-conforming one. A null {@code + * outputSchema} delegates to {@link #chat(List, List, Map)}. 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 @@ -98,12 +99,22 @@ public abstract ChatMessage chat( * @param outputSchema the schema the response should conform to, or null for an unconstrained * response * @return the chat response containing model outputs + * @throws UnsupportedOperationException if {@code outputSchema} is non-null and this connection + * has no native structured-output translation */ public ChatMessage chat( List messages, List tools, Map modelParams, @Nullable Object outputSchema) { + if (outputSchema != null) { + throw new UnsupportedOperationException( + getClass().getName() + + " has no native structured-output translation, so it cannot honor" + + " the given output schema. Override chat(List, List, Map, Object) to" + + " translate the schema natively, or pass no schema so the caller" + + " applies the prompt-engineering fallback."); + } return chat(messages, tools, modelParams); } } 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 117a70174..71f4fab06 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 @@ -323,8 +323,29 @@ void testChatRefillsTemplateOnSubsequentInvocations() { } @Test - @DisplayName("Default chat() overload ignores outputSchema and delegates to the 3-arg chat()") - void testDefaultChatOverloadIgnoresOutputSchema() { + @DisplayName("Default chat() overload rejects an outputSchema it cannot translate") + void testDefaultChatOverloadRejectsOutputSchema() { + RecordingConnection connection = new RecordingConnection(); + + // Dropping the schema instead would return an unconstrained response that the + // caller has no way to tell apart from a schema-conforming one. + assertThrows( + UnsupportedOperationException.class, + () -> + connection.chat( + List.of(new ChatMessage(MessageRole.USER, "hi")), + List.of(), + new HashMap<>(), + new Object())); + + // The rejection has to precede the delegation: a delegate-then-throw ordering + // would still issue a real provider request before failing. + assertNull(connection.capturedMessages); + } + + @Test + @DisplayName("Default chat() overload delegates to the 3-arg chat() for a null outputSchema") + void testDefaultChatOverloadDelegatesForNullOutputSchema() { RecordingConnection connection = new RecordingConnection(); Map modelParams = new HashMap<>(); modelParams.put("temperature", 0.5); @@ -334,10 +355,10 @@ void testDefaultChatOverloadIgnoresOutputSchema() { List.of(new ChatMessage(MessageRole.USER, "hi")), List.of(), modelParams, - new Object()); + null); - // 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. + // The 3-arg chat() ran (it is what produces "ok") and the overload added nothing + // to modelParams that could travel on to a provider SDK request. assertEquals("ok", response.getContent()); assertEquals(Map.of("temperature", 0.5), connection.capturedModelParams); } diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index a3443cf3d..c44561fb5 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -146,6 +146,35 @@ def supports_native_structured_output(self, effective_model: str | None) -> bool """ return False + def _reject_unsupported_output_schema( + self, output_schema: OutputSchema | None + ) -> None: + """Refuse an output schema this connection cannot translate natively. + + ``chat`` is abstract here, so there is no inherited body that could absorb a + schema loudly. A connection without a native structured-output translation + calls this as the first statement of its ``chat`` instead, which turns a + schema it could only drop into an error rather than an unconstrained response + that the caller would mistake for a schema-conforming one. + + Args: + output_schema: The schema the response should conform to. ``None`` returns + without effect. + + Raises: + NotImplementedError: If ``output_schema`` is not ``None``. + """ + if output_schema is None: + return + cls = type(self) + msg = ( + f"{cls.__module__}.{cls.__qualname__} has no native structured-output" + " translation, so it cannot honor the given output schema. Override chat()" + " to translate the schema natively, or pass no schema so the caller applies" + " the prompt-engineering fallback." + ) + raise NotImplementedError(msg) + DEFAULT_REASONING_PATTERNS: ClassVar[Tuple[re.Pattern[str], ...]] = ( re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), @@ -215,8 +244,9 @@ def chat( 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. + without a native structured-output translation reject a non-``None`` value + via ``_reject_unsupported_output_schema``, so a caller that wants the + prompt-engineering fallback must pass ``None``. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) 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 d6ec5fe5f..0c284b3b7 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 @@ -149,6 +149,26 @@ def test_default_capability_predicate_is_false() -> None: assert connection.supports_native_structured_output(None) is False +def test_output_schema_guard_rejects_a_schema() -> None: + """The guard refuses a schema a connection cannot translate natively. + + Dropping it instead would return an unconstrained response that the caller has no + way to tell apart from a schema-conforming one. + """ + connection = _RecordingConnection() + schema = OutputSchema(output_schema=_Answer) + + with pytest.raises(NotImplementedError, match="_RecordingConnection"): + connection._reject_unsupported_output_schema(schema) + + +def test_output_schema_guard_passes_through_none() -> None: + """A caller on the prompt-engineering fallback passes None and is let through.""" + connection = _RecordingConnection() + + assert connection._reject_unsupported_output_schema(None) is None + + def test_setup_routes_output_schema_through_to_connection() -> None: """chat() forwards a caller's ``output_schema`` on to the connection intact. 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 0055996f2..2aa03c264 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 @@ -99,7 +99,13 @@ def chat( output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Generate a tool call or a response according to input.""" + """Generate a tool call or a response according to input. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + structured-output translation. Declaring the parameter keeps a caller-supplied + schema out of ``**kwargs``. + """ + self._reject_unsupported_output_schema(output_schema) if "sum" in messages[-1].content: input = messages[-1].content # Validate the tool was bound before the model was invoked. 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 c046da9cb..9358fec9d 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 @@ -117,7 +117,13 @@ def chat( output_schema: OutputSchema | None = None, **kwargs: object, ) -> ChatMessage: - """Return a tool call for user input, or echo the tool response.""" + """Return a tool call for user input, or echo the tool response. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + structured-output translation. Declaring the parameter keeps a caller-supplied + schema out of ``**kwargs``. + """ + self._reject_unsupported_output_schema(output_schema) last_message = messages[-1] if last_message.role == MessageRole.TOOL: return ChatMessage(role=MessageRole.ASSISTANT, content=last_message.content) 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 ed9573bb7..2256f3502 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 @@ -174,11 +174,12 @@ def chat( ) -> ChatMessage: """Direct communication with Anthropic model service for chat conversation. - ``output_schema`` is accepted and ignored: this connection has no native + 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. """ + self._reject_unsupported_output_schema(output_schema) 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 0bc5f9f40..b653b4d67 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 @@ -130,7 +130,7 @@ 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 + Rejected when non-``None``: 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. @@ -143,6 +143,7 @@ def chat( ChatMessage Model response message """ + self._reject_unsupported_output_schema(output_schema) tool_specs = None if tools is not None: tool_specs = [to_openai_tool(metadata=tool.metadata) for tool in tools] 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 1adf535ee..b0061fb8d 100644 --- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py +++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py @@ -91,11 +91,12 @@ def chat( ) -> ChatMessage: """Process a sequence of messages, and return a response. - ``output_schema`` is accepted and ignored: this connection has no native + 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. """ + self._reject_unsupported_output_schema(output_schema) 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 f5aa7c2c3..68c2414ba 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 @@ -151,7 +151,7 @@ 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 + Rejected when non-``None``: 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. @@ -164,6 +164,7 @@ def chat( ChatMessage Model response message """ + self._reject_unsupported_output_schema(output_schema) tool_specs = None if tools is not None: tool_specs = [to_openai_tool(metadata=tool.metadata) for tool in tools] 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 index f526eddb0..e31d37535 100644 --- 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 @@ -20,6 +20,9 @@ import inspect from typing import Iterator, List, Type +from pydantic import BaseModel + +from flink_agents.api.agents.types import OutputSchema 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 ( @@ -48,12 +51,42 @@ ) +class _Answer(BaseModel): + """A representative output schema to hand to a connection.""" + + text: str + + def _subclasses_recursive(cls: Type[object]) -> Iterator[Type[object]]: for subclass in cls.__subclasses__(): yield subclass yield from _subclasses_recursive(subclass) +def _is_test_local(cls: Type[object]) -> bool: + """Whether ``cls`` is defined inside a ``tests`` package. + + A whole-tree pytest run imports every test module into one process, so doubles + declared inside a test file also turn up in the walk. Such a double may + deliberately accept a schema in order to assert on what a caller routed to it, + which is the opposite of the contract asserted here. + """ + return "tests" in cls.__module__.split(".") + + +def _connections_implementing_chat() -> List[Type[BaseChatModelConnection]]: + """Connections outside a ``tests`` package that supply their own ``chat`` body. + + A class that never overrides ``chat`` inherits only the abstract declaration, so + there is no body to assert on. + """ + return [ + cls + for cls in _subclasses_recursive(BaseChatModelConnection) + if not _is_test_local(cls) and cls.chat is not BaseChatModelConnection.chat + ] + + def test_every_connection_declares_output_schema_param() -> None: """Every connection in the tree must declare ``output_schema`` defaulting to None. @@ -85,3 +118,51 @@ def test_every_connection_declares_output_schema_param() -> None: assert not offenders, ( "connections violating the output_schema contract:\n" + "\n".join(offenders) ) + + +def _schema_rejection_failure( + cls: Type[BaseChatModelConnection], schema: OutputSchema +) -> str | None: + """Describe how ``cls.chat`` mishandles ``schema``, or ``None`` if it rejects it. + + ``__new__`` skips ``__init__``, so no credentials, no client and no network are + involved: the rejection has to happen before the first attribute access for the + call to get this far, which is what pins the guard to the top of ``chat``. + """ + name = f"{cls.__module__}.{cls.__qualname__}" + connection = cls.__new__(cls) + try: + cls.chat(connection, messages=[], tools=None, output_schema=schema) + except NotImplementedError: + return None + except Exception as exc: + return f"{name}: raised {type(exc).__name__} before rejecting the schema" + return f"{name}: accepted the schema instead of rejecting it" + + +def test_every_connection_rejects_an_output_schema_it_cannot_translate() -> None: + """A connection with no native translation must reject a schema, not drop it. + + Declaring the parameter only keeps the schema out of the provider request; on its + own it lets a connection silently return an unconstrained response that the caller + would treat as schema-conforming. Rejecting turns that into an error at the call. + + The tree is walked rather than hand-listed, so a connection added to a module that + is already imported is held to the contract for free. Reach still stops at those + imports: ``__subclasses__()`` only sees classes that have been imported, so a + connection in a brand-new module needs that module added at the top of this file. + """ + schema = OutputSchema(output_schema=_Answer) + connections = _connections_implementing_chat() + + assert connections, "the connection walk found nothing to check" + offenders = [ + failure + for cls in connections + if (failure := _schema_rejection_failure(cls, schema)) is not None + ] + + assert not offenders, ( + "connections that do not reject an untranslatable output_schema:\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 5dba9edb3..5e7b52a1c 100644 --- a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py @@ -107,11 +107,12 @@ def chat( ) -> ChatMessage: """Process a sequence of messages, and return a response. - ``output_schema`` is accepted and ignored: this connection has no native + 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. """ + self._reject_unsupported_output_schema(output_schema) 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 49d0c04b4..4073ce22e 100644 --- a/python/flink_agents/runtime/java/java_chat_model.py +++ b/python/flink_agents/runtime/java/java_chat_model.py @@ -60,16 +60,19 @@ def chat( output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Chat method that throws UnsupportedOperationException. + """Chat by forwarding the request to the wrapped Java connection. 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. + A non-``None`` ``output_schema`` is rejected: only messages, tools and kwargs + cross to the Java three-argument ``chat``, so a schema forwarded here would be + dropped on the way. Declaring the parameter 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. """ + self._reject_unsupported_output_schema(output_schema) java_messages = [ self._j_resource_adapter.fromPythonChatMessage(message) for message in messages From e537e4d875dd2d479285918c2935f4f559e753b0 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Fri, 24 Jul 2026 01:59:27 -0700 Subject: [PATCH 3/3] [api][python] Normalize a null structured output strategy to AUTO in Python An explicitly configured null structured output strategy resolved to AUTO in Java but raised a validation error in Python, so the same configuration behaved differently across the two languages. Java reads the value through ResourceDescriptor.getArgument, a plain map lookup, which cannot distinguish an absent key from one explicitly set to null. Rejecting an explicit null there would also reject an omitted one, so Python moves to the Java behavior rather than the reverse. A validator maps only null to AUTO, leaving every other unrecognized value to raise as before, and the field stays non-optional so a validated setup always carries a real strategy. The Java side gains the test for an explicitly null argument that its neighbours could not express, since Map.of rejects a null value. --- .../api/chat/model/BaseChatModelTest.java | 15 ++++++++++ .../api/chat_models/chat_model.py | 20 +++++++++++-- .../chat_models/tests/test_chat_model_base.py | 30 +++++++++++++++++++ 3 files changed, 63 insertions(+), 2 deletions(-) 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 71f4fab06..d27d79cb8 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 @@ -382,6 +382,21 @@ void testStructuredOutputStrategyDefaultsToAuto() { assertEquals(StructuredOutputStrategy.AUTO, setup.getStructuredOutputStrategy()); } + @Test + @DisplayName("Structured-output strategy defaults to AUTO when the descriptor argument is null") + void testStructuredOutputStrategyDefaultsToAutoForNullArgument() { + // A descriptor argument present with a null value is indistinguishable from an + // absent one here, so it resolves to the same default rather than failing. + TestChatModel model = + new TestChatModel( + new ResourceDescriptor( + TestChatModel.class.getName(), + Collections.singletonMap("structured_output_strategy", null)), + null); + + assertEquals(StructuredOutputStrategy.AUTO, model.getStructuredOutputStrategy()); + } + @Test @DisplayName("Structured-output strategy is read from the descriptor argument") void testStructuredOutputStrategyReadFromDescriptor() { diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index c44561fb5..d40f238ee 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -20,7 +20,7 @@ from enum import Enum from typing import Any, ClassVar, Dict, List, Mapping, Sequence, Tuple, cast -from pydantic import Field, PrivateAttr +from pydantic import Field, PrivateAttr, field_validator from typing_extensions import override from flink_agents.api.agents.types import OutputSchema @@ -289,10 +289,26 @@ class BaseChatModelSetup(Resource): 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." + "connection's model-dependent capability. An explicitly null value is " + "normalized to AUTO, so a validated setup always carries a real strategy." ), ) + @field_validator("structured_output_strategy", mode="before") + @classmethod + def _normalize_null_strategy(cls, value: Any) -> Any: + """Normalize an explicitly null strategy to the ``AUTO`` default. + + A configuration source can carry the key with a null value instead of + omitting it. Java cannot tell those two apart — its descriptor argument + lookup returns null in both cases and resolves them to ``AUTO`` — so an + explicit null resolves to ``AUTO`` here too rather than being rejected. + Unknown non-null values still fail validation. + """ + if value is None: + return StructuredOutputStrategy.AUTO + return value + @property @abstractmethod def model_kwargs(self) -> Dict[str, Any]: 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 0c284b3b7..462674bd7 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 @@ -211,6 +211,36 @@ def test_structured_output_strategy_coerces_name_and_value_case_insensitively( assert setup.structured_output_strategy is StructuredOutputStrategy.NATIVE +def test_structured_output_strategy_normalizes_explicit_none_to_auto() -> None: + """An explicitly null policy resolves to AUTO instead of being rejected. + + Java cannot distinguish a configuration that carries the key as null from one + that omits it, and resolves both to AUTO, so a null arriving on Python must not + fail validation or leave the attribute None. + """ + setup = _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=None + ) + + assert setup.structured_output_strategy is StructuredOutputStrategy.AUTO + + +@pytest.mark.parametrize("raw", ["bogus", ""]) +def test_structured_output_strategy_rejects_unrecognized_value(raw: str) -> None: + """Only null normalizes to AUTO; every other unrecognized value still raises. + + An empty string reaches this field in practice from an empty YAML scalar or an + unset environment substitution, and it must not be mistaken for an omitted value: + normalizing on falsiness rather than on null would accept it as AUTO. A non-empty + unrecognized name survives that same falsiness check, so it takes `"bogus"` to + catch a resolver that coerces any unknown string to AUTO. + """ + with pytest.raises(ValidationError): + _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=raw + ) + + 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