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..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 @@ -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,44 @@ 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. + * + *

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 + * @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 + * @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/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..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 @@ -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,124 @@ void testChatRefillsTemplateOnSubsequentInvocations() { assertEquals("tool result", connection.capturedMessages.get(1).getContent()); } + @Test + @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); + + ChatMessage response = + connection.chat( + List.of(new ChatMessage(MessageRole.USER, "hi")), + List.of(), + modelParams, + null); + + // 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); + } + + @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 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() { + 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..d40f238ee 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 pydantic import Field, PrivateAttr, field_validator 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,61 @@ 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 + + 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), @@ -106,6 +227,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 +238,15 @@ 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 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.) @@ -153,6 +284,30 @@ 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. 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 @@ -256,9 +411,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..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 @@ -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,120 @@ 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_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. + + 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_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 + 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..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 @@ -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,9 +96,16 @@ 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.""" + """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 0b30f9afb..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 @@ -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,9 +114,16 @@ 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.""" + """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 c077c6c8e..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 @@ -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,17 @@ 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. + + 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 18a092128..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 @@ -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 + 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -136,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 7c36ec38e..b0061fb8d 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,17 @@ 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. + + 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 0d487772b..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 @@ -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 + 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -157,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 new file mode 100644 index 000000000..e31d37535 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py @@ -0,0 +1,168 @@ +################################################################################ +# 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 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 ( + 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, +) + + +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. + + 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) + ) + + +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 6587a8cba..5e7b52a1c 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,17 @@ 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. + + 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 10bc169c1..4073ce22e 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,14 +57,22 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + 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. + + 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