Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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.
*
* <p>Capability is <b>model-dependent</b>, 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 <i>effective</i> model at request-build time — the model
* actually being called, which per-request parameters may override.
*
* <p>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.
*
Expand All @@ -55,4 +77,33 @@ public ResourceType getResourceType() {
*/
public abstract ChatMessage chat(
List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams);

/**
* Process a chat request that carries an output schema, and return a chat response.
*
* <p>{@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.
*
* <p>This default implementation <b>ignores</b> {@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<ChatMessage> messages,
List<Tool> tools,
Map<String, Object> modelParams,
@Nullable Object outputSchema) {
return chat(messages, tools, modelParams);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ public abstract class BaseChatModelSetup extends Resource {
@Nullable protected String skillDiscoveryPrompt;
protected List<String> allowedCommands;
protected List<String> allowedScriptDirs;
protected StructuredOutputStrategy structuredOutputStrategy;

@Nullable protected BaseChatModelConnection connection;
protected final List<Tool> tools = new ArrayList<>();
Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -219,4 +224,15 @@ public List<String> getAllowedCommands() {
public List<String> 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;
}
}
Original file line number Diff line number Diff line change
@@ -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.
*
* <p>This expresses <b>policy</b> only. Whether a connection <i>can</i> apply the provider's native
* structured-output API is a separate, model-dependent <b>capability</b> 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.
*
* <ul>
* <li>{@code AUTO} defers to {@code modelCapable}: native when the effective model can, else
* the prompt-engineering fallback.
* <li>{@code NATIVE} always resolves to native, ignoring {@code modelCapable}, so an explicit
* user intent surfaces a provider error rather than silently degrading.
* <li>{@code PROMPT} never resolves to native.
* </ul>
*
* @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()));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,7 @@ void testChatResponseFormat() {
/** Connection that captures the messages passed to it for assertions. */
private static class RecordingConnection extends BaseChatModelConnection {
List<ChatMessage> capturedMessages;
Map<String, Object> capturedModelParams;

RecordingConnection() {
super(
Expand All @@ -249,6 +250,7 @@ private static class RecordingConnection extends BaseChatModelConnection {
public ChatMessage chat(
List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {
this.capturedMessages = new ArrayList<>(messages);
this.capturedModelParams = new HashMap<>(modelParams);
return new ChatMessage(MessageRole.ASSISTANT, "ok");
}
}
Expand Down Expand Up @@ -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<String, Object> 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() {
Expand Down
Loading
Loading