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 3cb2e655b..4f9daaf84 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 @@ -27,23 +27,32 @@ import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.skills.Skills; +import org.apache.flink.agents.api.subagent.SubagentSetup; import org.apache.flink.agents.api.tools.Tool; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.util.Preconditions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nullable; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Set; public abstract class BaseChatModelSetup extends Resource { + + private static final Logger LOG = LoggerFactory.getLogger(BaseChatModelSetup.class); + protected final String connectionName; protected String model; protected Object prompt; protected List toolNames; + protected final List subagentNames; @Nullable protected List skills; @Nullable protected String skillDiscoveryPrompt; protected List allowedCommands; @@ -59,6 +68,9 @@ public BaseChatModelSetup(ResourceDescriptor descriptor, ResourceContext resourc this.model = descriptor.getArgument("model"); this.prompt = descriptor.getArgument("prompt"); this.toolNames = descriptor.getArgument("tools"); + List declaredSubagents = descriptor.getArgument("subagents"); + this.subagentNames = + declaredSubagents == null ? new ArrayList<>() : new ArrayList<>(declaredSubagents); this.skills = descriptor.getArgument("skills"); List declaredCommands = descriptor.getArgument("allowed_commands"); this.allowedCommands = @@ -93,7 +105,7 @@ public void open() throws Exception { } if (this.skills != null) { this.skillDiscoveryPrompt = - this.resourceContext.generateAvailableSkillsPrompt(this.skills); + nullIfEmpty(this.resourceContext.generateAvailableSkillsPrompt(this.skills)); List mutable = this.toolNames == null ? new ArrayList<>() : new ArrayList<>(this.toolNames); if (!mutable.contains(Skills.LOAD_SKILL_TOOL)) { @@ -104,11 +116,54 @@ public void open() throws Exception { } this.toolNames = mutable; } + // Rebuilt from scratch: open() may run again on the same instance, and the callables must + // not accumulate. + this.tools.clear(); + Set callableNames = new LinkedHashSet<>(); if (this.toolNames != null) { for (String name : this.toolNames) { + Preconditions.checkState( + callableNames.add(name), "Duplicate callable name: %s", name); this.tools.add((Tool) this.resourceContext.getResource(name, ResourceType.TOOL)); } } + for (String name : this.subagentNames) { + // Tools are forbidden to carry the reserved prefix at registration, so a prefixed + // callable name can only come from this loop and a clash with a tool is impossible. + // Checked before the schema below, because a name declared twice is a mistake in the + // declaration whether or not it ends up registered. + Preconditions.checkState( + callableNames.add(SubagentSetup.CALLABLE_NAME_PREFIX + name), + "Duplicate callable name: %s", + SubagentSetup.CALLABLE_NAME_PREFIX + name); + Resource resource = this.resourceContext.getResource(name, ResourceType.AGENT); + // A sub-agent owned by the other language resolves to a bridge handle here, which + // carries no schema to declare, so it is rejected instead of silently dropped. + Preconditions.checkState( + resource instanceof SubagentSetup, + "Sub-agent %s must resolve to a SubagentSetup, but was %s", + name, + resource.getClass().getName()); + SubagentSetup setup = (SubagentSetup) resource; + String inputSchema = setup.getInputSchema(); + if (inputSchema == null) { + // Unlike a bridge handle this is a sub-agent the caller could have described, so + // it is dropped with a warning rather than failing the job: the rest of the + // callables stay usable. + LOG.warn( + "Sub-agent {} declares neither an input schema nor an input type, so there" + + " are no arguments for the model to build a call from and it is" + + " not offered as a callable.", + name); + continue; + } + this.tools.add(new SubagentTool(name, setup.getDescription(), inputSchema)); + } + } + + @Nullable + private static String nullIfEmpty(@Nullable String value) { + return value == null || value.isEmpty() ? null : value; } public abstract Map getParameters(); @@ -171,10 +226,11 @@ public ChatMessage chat( messages = promptMessages; } - if (this.skillDiscoveryPrompt != null && !this.skillDiscoveryPrompt.isEmpty()) { - int idx = ChatMessage.findFirstSystemMessage(messages); + if (this.skillDiscoveryPrompt != null) { + // Right after the first system message, or at the head when there is none. + int idx = ChatMessage.findFirstSystemMessage(messages) + 1; List mutated = new ArrayList<>(messages); - mutated.add(idx + 1, new ChatMessage(MessageRole.SYSTEM, this.skillDiscoveryPrompt)); + mutated.add(idx, new ChatMessage(MessageRole.SYSTEM, this.skillDiscoveryPrompt)); messages = mutated; } @@ -210,6 +266,16 @@ public List getToolNames() { return toolNames; } + /** Names of the {@code AGENT} resources this setup declares as delegable. */ + public List getSubagentNames() { + return subagentNames; + } + + @VisibleForTesting + public List getTools() { + return tools; + } + @Nullable public List getSkills() { return skills; diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/SubagentTool.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/SubagentTool.java new file mode 100644 index 000000000..5cbad8c4f --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/SubagentTool.java @@ -0,0 +1,75 @@ +/* + * 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 org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; + +/** + * Presents an {@code AGENT} resource to a chat model as a callable, so that the model can delegate + * a task by issuing a function call. The callable name carries the reserved {@link + * SubagentSetup#CALLABLE_NAME_PREFIX}, which is how the executing side tells a delegation apart + * from a plain tool call. + * + *

Metadata only: it carries the schema the model needs to build the call, and nothing else. The + * call itself is dispatched by resolving the {@code AGENT} resource at execution time, so {@link + * #call} is never reached. + */ +class SubagentTool extends Tool { + + SubagentTool(String agentName, String description, String inputSchema) { + super( + new ToolMetadata( + SubagentSetup.CALLABLE_NAME_PREFIX + agentName, + effectiveDescription(agentName, description), + inputSchema)); + } + + /** + * Falls back to a generic delegation description, so an undescribed sub-agent stays usable. + * Every description ends with the sub-agent marker, which replaces a separate listing message: + * the model learns that the callable is a delegation from the description alone. + */ + private static String effectiveDescription(String agentName, String description) { + String effective = + description == null || description.isBlank() + ? "Delegate a standalone task to sub-agent " + agentName + : description; + return effective + " This is subagent."; + } + + /** + * Sub-agents are declared to the model as plain functions: the model builds the call the same + * way it builds a tool call, and only the executing side tells them apart. + */ + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + throw new UnsupportedOperationException( + "SubagentTool is metadata-only; resolve the AGENT resource at execution time."); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/InputSchemas.java b/api/src/main/java/org/apache/flink/agents/api/subagent/InputSchemas.java new file mode 100644 index 000000000..abe2042ab --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/InputSchemas.java @@ -0,0 +1,89 @@ +/* + * 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.subagent; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +import javax.annotation.Nullable; + +/** + * Renders the input type a sub-agent declares as the JSON Schema a chat model is told about, so + * that a sub-agent which types its arguments does not also have to spell out their schema. + * + *

Rendering goes through the same Jackson generator {@code ReActAgent} renders a POJO output + * schema with, which keeps the two type-to-schema paths in this module on one implementation and + * adds no dependency. + */ +final class InputSchemas { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private InputSchemas() {} + + /** + * The schema of {@code type}, or {@code null} when the type states no shape a model could build + * a call from. That is {@link Object}, the type a sub-agent declares when it declares none, and + * any type that does not render as a JSON object, because the parameters of a callable must be + * one. + * + * @throws IllegalArgumentException if rendering the type fails, which is a declaration mistake + * worth failing on rather than dropping silently. + */ + @Nullable + static String fromType(@Nullable Class type) { + if (type == null || type == Object.class) { + return null; + } + JsonNode schema = render(type); + return "object".equals(schema.path("type").asText()) ? schema.toString() : null; + } + + private static JsonNode render(Class type) { + try { + return MAPPER.generateJsonSchema(type).getSchemaNode(); + } catch (JsonMappingException | IllegalArgumentException e) { + // Both are reachable: a class whose getters disagree on a property name fails the + // mapping, and one the generator has no JSON-object serializer for is refused with an + // IllegalArgumentException naming no remedy. + throw new IllegalArgumentException( + String.format( + "Sub-agent input type %s cannot be rendered as a JSON Schema, so it" + + " cannot be declared to a chat model. Declare an input schema" + + " explicitly, or use an input type whose fields are all" + + " JSON-Schema-renderable. Rendering it reported: %s", + type.getName(), e.getMessage()), + e); + } catch (StackOverflowError e) { + // The generator carries no cycle guard, so a class that reaches itself through its own + // members recurses until the stack is gone. A separate clause rather than another type + // on the union above because the error carries no message to quote, so this case has to + // name the cause itself. + throw new IllegalArgumentException( + String.format( + "Sub-agent input type %s is self-referential, so rendering it as a" + + " JSON Schema does not terminate and it cannot be declared to" + + " a chat model. Declare an input schema explicitly, or use an" + + " input type that does not refer back to itself.", + type.getName()), + e); + } + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java index 6357e1e67..45626d623 100644 --- a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java @@ -19,22 +19,104 @@ package org.apache.flink.agents.api.subagent; import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.resource.SerializableResource; +import javax.annotation.Nullable; + /** * Caller-facing definition of a sub-agent, registered in the agent plan as an {@code AGENT} * resource. */ public abstract class SubagentSetup extends SerializableResource { + /** + * Prefix of the callable name a sub-agent is exposed to a chat model under. Tools are forbidden + * to register under this prefix, so a prefixed callable name unambiguously addresses a + * sub-agent and the executing side routes it to the {@code AGENT} namespace. + */ + public static final String CALLABLE_NAME_PREFIX = "subagent_"; + + /** + * Tells a caller what this sub-agent is for, so that it can decide whether to delegate to it. + * This is routing information for the caller, not an instruction for the sub-agent itself. + */ + @JsonProperty("description") + private String description; + + /** + * JSON Schema of the arguments this sub-agent accepts, as declared explicitly. Null when it was + * not, in which case {@link #getInputSchema()} derives it from {@link #getInputType()}. + */ + @JsonProperty("input_schema") + @Nullable + private String inputSchema; + + protected SubagentSetup() { + this(""); + } + + protected SubagentSetup(String description) { + this(description, null); + } + + protected SubagentSetup(String description, @Nullable String inputSchema) { + this.description = description == null ? "" : description; + if (inputSchema != null && inputSchema.isBlank()) { + throw new IllegalArgumentException("Sub-agent input schema must not be blank."); + } + this.inputSchema = inputSchema; + } + @Override @JsonIgnore public ResourceType getResourceType() { return ResourceType.AGENT; } + public String getDescription() { + return description; + } + + /** + * Type of the arguments this sub-agent accepts, from which the schema declared to a chat model + * is derived. Override to type the arguments instead of spelling out their schema; {@link + * Object}, the default, states no shape. + * + *

Ignored for JSON, like {@link #getResourceType()}: it is behavior, not state, and writing + * it would make the plan JSON carry a Java class name the Python side cannot read. + */ + @JsonIgnore + public Class getInputType() { + return Object.class; + } + + /** + * Type the result is converted to before it is reported back to the caller. Override to state + * the shape of a result that is not already JSON-compatible, or to narrow a wider one; {@link + * Object}, the default, reports the result as it arrived. + * + *

Ignored for JSON for the same reason as {@link #getInputType()}. + */ + @JsonIgnore + public Class getResultType() { + return Object.class; + } + + /** + * JSON Schema of the arguments this sub-agent accepts: the one declared explicitly, else the + * one derived from {@link #getInputType()}. + * + * @return the schema, or {@code null} when neither says anything a model could build a call + * from, in which case this sub-agent is not declared to a chat model at all. + */ + @Nullable + public String getInputSchema() { + return inputSchema != null ? inputSchema : InputSchemas.fromType(getInputType()); + } + /** * Issues a new invocation with an implementation-assigned identity. This is the preferred form. */ diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java new file mode 100644 index 000000000..b4546bc81 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java @@ -0,0 +1,378 @@ +/* + * 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 com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentFuture; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Covers how a setup declares its {@code AGENT} resources to a chat model. */ +class BaseChatModelSetupSubagentTest { + + private static final String CUSTOM_SCHEMA = + "{\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}}}"; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private final Map store = new HashMap<>(); + + /** + * AGENT-type overrides, so one name can hold a tool and a sub-agent at the same time; an AGENT + * lookup falls back to {@link #store} when there is no override. + */ + private final Map agentStore = new HashMap<>(); + + private final StubConnection connection = + new StubConnection(new ResourceDescriptor("X", Map.of()), null); + + private final ResourceContext resourceContext = + new ResourceContext() { + @Override + public Resource getResource(String name, ResourceType type) { + Resource resource = + type == ResourceType.AGENT + ? agentStore.getOrDefault(name, store.get(name)) + : store.get(name); + if (resource == null) { + throw new IllegalArgumentException("No such resource: " + name); + } + return resource; + } + + @Override + public String generateAvailableSkillsPrompt(List skillNames) { + return "" + skillNames + ""; + } + + @Override + public List getSkillDirs(List skillNames) { + return List.of(); + } + }; + + private static class StubChatSetup extends BaseChatModelSetup { + StubChatSetup(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public Map getParameters() { + return new HashMap<>(); + } + } + + private static class StubConnection extends BaseChatModelConnection { + List capturedMessages; + List capturedTools; + + StubConnection(ResourceDescriptor d, ResourceContext c) { + super(d, c); + } + + @Override + public ChatMessage chat( + List messages, List tools, Map modelParams) { + this.capturedMessages = new ArrayList<>(messages); + this.capturedTools = new ArrayList<>(tools); + return new ChatMessage(MessageRole.ASSISTANT, "ok"); + } + } + + private static class StubTool extends Tool { + StubTool(String name) { + super(new ToolMetadata(name, "stub", "{}")); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return ToolResponse.success(""); + } + } + + /** Metadata-carrying sub-agent double: declaring it is all these tests exercise. */ + private static class StubSubagentSetup extends SubagentSetup { + + private static final long serialVersionUID = 1L; + + StubSubagentSetup(String description) { + super(description); + } + + StubSubagentSetup(String description, String inputSchema) { + super(description, inputSchema); + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt) { + throw new UnsupportedOperationException(); + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId) { + throw new UnsupportedOperationException(); + } + + @Override + public SubagentFuture submit( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + throw new UnsupportedOperationException(); + } + } + + /** Types its arguments, so the schema it is declared with is derived rather than written. */ + private static class TypedStubSubagentSetup extends StubSubagentSetup { + + private static final long serialVersionUID = 1L; + + TypedStubSubagentSetup(String description) { + super(description); + } + + @Override + public Class getInputType() { + return Review.class; + } + } + + /** The arguments of {@link TypedStubSubagentSetup}. */ + public static class Review { + private String path; + + public String getPath() { + return path; + } + } + + private StubChatSetup setupWith(Map extraArgs) { + store.put("conn", connection); + Map args = new HashMap<>(); + args.put("connection", "conn"); + args.putAll(extraArgs); + return new StubChatSetup(new ResourceDescriptor("X", args), resourceContext); + } + + @Test + void declaredSubagentsReachTheModelAsCallablesAfterTheTools() throws Exception { + store.put("lookup", new StubTool("lookup")); + store.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + StubChatSetup setup = + setupWith(Map.of("tools", List.of("lookup"), "subagents", List.of("reviewer"))); + + setup.open(); + setup.chat(new ArrayList<>()); + + assertThat(connection.capturedTools).hasSize(2); + assertThat(connection.capturedTools.get(0).getMetadata().getName()).isEqualTo("lookup"); + ToolMetadata delegated = connection.capturedTools.get(1).getMetadata(); + assertThat(delegated.getName()).isEqualTo("subagent_reviewer"); + assertThat(delegated.getDescription()).isEqualTo("Reviews a file. This is subagent."); + assertThat(delegated.getInputSchema()).isEqualTo(CUSTOM_SCHEMA); + } + + @Test + void anUndescribedSubagentIsStillDelegable() throws Exception { + store.put("reviewer", new StubSubagentSetup("", CUSTOM_SCHEMA)); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer"))); + + setup.open(); + + assertThat(setup.getTools().get(0).getMetadata().getName()).isEqualTo("subagent_reviewer"); + assertThat(setup.getTools().get(0).getMetadata().getDescription()) + .isEqualTo("Delegate a standalone task to sub-agent reviewer This is subagent."); + assertThat(setup.getTools().get(0).getMetadata().getInputSchema()).isEqualTo(CUSTOM_SCHEMA); + } + + @Test + void aSubagentThatTypesItsArgumentsIsDeclaredWithTheDerivedSchema() throws Exception { + store.put("reviewer", new TypedStubSubagentSetup("Reviews a file.")); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer"))); + + setup.open(); + + String declared = setup.getTools().get(0).getMetadata().getInputSchema(); + JsonNode schema = MAPPER.readTree(declared); + assertThat(schema.path("type").asText()).isEqualTo("object"); + assertThat(schema.path("properties").path("path").path("type").asText()) + .isEqualTo("string"); + } + + /** + * A sub-agent that states no shape for its arguments leaves the model nothing to build a call + * from, so it is dropped rather than declared as a callable it could only misuse. + */ + @Test + void aSubagentThatStatesNoInputShapeIsNotOfferedToTheModel() throws Exception { + store.put("reviewer", new StubSubagentSetup("Reviews a file.")); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer"))); + + setup.open(); + + assertThat(setup.getTools()).isEmpty(); + } + + /** Dropping one is not dropping the rest: the other callables stay usable. */ + @Test + void aSubagentWithoutAnInputShapeDoesNotStopTheOthersFromBeingDeclared() throws Exception { + store.put("opaque", new StubSubagentSetup("Reviews a file.")); + store.put("coder", new StubSubagentSetup("Writes a patch.", CUSTOM_SCHEMA)); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("opaque", "coder"))); + + setup.open(); + + assertThat(setup.getTools()).hasSize(1); + assertThat(setup.getTools().get(0).getMetadata().getName()).isEqualTo("subagent_coder"); + } + + /** The reserved prefix keeps the two namespaces apart, so one name may serve both. */ + @Test + void aToolAndASubagentMayShareAName() throws Exception { + store.put("reviewer", new StubTool("reviewer")); + agentStore.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + StubChatSetup setup = + setupWith(Map.of("tools", List.of("reviewer"), "subagents", List.of("reviewer"))); + + setup.open(); + + assertThat(setup.getTools()).hasSize(2); + assertThat(setup.getTools().get(0).getMetadata().getName()).isEqualTo("reviewer"); + assertThat(setup.getTools().get(1).getMetadata().getName()).isEqualTo("subagent_reviewer"); + } + + @Test + void aRepeatedToolNameIsRejected() { + store.put("lookup", new StubTool("lookup")); + StubChatSetup setup = setupWith(Map.of("tools", List.of("lookup", "lookup"))); + + assertThatThrownBy(setup::open) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Duplicate callable name: lookup"); + } + + @Test + void aRepeatedSubagentNameIsRejected() { + store.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer", "reviewer"))); + + assertThatThrownBy(setup::open) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Duplicate callable name: subagent_reviewer"); + } + + /** + * A sub-agent owned by the other language resolves to a bridge handle with no schema on it, so + * declaring it must fail loudly rather than leave the model with a callable it cannot call. + */ + @Test + void aSubagentWithoutASetupIsRejected() { + store.put("reviewer", new StubTool("reviewer")); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer"))); + + assertThatThrownBy(setup::open) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("must resolve to a SubagentSetup"); + } + + @Test + void reopeningDoesNotDuplicateTheCallables() throws Exception { + store.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer"))); + + setup.open(); + setup.open(); + + assertThat(setup.getTools()).hasSize(1); + } + + /** Delegation is offered through the callables alone: no listing message is injected. */ + @Test + void declaredSubagentsInjectNoListingMessage() throws Exception { + store.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + store.put("coder", new StubSubagentSetup("Writes a patch.", CUSTOM_SCHEMA)); + StubChatSetup setup = setupWith(Map.of("subagents", List.of("reviewer", "coder"))); + setup.open(); + + setup.chat( + new ArrayList<>( + List.of( + new ChatMessage(MessageRole.SYSTEM, "You are helpful."), + new ChatMessage(MessageRole.USER, "review it")))); + + assertThat(connection.capturedMessages).hasSize(2); + assertThat(connection.capturedMessages.get(0).getContent()).isEqualTo("You are helpful."); + assertThat(connection.capturedTools).hasSize(2); + } + + @Test + void onlyTheSkillListingIsInjectedWhenSubagentsAreDeclared() throws Exception { + store.put("reviewer", new StubSubagentSetup("Reviews a file.", CUSTOM_SCHEMA)); + store.put("load_skill", new StubTool("load_skill")); + store.put("bash", new StubTool("bash")); + StubChatSetup setup = + setupWith( + Map.of( + "subagents", List.of("reviewer"), + "skills", List.of("github"))); + setup.open(); + + setup.chat(new ArrayList<>(List.of(new ChatMessage(MessageRole.USER, "review it")))); + + assertThat(connection.capturedMessages).hasSize(2); + assertThat(connection.capturedMessages.get(0).getContent()) + .startsWith(""); + assertThat(connection.capturedMessages.get(1).getContent()).isEqualTo("review it"); + } + + @Test + void aSetupWithoutSubagentsInjectsNothing() throws Exception { + StubChatSetup setup = setupWith(Map.of()); + setup.open(); + + setup.chat(new ArrayList<>(List.of(new ChatMessage(MessageRole.USER, "hi")))); + + assertThat(connection.capturedMessages).hasSize(1); + assertThat(connection.capturedTools).isEmpty(); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java new file mode 100644 index 000000000..cae0801e2 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java @@ -0,0 +1,218 @@ +/* + * 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.subagent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.context.RunnerContext; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Pins the routing metadata {@link SubagentSetup} carries for a caller. */ +public class SubagentSetupTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** A setup that only carries metadata: invocation lives in the runtime layer. */ + private static class MetadataOnlySetup extends SubagentSetup { + + private static final long serialVersionUID = 1L; + + MetadataOnlySetup() { + super(); + } + + MetadataOnlySetup(String description) { + super(description); + } + + MetadataOnlySetup(String description, @Nullable String inputSchema) { + super(description, inputSchema); + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt) { + throw new UnsupportedOperationException(); + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId) { + throw new UnsupportedOperationException(); + } + + @Override + public SubagentFuture submit( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + throw new UnsupportedOperationException(); + } + } + + /** Types what it takes and what it returns instead of spelling out a schema. */ + private static class TypedSetup extends MetadataOnlySetup { + + private static final long serialVersionUID = 1L; + + private final Class inputType; + private final Class resultType; + + TypedSetup(Class inputType, Class resultType) { + super("Reviews a file."); + this.inputType = inputType; + this.resultType = resultType; + } + + @Override + public Class getInputType() { + return inputType; + } + + @Override + public Class getResultType() { + return resultType; + } + } + + /** The arguments of the typed setups above. */ + public static class Review { + private String path; + private int lines; + + public String getPath() { + return path; + } + + public int getLines() { + return lines; + } + } + + /** Reaches itself through its own member, so rendering it as a schema does not terminate. */ + public static class Cyclic { + private Cyclic next; + + public Cyclic getNext() { + return next; + } + } + + @Test + void aSetupThatDeclaresNothingStatesNoShapeForItsArguments() { + MetadataOnlySetup setup = new MetadataOnlySetup(); + + assertThat(setup.getDescription()).isEmpty(); + assertThat(setup.getInputType()).isEqualTo(Object.class); + assertThat(setup.getResultType()).isEqualTo(Object.class); + assertThat(setup.getInputSchema()).isNull(); + } + + @Test + void aDescriptionAloneStillStatesNoInputShape() { + MetadataOnlySetup setup = new MetadataOnlySetup("Reviews a changed file."); + + assertThat(setup.getDescription()).isEqualTo("Reviews a changed file."); + assertThat(setup.getInputSchema()).isNull(); + } + + @Test + void anAbsentDescriptionReadsAsEmptyRatherThanNull() { + assertThat(new MetadataOnlySetup(null).getDescription()).isEmpty(); + } + + @Test + void aBlankInputSchemaIsRejectedAtConstruction() { + assertThatThrownBy(() -> new MetadataOnlySetup("desc", " ")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("input schema must not be blank"); + } + + /** Absent is not blank: it leaves the schema to be derived from the input type. */ + @Test + void anAbsentInputSchemaIsLeftToTheInputType() { + assertThat(new MetadataOnlySetup("desc", null).getInputSchema()).isNull(); + } + + @Test + void anInputTypeIsRenderedAsTheInputSchema() throws Exception { + TypedSetup setup = new TypedSetup(Review.class, Object.class); + + JsonNode schema = MAPPER.readTree(setup.getInputSchema()); + assertThat(schema.path("type").asText()).isEqualTo("object"); + assertThat(schema.path("properties").path("path").path("type").asText()) + .isEqualTo("string"); + assertThat(schema.path("properties").path("lines").path("type").asText()) + .isEqualTo("integer"); + } + + @Test + void anExplicitInputSchemaWinsOverTheInputType() { + String declared = "{\"type\":\"object\",\"properties\":{\"prompt\":{\"type\":\"string\"}}}"; + + assertThat(new MetadataOnlySetup("desc", declared).getInputSchema()).isEqualTo(declared); + } + + /** + * The parameters of a callable must be a JSON object, so a type that renders as anything else + * declares no shape a model could build a call from. + */ + @Test + void anInputTypeThatRendersAsNoObjectStatesNoSchema() { + assertThat(new TypedSetup(String.class, Object.class).getInputSchema()).isNull(); + } + + @Test + void aSelfReferentialInputTypeIsRejected() { + assertThatThrownBy(() -> new TypedSetup(Cyclic.class, Object.class).getInputSchema()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("is self-referential"); + } + + /** The declared types drive behavior, so they must not leak into the cross-language plan. */ + @Test + void theDeclaredTypesStayOutOfThePlanJson() throws Exception { + String json = MAPPER.writeValueAsString(new TypedSetup(Review.class, Review.class)); + + assertThat(json).doesNotContain("inputType").doesNotContain("resultType"); + } + + /** + * The plan JSON is a cross-language contract: these two keys are what the Python side reads, so + * they are pinned literally rather than through the getters. + */ + @Test + void theMetadataSerializesUnderTheCrossLanguageKeys() throws Exception { + String customSchema = + "{\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}}}"; + + String json = + MAPPER.writeValueAsString(new MetadataOnlySetup("Reviews a file.", customSchema)); + + assertThat(json) + .contains("\"description\":\"Reviews a file.\"") + .contains("\"input_schema\""); + assertThat(json).doesNotContain("inputSchema"); + Map parsed = MAPPER.readValue(json, Map.class); + assertThat(parsed).containsEntry("input_schema", customSchema); + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java index 051f1fff8..869f4f9a2 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java @@ -684,11 +684,29 @@ private ResourceProvider createResourceProvider( /** Adds a resource provider to the resourceProviders map. */ private void addResourceProvider(ResourceProvider provider) { checkNoRouterModelNameClash(provider); + checkToolNameNotReserved(provider); resourceProviders .computeIfAbsent(provider.getType(), k -> new HashMap<>()) .put(provider.getName(), provider); } + /** + * A tool name must not carry the reserved {@code subagent_} prefix: sub-agent callables are + * exposed to the model under that prefix, and dispatch routes any prefixed call to the {@code + * AGENT} namespace, so a tool registered under the prefix could never be called. Fail clearly + * at plan-construction time rather than at call time. + */ + private void checkToolNameNotReserved(ResourceProvider provider) { + if (provider.getType() == TOOL + && provider.getName().startsWith(SubagentSetup.CALLABLE_NAME_PREFIX)) { + throw new IllegalArgumentException( + String.format( + "Tool name '%s' must not start with the reserved prefix '%s'," + + " which identifies sub-agent callables.", + provider.getName(), SubagentSetup.CALLABLE_NAME_PREFIX)); + } + } + /** * A name must not be registered as both a {@link ResourceType#CHAT_MODEL} and a {@link * ResourceType#MODEL_ROUTER}: an agent references either by putting it in {@code diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index a67571a8f..cfa538075 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -26,7 +26,10 @@ import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.resource.Resource; import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.api.subagent.SubagentSetup; import org.apache.flink.agents.api.tools.Tool; import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; import org.apache.flink.agents.api.tools.ToolParameterInjection; @@ -39,9 +42,12 @@ import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.tools.FunctionTool; +import org.apache.flink.agents.plan.utils.ToolResultUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; @@ -108,12 +114,28 @@ private static List buildToolCallExecutions( } Tool tool = null; + SubagentSetup agent = null; Exception preparationError = null; + // The reserved subagent_ prefix separates the two namespaces: tool registration rejects + // the prefix, so a prefixed callable name can only address a sub-agent. Matched once + // here and carried down, because resolving the AGENT resource throws when the name is + // absent and would otherwise have to be attempted for every plain tool call. + boolean delegated = name.startsWith(SubagentSetup.CALLABLE_NAME_PREFIX); try { - tool = (Tool) ctx.getResource(name, ResourceType.TOOL); + if (delegated) { + agent = + resolveSubagent( + name.substring(SubagentSetup.CALLABLE_NAME_PREFIX.length()), + ctx); + } else { + tool = (Tool) ctx.getResource(name, ResourceType.TOOL); + } } catch (Exception e) { preparationError = e; } + + // Injection is a tool-only contract, so a sub-agent call carries the model arguments + // unchanged. if (tool != null) { try { // Framework-owned injected args must win over model-provided values so hidden @@ -136,7 +158,8 @@ private static List buildToolCallExecutions( ExecutionReporters.started( ctx, ExecutionReporter.EntityTypes.TOOL, name, entityMetadata); - if (tool == null || preparationError != null) { + boolean unresolved = tool == null && agent == null; + if (unresolved || preparationError != null) { Exception failure = preparationError != null ? preparationError @@ -148,11 +171,7 @@ private static List buildToolCallExecutions( recordInlineResponse( id, ToolResponse.error( - String.format( - tool == null - ? "Tool %s does not exist." - : "Tool %s execute failed.", - name)), + prepFailureMessage(name, delegated, unresolved, failure)), diagnosticError, success, error, @@ -186,7 +205,9 @@ public ToolResponse call() throws Exception { return toolRef.call(new ToolParameters(callArguments)); } }; - executions.add(new ToolCallExecution(id, name, callable, entityMetadata)); + executions.add( + new ToolCallExecution( + id, name, callable, entityMetadata, agent, callArguments)); } return executions; } @@ -197,21 +218,32 @@ private static void executeParallel( Map success, Map error, Map responses) { - List> callables = new ArrayList<>(executions.size()); + // Sub-agent calls already run through durable execution inside the setup, so they stay + // synchronous here and only the tool calls enter the durable batch. + List toolExecutions = new ArrayList<>(); for (ToolCallExecution execution : executions) { + if (execution.agent != null) { + dispatchAgentExecution(execution, ctx, success, error, responses); + } else { + toolExecutions.add(execution); + } + } + List> callables = new ArrayList<>(toolExecutions.size()); + for (ToolCallExecution execution : toolExecutions) { callables.add(execution.callable); } try { List> outcomes = ctx.durableExecuteAllAsync(callables); for (int i = 0; i < outcomes.size(); i++) { - recordOutcome(executions.get(i), outcomes.get(i), ctx, success, error, responses); + recordOutcome( + toolExecutions.get(i), outcomes.get(i), ctx, success, error, responses); } } catch (Exception e) { - for (ToolCallExecution execution : executions) { + for (ToolCallExecution execution : toolExecutions) { recordExecutionException(execution, e, success, error, responses); } } catch (Error e) { - for (ToolCallExecution execution : executions) { + for (ToolCallExecution execution : toolExecutions) { ExecutionReporters.failed( ctx, ExecutionReporter.EntityTypes.TOOL, @@ -232,6 +264,10 @@ private static void executeSequentially( Map error, Map responses) { for (ToolCallExecution execution : executions) { + if (execution.agent != null) { + dispatchAgentExecution(execution, ctx, success, error, responses); + continue; + } try { ToolResponse response = toolCallAsync @@ -325,6 +361,71 @@ private static void recordInlineResponse( } } + private static void dispatchAgentExecution( + ToolCallExecution execution, + RunnerContext ctx, + Map success, + Map error, + Map responses) { + try { + // submit() and await() already run through durable execution inside the setup, so + // wrapping the call again here would nest durable cursors. + SubagentResult result = execution.agent.submit(ctx, execution.agentArguments).await(); + recordAgentResult(execution, result, ctx, success, error, responses); + } catch (Exception e) { + recordExecutionException(execution, e, success, error, responses); + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + execution.name, + execution.entityMetadata, + e, + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + } + } + + private static void recordAgentResult( + ToolCallExecution execution, + SubagentResult result, + RunnerContext ctx, + Map success, + Map error, + Map responses) + throws Exception { + if (result.isSuccess()) { + success.put(execution.id, true); + responses.put( + execution.id, + ToolResponse.success( + ToolResultUtils.toChatMessageContent( + ToolResultUtils.normalizeAgentResult( + result.getResult(), execution.agent.getResultType())))); + ExecutionReporters.succeeded( + ctx, + ExecutionReporter.EntityTypes.TOOL, + execution.name, + execution.entityMetadata); + } else { + // The model sees why the delegation failed, so it can correct the call instead of + // repeating it blindly; the error map keeps the same detail for observability. + success.put(execution.id, false); + responses.put( + execution.id, + ToolResponse.error( + withReason( + String.format("Sub-agent %s execute failed", execution.name), + result.getErrorMessage()))); + error.put(execution.id, result.getErrorMessage()); + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + execution.name, + execution.entityMetadata, + result.getException(), + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + } + } + private static void recordExecutionException( ToolCallExecution execution, Exception exception, @@ -334,10 +435,37 @@ private static void recordExecutionException( success.put(execution.id, false); responses.put( execution.id, - ToolResponse.error(String.format("Tool %s execute failed.", execution.name))); + ToolResponse.error( + execution.agent != null + ? withReason( + String.format( + "Sub-agent %s execute failed", execution.name), + exception.getMessage()) + : String.format("Tool %s execute failed.", execution.name))); error.put(execution.id, exception.getMessage()); } + /** + * The message the model sees when a call could not even be prepared. A rejected sub-agent call + * carries the reason, so the model can correct the call instead of repeating it blindly. + * + * @param delegated whether the callable name addressed a sub-agent, decided once by the caller + * rather than matched again here + */ + private static String prepFailureMessage( + String name, boolean delegated, boolean unresolved, Exception failure) { + if (delegated) { + return withReason( + String.format("Sub-agent %s execute failed", name), failure.getMessage()); + } + return String.format( + unresolved ? "Tool %s does not exist." : "Tool %s execute failed.", name); + } + + private static String withReason(String message, @Nullable String reason) { + return reason == null || reason.isBlank() ? message + "." : message + ": " + reason; + } + private static void recordToolResponse( String id, ToolResponse response, @@ -356,17 +484,40 @@ private static final class ToolCallExecution { private final String name; private final DurableCallable callable; private final Map entityMetadata; + private final SubagentSetup agent; + private final Map agentArguments; private ToolCallExecution( String id, String name, DurableCallable callable, - Map entityMetadata) { + Map entityMetadata, + SubagentSetup agent, + Map agentArguments) { this.id = id; this.name = name; this.callable = callable; this.entityMetadata = entityMetadata; + this.agent = agent; + this.agentArguments = agentArguments; + } + } + + /** + * Resolves a sub-agent, in one lookup: the {@code AGENT} resource is fetched once and checked + * once here, and the caller carries the setup from then on. + */ + private static SubagentSetup resolveSubagent(String name, RunnerContext ctx) throws Exception { + Resource resource = ctx.getResource(name, ResourceType.AGENT); + if (!(resource instanceof SubagentSetup)) { + // A sub-agent owned by the other language resolves to a bridge handle here, which + // cannot be called through this path. + throw new IllegalArgumentException( + String.format( + "Sub-agent %s must resolve to a SubagentSetup, but was %s.", + name, resource == null ? "null" : resource.getClass().getName())); } + return (SubagentSetup) resource; } private static Map toolEntityMetadata( diff --git a/plan/src/main/java/org/apache/flink/agents/plan/utils/ToolResultUtils.java b/plan/src/main/java/org/apache/flink/agents/plan/utils/ToolResultUtils.java new file mode 100644 index 000000000..7ebf568f5 --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/utils/ToolResultUtils.java @@ -0,0 +1,182 @@ +/* + * 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.plan.utils; + +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +import javax.annotation.Nullable; + +import java.lang.reflect.Array; +import java.util.List; +import java.util.Map; + +/** + * Turns a sub-agent result into something a chat model can be told. + * + *

A sub-agent result reaches the caller as an opaque object, but from here on it is carried in a + * tool message and re-bound after a failover, so it must hold nothing that JSON cannot express. + * Such a payload is rejected with the path where it was found instead of being dropped or silently + * stringified. + * + *

A sub-agent that declares a result type is read through it: the type says how to interpret a + * result the checks below would refuse, and what comes out is only what the type declares. + */ +public final class ToolResultUtils { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** + * Reads a result as the type its sub-agent declares. A property the type does not declare is + * ignored rather than refused, so that a sub-agent returning more than it declares is narrowed + * to what it declares instead of failing the call; the Python side reads through pydantic, + * which ignores extra fields the same way. + */ + private static final ObjectMapper TYPED_MAPPER = + new ObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES); + + private ToolResultUtils() {} + + /** + * Rejects a payload JSON cannot express, and reduces the rest to plain maps, lists and scalars. + * + * @param raw the result a sub-agent produced + * @return the same value in generic form + */ + public static Object normalizeAgentResult(Object raw) { + return normalizeAgentResult(raw, null); + } + + /** + * As {@link #normalizeAgentResult(Object)}, but read through {@code resultType} first when the + * sub-agent declares one. + * + * @param raw the result a sub-agent produced + * @param resultType the type the sub-agent declares for its result, or {@code null} or {@link + * Object} when it declares none, in which case {@code raw} is taken as it arrived + * @return the result in generic form + */ + public static Object normalizeAgentResult(Object raw, @Nullable Class resultType) { + if (resultType == null || resultType == Object.class) { + requireJsonCompatible(raw, "result"); + return MAPPER.convertValue(MAPPER.valueToTree(raw), Object.class); + } + // Two conversions: the first reads the result as the declared type, which is what admits a + // payload the compatibility check below would refuse and what drops everything the type + // does not declare; the second reduces that back to plain maps, lists and scalars, so what + // is reported and re-bound after a failover is the same generic form the undeclared path + // produces. + Object typed = TYPED_MAPPER.convertValue(raw, resultType); + Object generic = MAPPER.convertValue(typed, Object.class); + // Still required: a declared type can render a field JSON cannot express, and the result + // outlives this call in a tool message and in state. + requireJsonCompatible(generic, "result"); + return generic; + } + + /** Renders a value as the content of the tool message handed back to the model. */ + public static String toChatMessageContent(Object value) throws Exception { + if (value == null) { + return "null"; + } + if (value instanceof List + || value instanceof Map + || value instanceof JsonNode + || value.getClass().isArray()) { + return MAPPER.writeValueAsString(value); + } + return String.valueOf(value); + } + + private static void requireJsonCompatible(Object value, String path) { + if (value == null || value instanceof String || value instanceof Boolean) { + return; + } + if (value instanceof Number) { + requireFiniteNumber((Number) value, path); + return; + } + if (value instanceof JsonNode) { + requireJsonNodeCompatible((JsonNode) value, path); + return; + } + if (value instanceof Map) { + for (Map.Entry entry : ((Map) value).entrySet()) { + if (!(entry.getKey() instanceof String)) { + throw new IllegalArgumentException( + "Map keys in sub-agent result must be strings at " + path); + } + requireJsonCompatible(entry.getValue(), path + "." + entry.getKey()); + } + return; + } + if (value instanceof List) { + List list = (List) value; + for (int i = 0; i < list.size(); i++) { + requireJsonCompatible(list.get(i), path + "[" + i + "]"); + } + return; + } + if (value.getClass().isArray()) { + int length = Array.getLength(value); + for (int i = 0; i < length; i++) { + requireJsonCompatible(Array.get(value, i), path + "[" + i + "]"); + } + return; + } + throw invalid(path, "found " + value.getClass().getName()); + } + + private static void requireJsonNodeCompatible(JsonNode node, String path) { + if (node.isPojo()) { + throw invalid(path, "POJONode is not supported"); + } + if (node.isFloatingPointNumber() && !Double.isFinite(node.doubleValue())) { + throw new IllegalArgumentException("Non-finite number in sub-agent result at " + path); + } + if (node.isArray()) { + for (int i = 0; i < node.size(); i++) { + requireJsonNodeCompatible(node.get(i), path + "[" + i + "]"); + } + return; + } + if (node.isObject()) { + node.fields() + .forEachRemaining( + entry -> + requireJsonNodeCompatible( + entry.getValue(), path + "." + entry.getKey())); + } + } + + private static IllegalArgumentException invalid(String path, String detail) { + return new IllegalArgumentException( + "Sub-agent result must be JSON-compatible at " + path + ", " + detail); + } + + private static void requireFiniteNumber(Number number, String path) { + if (number instanceof Double && !Double.isFinite(number.doubleValue())) { + throw new IllegalArgumentException("Non-finite number in sub-agent result at " + path); + } + if (number instanceof Float && !Float.isFinite(number.floatValue())) { + throw new IllegalArgumentException("Non-finite number in sub-agent result at " + path); + } + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java index 8c801ae43..b5a1b3b04 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java @@ -25,6 +25,7 @@ import org.apache.flink.agents.api.subagent.SubagentSetup; import org.apache.flink.agents.api.subagent.TestSubagentSetup; import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; +import org.apache.flink.agents.plan.tools.bash.BashTool; import org.junit.jupiter.api.Test; import java.util.Map; @@ -84,4 +85,23 @@ void nonSubagentAgentResourceIsRejected() { .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("must be a SubagentSetup or a ResourceDescriptor"); } + + /** + * Sub-agent callables reach the model under the reserved {@code subagent_} prefix, so a tool + * registered under that prefix could never be called and is rejected at plan-construction time. + */ + @Test + void toolNameWithTheReservedSubagentPrefixIsRejected() { + Agent agent = new Agent(); + agent.addResource( + "subagent_helper", + ResourceType.TOOL, + new BashTool( + ResourceDescriptor.Builder.newBuilder(BashTool.class.getName()).build(), + null)); + + assertThatThrownBy(() -> new AgentPlan(agent)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must not start with the reserved prefix 'subagent_'"); + } } diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java new file mode 100644 index 000000000..fa956243d --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java @@ -0,0 +1,438 @@ +/* + * 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.plan.actions; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.configuration.ReadableConfiguration; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.Outcome; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ToolRequestEvent; +import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.memory.BaseLongTermMemory; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentFuture; +import org.apache.flink.agents.api.subagent.SubagentFutures; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for dispatching a tool call to an {@code AGENT} resource. */ +class ToolCallActionSubagentTest { + + @Test + void delegatesToTheSubagentAndReportsItsNormalizedResult() throws Exception { + Map payload = new LinkedHashMap<>(); + payload.put("verdict", "approved"); + payload.put("findings", List.of("style")); + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok(payload)); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("{\"verdict\":\"approved\",\"findings\":[\"style\"]}"); + assertThat(response.getError()).doesNotContainKey("call-1"); + } + + @Test + void handsTheModelArgumentsToTheSubagentAsThePrompt() throws Exception { + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok("done")); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + assertThat(agent.prompts).containsExactly(Map.of("prompt", "review the diff")); + // A sub-agent call resolves through the setup, which owns its own durable execution. + assertThat(ctx.durableExecutions).isZero(); + } + + @Test + void reportsAFailedSubagentResultWithTheDetailExposedToTheModel() throws Exception { + RecordingSubagentSetup agent = + new RecordingSubagentSetup(SubagentResult.error("upstream refused")); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo("Sub-agent subagent_reviewer execute failed: upstream refused"); + assertThat(response.getError()).containsEntry("call-1", "upstream refused"); + } + + @Test + void reportsAFailureRaisedWhileSubmitting() throws Exception { + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok("unreachable")); + agent.submitFailure = new IllegalStateException("mailbox is full"); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo("Sub-agent subagent_reviewer execute failed: mailbox is full"); + assertThat(response.getError()).containsEntry("call-1", "mailbox is full"); + } + + @Test + void rejectsAResultJsonCannotExpress() throws Exception { + RecordingSubagentSetup agent = + new RecordingSubagentSetup(SubagentResult.ok(Map.of("handle", new Object()))); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .startsWith("Sub-agent subagent_reviewer execute failed") + .contains("result.handle"); + assertThat(response.getError().get("call-1")).contains("result.handle"); + } + + /** A declared result type is what admits a result JSON cannot express on its own. */ + @Test + void readsAResultThroughTheTypeTheSubagentDeclares() throws Exception { + RecordingSubagentSetup agent = + new TypedRecordingSubagentSetup(SubagentResult.ok(new Verdict(true, "clean"))); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("{\"approved\":true,\"note\":\"clean\"}"); + } + + /** The reserved prefix routes each namespace on its own, even under one shared name. */ + @Test + void routesAToolAndASubagentSharingANameToTheirOwnNamespace() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withAgent( + "reviewer", new RecordingSubagentSetup(SubagentResult.ok("done"))) + .withTool("reviewer", new StubTool("reviewer")); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent delegated = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(delegated.getSuccess()).containsEntry("call-1", true); + assertThat(delegated.getResponses().get("call-1").getResult()).isEqualTo("done"); + // A sub-agent call resolves through the setup, which owns its own durable execution. + assertThat(ctx.durableExecutions).isZero(); + + ToolCallAction.processToolRequest(toolRequest("reviewer"), ctx); + + ToolResponseEvent direct = ToolResponseEvent.fromEvent(ctx.sentEvents.get(1)); + assertThat(direct.getSuccess()).containsEntry("call-1", true); + assertThat(direct.getResponses().get("call-1").getResult()).isEqualTo("reviewer called"); + assertThat(ctx.durableExecutions).isOne(); + } + + @Test + void refusesAnAgentResourceThatCarriesNoCallableSetup() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + ctx.agents.put("reviewer", new StubTool("reviewer")); + + ToolCallAction.processToolRequest(toolRequest("subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo( + "Sub-agent subagent_reviewer execute failed: Sub-agent reviewer must" + + " resolve to a SubagentSetup, but was " + + StubTool.class.getName() + + "."); + assertThat(response.getError().get("call-1")) + .isEqualTo( + "Sub-agent reviewer must resolve to a SubagentSetup, but was " + + StubTool.class.getName() + + "."); + } + + @Test + void stillDispatchesAToolWhenBothKindsAreRegisteredUnderDifferentNames() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withAgent( + "reviewer", new RecordingSubagentSetup(SubagentResult.ok("done"))) + .withTool("queryOrder", new StubTool("queryOrder")); + + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("queryOrder called"); + assertThat(ctx.durableExecutions).isOne(); + } + + private static ToolRequestEvent toolRequest(String callableName) { + return new ToolRequestEvent( + "model", + List.of( + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of( + "name", + callableName, + "arguments", + Map.of("prompt", "review the diff"))))); + } + + /** Captures every prompt it is handed and resolves to a preset outcome. */ + private static class RecordingSubagentSetup extends SubagentSetup { + private final SubagentResult outcome; + private final List prompts = new ArrayList<>(); + private Exception submitFailure; + + RecordingSubagentSetup(SubagentResult outcome) { + super("Reviews a diff."); + this.outcome = outcome; + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt) throws Exception { + return submit(ctx, prompt, "session", "call"); + } + + @Override + public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId) + throws Exception { + return submit(ctx, prompt, sessionId, "call"); + } + + @Override + public SubagentFuture submit( + RunnerContext ctx, Object prompt, String sessionId, String callId) + throws Exception { + if (submitFailure != null) { + throw submitFailure; + } + prompts.add(prompt); + return new ResolvedSubagentFuture(sessionId, callId, outcome); + } + } + + /** Declares a result type, so its result is read through it. */ + private static class TypedRecordingSubagentSetup extends RecordingSubagentSetup { + TypedRecordingSubagentSetup(SubagentResult outcome) { + super(outcome); + } + + @Override + public Class getResultType() { + return Verdict.class; + } + } + + /** The result {@link TypedRecordingSubagentSetup} declares. */ + public static class Verdict { + private boolean approved; + private String note; + + public Verdict() {} + + public Verdict(boolean approved, String note) { + this.approved = approved; + this.note = note; + } + + public boolean isApproved() { + return approved; + } + + public void setApproved(boolean approved) { + this.approved = approved; + } + + public String getNote() { + return note; + } + + public void setNote(String note) { + this.note = note; + } + } + + private static class ResolvedSubagentFuture extends SubagentFuture { + private final SubagentResult outcome; + + ResolvedSubagentFuture(String sessionId, String callId, SubagentResult outcome) { + super(sessionId, callId); + this.outcome = outcome; + } + + @Override + public boolean isDone() { + return true; + } + + @Override + public SubagentResult await() { + return outcome; + } + + @Override + public SubagentFutures combine(SubagentFuture... others) { + throw new UnsupportedOperationException(); + } + } + + private static class StubTool extends Tool { + StubTool(String name) { + super(new ToolMetadata(name, "Stub.", "{}")); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return ToolResponse.success(getMetadata().getName() + " called"); + } + } + + private static class FakeRunnerContext implements RunnerContext { + private final List sentEvents = new ArrayList<>(); + private final Map tools = new LinkedHashMap<>(); + private final Map agents = new LinkedHashMap<>(); + private final AgentConfiguration config = new AgentConfiguration(Map.of()); + private int durableExecutions; + + FakeRunnerContext withTool(String name, Resource tool) { + tools.put(name, tool); + return this; + } + + FakeRunnerContext withAgent(String name, SubagentSetup agent) { + agents.put(name, agent); + return this; + } + + @Override + public void sendEvent(Event event) { + sentEvents.add(event); + } + + @Override + public MemoryObject getSensoryMemory() { + return null; + } + + @Override + public MemoryObject getShortTermMemory() { + return null; + } + + @Override + public BaseLongTermMemory getLongTermMemory() { + return null; + } + + @Override + public FlinkAgentsMetricGroup getAgentMetricGroup() { + return null; + } + + @Override + public FlinkAgentsMetricGroup getActionMetricGroup() { + return null; + } + + @Override + public Resource getResource(String name, ResourceType type) throws Exception { + Map registry = type == ResourceType.AGENT ? agents : tools; + Resource resource = registry.get(name); + if (resource == null) { + throw new IllegalArgumentException("Resource does not exist: " + name); + } + return resource; + } + + @Override + public ReadableConfiguration getConfig() { + return config; + } + + @Override + public Map getActionConfig() { + return Map.of(); + } + + @Override + public Object getActionConfigValue(String key) { + return null; + } + + @Override + public T durableExecute(DurableCallable callable) throws Exception { + durableExecutions++; + return callable.call(); + } + + @Override + public T durableExecuteAsync(DurableCallable callable) throws Exception { + durableExecutions++; + return callable.call(); + } + + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + durableExecutions++; + outcomes.add(Outcome.success(callable.call())); + } + return outcomes; + } + + @Override + public void close() {} + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/utils/ToolResultUtilsTest.java b/plan/src/test/java/org/apache/flink/agents/plan/utils/ToolResultUtilsTest.java new file mode 100644 index 000000000..d5cb12ff4 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/utils/ToolResultUtilsTest.java @@ -0,0 +1,216 @@ +/* + * 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.plan.utils; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class ToolResultUtilsTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** The result a sub-agent that declares a result type produces. */ + public static class Verdict { + private boolean approved; + private String note; + private double score; + + public Verdict() {} + + public Verdict(boolean approved, String note, double score) { + this.approved = approved; + this.note = note; + this.score = score; + } + + public boolean isApproved() { + return approved; + } + + public void setApproved(boolean approved) { + this.approved = approved; + } + + public String getNote() { + return note; + } + + public void setNote(String note) { + this.note = note; + } + + public double getScore() { + return score; + } + + public void setScore(double score) { + this.score = score; + } + } + + @Test + void normalizeReducesNestedContainersToPlainMapsAndLists() { + Map raw = new LinkedHashMap<>(); + raw.put("items", List.of(1, "two")); + raw.put("node", MAPPER.createObjectNode().put("flag", true)); + + Object normalized = ToolResultUtils.normalizeAgentResult(raw); + + assertThat(normalized) + .isInstanceOf(Map.class) + .isEqualTo(Map.of("items", List.of(1, "two"), "node", Map.of("flag", true))); + } + + @Test + void normalizeKeepsScalarsAndNull() { + assertThat(ToolResultUtils.normalizeAgentResult("done")).isEqualTo("done"); + assertThat(ToolResultUtils.normalizeAgentResult(3)).isEqualTo(3); + assertThat(ToolResultUtils.normalizeAgentResult(null)).isNull(); + } + + @Test + void normalizeReportsThePathOfAValueJsonCannotExpress() { + Map raw = Map.of("outer", List.of(new Object())); + + assertThatThrownBy(() -> ToolResultUtils.normalizeAgentResult(raw)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("result.outer[0]") + .hasMessageContaining("java.lang.Object"); + } + + @Test + void normalizeRejectsNonStringMapKeys() { + assertThatThrownBy( + () -> + ToolResultUtils.normalizeAgentResult( + Collections.singletonMap(1, "one"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Map keys in sub-agent result must be strings at result"); + } + + @Test + void normalizeRejectsNonFiniteNumbers() { + assertThatThrownBy( + () -> + ToolResultUtils.normalizeAgentResult( + Map.of("ratio", Double.POSITIVE_INFINITY))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Non-finite number in sub-agent result at result.ratio"); + + assertThatThrownBy(() -> ToolResultUtils.normalizeAgentResult(Float.NaN)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Non-finite number in sub-agent result at result"); + } + + @Test + void normalizeRejectsAPojoCarriedInsideAJsonTree() { + JsonNode tree = + MAPPER.createObjectNode().set("wrapped", MAPPER.getNodeFactory().pojoNode(this)); + + assertThatThrownBy(() -> ToolResultUtils.normalizeAgentResult(tree)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("result.wrapped") + .hasMessageContaining("POJONode is not supported"); + } + + @Test + void normalizeWalksArrays() { + Object normalized = ToolResultUtils.normalizeAgentResult(new int[] {1, 2}); + + assertThat(normalized).isEqualTo(List.of(1, 2)); + + assertThatThrownBy(() -> ToolResultUtils.normalizeAgentResult(new Object[] {new Object()})) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("result[0]"); + } + + /** + * Declaring a result type is what makes a result JSON cannot express reportable: the type says + * how to read it, and what comes out is only what the type declares. + */ + @Test + void aDeclaredResultTypeReadsAPojoIntoGenericForm() { + Object normalized = + ToolResultUtils.normalizeAgentResult( + new Verdict(true, "clean", 1.5), Verdict.class); + + assertThat(normalized).isEqualTo(Map.of("approved", true, "note", "clean", "score", 1.5)); + } + + @Test + void aDeclaredResultTypeNarrowsAWiderResultToWhatItDeclares() { + Object normalized = + ToolResultUtils.normalizeAgentResult( + Map.of("approved", true, "note", "clean", "score", 1.5, "extra", 1), + Verdict.class); + + assertThat(normalized).isEqualTo(Map.of("approved", true, "note", "clean", "score", 1.5)); + } + + /** The undeclared path is unchanged: a POJO is still what it refuses. */ + @Test + void anUndeclaredResultTypeStillRejectsAPojo() { + assertThatThrownBy( + () -> ToolResultUtils.normalizeAgentResult(new Verdict(true, "clean", 1))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must be JSON-compatible at result"); + } + + @Test + void declaringObjectIsTheSameAsDeclaringNothing() { + Map raw = Map.of("items", List.of(1, 2)); + + assertThat(ToolResultUtils.normalizeAgentResult(raw, Object.class)) + .isEqualTo(ToolResultUtils.normalizeAgentResult(raw)); + assertThat(ToolResultUtils.normalizeAgentResult(raw, null)) + .isEqualTo(ToolResultUtils.normalizeAgentResult(raw)); + } + + /** A declared type can still render a field JSON cannot express, so the check stays. */ + @Test + void aDeclaredResultTypeStillRejectsWhatJsonCannotExpress() { + Map raw = + Map.of("approved", true, "note", "clean", "score", Double.POSITIVE_INFINITY); + + assertThatThrownBy(() -> ToolResultUtils.normalizeAgentResult(raw, Verdict.class)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Non-finite number in sub-agent result at result.score"); + } + + @Test + void chatMessageContentSerializesContainersAndStringifiesScalars() throws Exception { + assertThat(ToolResultUtils.toChatMessageContent(null)).isEqualTo("null"); + assertThat(ToolResultUtils.toChatMessageContent("done")).isEqualTo("done"); + assertThat(ToolResultUtils.toChatMessageContent(7)).isEqualTo("7"); + assertThat(ToolResultUtils.toChatMessageContent(Map.of("a", 1))).isEqualTo("{\"a\":1}"); + assertThat(ToolResultUtils.toChatMessageContent(List.of(1, 2))).isEqualTo("[1,2]"); + assertThat(ToolResultUtils.toChatMessageContent(new int[] {1, 2})).isEqualTo("[1,2]"); + assertThat(ToolResultUtils.toChatMessageContent(MAPPER.createObjectNode().put("a", 1))) + .isEqualTo("{\"a\":1}"); + } +} diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index 11cf65791..ad9ee3b61 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -15,6 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +import logging import re from abc import ABC, abstractmethod from enum import Enum @@ -29,12 +30,16 @@ MessageRole, find_first_system_message, ) +from flink_agents.api.chat_models.subagent_tool import SubagentTool from flink_agents.api.metric_group import MetricGroup from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import Resource, ResourceType from flink_agents.api.skills import BASH_TOOL, LOAD_SKILL_TOOL +from flink_agents.api.subagent import CALLABLE_NAME_PREFIX, SubagentSetup from flink_agents.api.tools.tool import Tool +_LOG = logging.getLogger(__name__) + class StructuredOutputStrategy(str, Enum): """User intent about how an output schema should be applied to a chat request. @@ -322,6 +327,10 @@ class BaseChatModelSetup(Resource): _resolved_connection: BaseChatModelConnection | None = PrivateAttr(default=None) prompt: Prompt | str | None = None tools: List[str] | List[Tool] = Field(default_factory=list) + subagents: List[str] = Field( + default_factory=list, + description="Names of the AGENT resources this setup may delegate to.", + ) skills: List[str] | None = None skill_discovery_prompt: str | None = None allowed_commands: List[str] = Field(default_factory=list) @@ -382,17 +391,76 @@ def open(self) -> None: if self.skills is not None: self.skill_discovery_prompt = ( self.resource_context.generate_available_skills_prompt(*self.skills) + or None ) - self.tools.extend([LOAD_SKILL_TOOL, BASH_TOOL]) - if len(self.tools) > 0: - self.tools = [ - cast( + # Rebuilt from scratch: open() may run again on the same instance, and the + # callables must not accumulate. Sub-agent callables are derived from + # ``subagents`` rather than declared under ``tools``, so a previous build of + # them is dropped here instead of being read back as a declaration. + declared: List[str | Tool] = [ + entry for entry in self.tools if not isinstance(entry, SubagentTool) + ] + declared_names = [ + entry if isinstance(entry, str) else entry.name for entry in declared + ] + if self.skills is not None: + for skill_tool in (LOAD_SKILL_TOOL, BASH_TOOL): + if skill_tool not in declared_names: + declared.append(skill_tool) + declared_names.append(skill_tool) + + callables: List[Tool] = [] + callable_names: set[str] = set() + for entry, name in zip(declared, declared_names, strict=True): + if name in callable_names: + msg = f"Duplicate callable name: {name}" + raise ValueError(msg) + callable_names.add(name) + callables.append( + entry + if isinstance(entry, Tool) + else cast( "Tool", - self.resource_context.get_resource(tool_name, ResourceType.TOOL), + self.resource_context.get_resource(name, ResourceType.TOOL), ) - for tool_name in self.tools - ] + ) + for name in self.subagents: + # Tools are forbidden to carry the reserved prefix at registration, so a + # prefixed callable name can only come from this loop and a clash with a + # tool is impossible. Checked before the schema below, because a name + # declared twice is a mistake in the declaration whether or not it ends + # up registered. + callable_name = CALLABLE_NAME_PREFIX + name + if callable_name in callable_names: + msg = f"Duplicate callable name: {callable_name}" + raise ValueError(msg) + callable_names.add(callable_name) + setup = self.resource_context.get_resource(name, ResourceType.AGENT) + # A sub-agent owned by the other language resolves to a bridge handle + # here, which carries no schema to declare, so it is rejected instead + # of silently dropped. + if not isinstance(setup, SubagentSetup): + msg = ( + f"Sub-agent {name} must resolve to a SubagentSetup, " + f"but was {type(setup).__name__}" + ) + raise TypeError(msg) + if setup.input_schema is None: + # Unlike a bridge handle this is a sub-agent the caller could have + # described, so it is dropped with a warning rather than failing the + # job: the rest of the callables stay usable. + _LOG.warning( + "Sub-agent %s declares neither an input schema nor an input" + " type, so there are no arguments for the model to build a" + " call from and it is not offered as a callable.", + name, + ) + continue + callables.append( + SubagentTool.of(name, setup.description, setup.input_schema) + ) + self.tools = callables def chat( self, @@ -438,17 +506,15 @@ def chat( prompt_messages.append(msg) messages = prompt_messages - if self.skills is not None: - index = find_first_system_message(messages) - messages = ( - messages[: index + 1] - + [ - ChatMessage( - role=MessageRole.SYSTEM, content=self.skill_discovery_prompt - ) - ] - + messages[index + 1 :] - ) + if self.skill_discovery_prompt: + # Right after the first system message, or at the head when there is none. + index = find_first_system_message(messages) + 1 + injected = [ + ChatMessage( + role=MessageRole.SYSTEM, content=self.skill_discovery_prompt + ) + ] + messages = list(messages[:index]) + injected + list(messages[index:]) # Call chat model connection to execute chat merged_kwargs = self.model_kwargs.copy() diff --git a/python/flink_agents/api/chat_models/subagent_tool.py b/python/flink_agents/api/chat_models/subagent_tool.py new file mode 100644 index 000000000..1e579b2ba --- /dev/null +++ b/python/flink_agents/api/chat_models/subagent_tool.py @@ -0,0 +1,88 @@ +################################################################################ +# 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. +################################################################################ +"""Presenting an AGENT resource to a chat model as a callable.""" + +import json +from typing import Any + +from typing_extensions import override + +from flink_agents.api.subagent import CALLABLE_NAME_PREFIX +from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType +from flink_agents.api.tools.utils import create_model_from_schema + + +class SubagentTool(Tool): + """Presents an AGENT resource to a chat model as a callable, so that the + model can delegate a task by issuing a function call. The callable name + carries the reserved :data:`CALLABLE_NAME_PREFIX`, which is how the + executing side tells a delegation apart from a plain tool call. + + Metadata only: it carries the schema the model needs to build the call, and + nothing else. The call itself is dispatched by resolving the AGENT resource + at execution time, so :meth:`call` is never reached. + """ + + @staticmethod + def of(name: str, description: str, input_schema: str) -> "SubagentTool": + """Build the callable declaration of the named sub-agent. + + Parameters + ---------- + name : str + The name the AGENT resource is registered under. The function name + the model issues is this name behind the reserved prefix. + description : str + What the sub-agent is for. Falls back to a generic delegation + description, so an undescribed sub-agent stays usable. Every + description ends with the sub-agent marker, which replaces a + separate listing message: the model learns that the callable is a + delegation from the description alone. + input_schema : str + JSON Schema of the arguments the sub-agent accepts. + """ + callable_name = CALLABLE_NAME_PREFIX + name + return SubagentTool( + metadata=ToolMetadata( + name=callable_name, + description=( + description or f"Delegate a standalone task to sub-agent {name}" + ) + + " This is subagent.", + # The cross-language contract carries a JSON schema string, while + # tool metadata holds a model type, so it is built here. + args_schema=create_model_from_schema( + callable_name, json.loads(input_schema) + ), + ) + ) + + @classmethod + @override + def tool_type(cls) -> ToolType: + """Sub-agents are declared to the model as plain functions: the model + builds the call the same way it builds a tool call, and only the + executing side tells them apart. + """ + return ToolType.FUNCTION + + @override + def call(self, *args: Any, **kwargs: Any) -> Any: + """Never reached: dispatch resolves the AGENT resource instead.""" + msg = "SubagentTool is metadata-only; resolve the AGENT resource at execution time." + raise NotImplementedError(msg) diff --git a/python/flink_agents/api/chat_models/tests/test_chat_model_subagents.py b/python/flink_agents/api/chat_models/tests/test_chat_model_subagents.py new file mode 100644 index 000000000..40b95e2e2 --- /dev/null +++ b/python/flink_agents/api/chat_models/tests/test_chat_model_subagents.py @@ -0,0 +1,387 @@ +################################################################################ +# 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. +################################################################################ +"""Covers how a setup declares its AGENT resources to a chat model.""" + +from typing import Any, Dict, List, Sequence + +import pytest +from pydantic import BaseModel, 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, + BaseChatModelSetup, +) +from flink_agents.api.chat_models.subagent_tool import SubagentTool +from flink_agents.api.resource import Resource, ResourceType +from flink_agents.api.resource_context import ResourceContext +from flink_agents.api.subagent import SubagentSetup +from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType + +CUSTOM_SCHEMA = '{"type":"object","properties":{"path":{"type":"string"}}}' + + +class _RecordingConnection(BaseChatModelConnection): + """Connection that captures what the setup hands to the model.""" + + captured_messages: List[ChatMessage] = Field(default_factory=list) + captured_tools: List[Tool] = Field(default_factory=list) + + def chat( + self, + messages: Sequence[ChatMessage], + tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, + **kwargs: Any, + ) -> ChatMessage: + """Record the request and answer with a fixed message.""" + self.captured_messages = list(messages) + self.captured_tools = list(tools or []) + return ChatMessage(role=MessageRole.ASSISTANT, content="ok") + + +class _StubSetup(BaseChatModelSetup): + """Setup whose model parameters are irrelevant to these tests.""" + + @property + def model_kwargs(self) -> Dict[str, Any]: + """Return no model settings.""" + return {} + + +class _EmptyArgs(BaseModel): + """Argument model of the stub tool.""" + + +class _StubTool(Tool): + """A plain tool, standing in for anything declared under ``tools``.""" + + @classmethod + def tool_type(cls) -> ToolType: + """Return the function tool type.""" + return ToolType.FUNCTION + + def call(self, *args: Any, **kwargs: Any) -> Any: + """Never called by these tests.""" + raise NotImplementedError + + @staticmethod + def of(name: str) -> "_StubTool": + """Build a stub tool under the given name.""" + return _StubTool( + metadata=ToolMetadata(name=name, description="stub", args_schema=_EmptyArgs) + ) + + +class _StubSubagentSetup(SubagentSetup): + """Metadata-carrying sub-agent double: declaring it is all these tests exercise.""" + + async def submit( + self, + ctx: Any, + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> Any: + """Never invoked: these tests stop at declaration.""" + raise NotImplementedError + + +class _Review(BaseModel): + """Arguments of the typed double below.""" + + path: str + lines: int = 0 + + +class _TypedStubSubagentSetup(_StubSubagentSetup): + """Types its arguments, so the schema it is declared with is derived.""" + + @classmethod + def input_type(cls) -> type: + """Return the declared argument type.""" + return _Review + + +class _StubResourceContext(ResourceContext): + """Resource context backed by a plain name-keyed store. + + ``agent_store`` holds AGENT-type overrides, so one name can carry a tool + and a sub-agent at the same time; an AGENT lookup falls back to the plain + store when there is no override. + """ + + def __init__( + self, + store: Dict[str, Resource], + agent_store: Dict[str, Resource] | None = None, + ) -> None: + self._store = store + self._agent_store = agent_store or {} + + def get_resource(self, name: str, resource_type: ResourceType) -> Resource: + """Return the stored resource, honoring the AGENT overrides.""" + if resource_type == ResourceType.AGENT and name in self._agent_store: + return self._agent_store[name] + if name not in self._store: + msg = f"No such resource: {name}" + raise KeyError(msg) + return self._store[name] + + def generate_available_skills_prompt(self, *skill_names: str) -> str: + """Return a recognizable stand-in for the skill listing.""" + return f"{list(skill_names)}" + + def get_skill_dirs(self, *skill_names: str) -> List[str]: + """Return no skill directories.""" + return [] + + +def _build( + store: Dict[str, Resource], + agent_store: Dict[str, Resource] | None = None, + **setup_args: Any, +) -> tuple[_StubSetup, _RecordingConnection]: + connection = _RecordingConnection() + store["conn"] = connection + setup = _StubSetup( + connection="conn", + model="m", + resource_context=_StubResourceContext(store, agent_store), + **setup_args, + ) + return setup, connection + + +def test_declared_subagents_reach_the_model_as_callables_after_the_tools() -> None: + """Sub-agents are declared to the model after the plain tools.""" + store: Dict[str, Resource] = { + "lookup": _StubTool.of("lookup"), + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ), + } + setup, connection = _build(store, subagents=["reviewer"], tools=["lookup"]) + + setup.open() + setup.chat([]) + + assert [tool.metadata.name for tool in connection.captured_tools] == [ + "lookup", + "subagent_reviewer", + ] + delegated = connection.captured_tools[1] + assert isinstance(delegated, SubagentTool) + assert delegated.metadata.description == "Reviews a file. This is subagent." + assert delegated.metadata.get_parameters_dict()["properties"] == { + "path": {"title": "Path", "type": "string"} + } + + +def test_an_undescribed_subagent_is_still_delegable() -> None: + """Without a description, a generic delegation description is declared.""" + store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup(input_schema=CUSTOM_SCHEMA) + } + setup, _ = _build(store, subagents=["reviewer"]) + + setup.open() + + tool = setup.tools[0] + assert tool.metadata.name == "subagent_reviewer" + assert tool.metadata.description == ( + "Delegate a standalone task to sub-agent reviewer This is subagent." + ) + assert tool.metadata.get_parameters_dict()["properties"] == { + "path": {"title": "Path", "type": "string"} + } + + +def test_a_typed_subagent_is_declared_with_the_derived_schema() -> None: + """A declared argument type is rendered, so it need not be written out.""" + store: Dict[str, Resource] = { + "reviewer": _TypedStubSubagentSetup(description="Reviews a file.") + } + setup, _ = _build(store, subagents=["reviewer"]) + + setup.open() + + assert setup.tools[0].metadata.get_parameters_dict()["properties"] == { + "path": {"title": "Path", "type": "string"}, + "lines": {"default": 0, "title": "Lines", "type": "integer"}, + } + + +def test_a_subagent_that_states_no_input_shape_is_not_offered_to_the_model() -> None: + """A sub-agent that states no shape leaves the model nothing to build a call + from, so it is dropped rather than declared as a callable it could only + misuse. + """ + store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup(description="Reviews a file.") + } + setup, _ = _build(store, subagents=["reviewer"]) + + setup.open() + + assert setup.tools == [] + + +def test_a_subagent_without_an_input_shape_does_not_stop_the_others() -> None: + """Dropping one is not dropping the rest: the other callables stay usable.""" + store: Dict[str, Resource] = { + "opaque": _StubSubagentSetup(description="Reviews a file."), + "coder": _StubSubagentSetup( + description="Writes a patch.", input_schema=CUSTOM_SCHEMA + ), + } + setup, _ = _build(store, subagents=["opaque", "coder"]) + + setup.open() + + assert [tool.metadata.name for tool in setup.tools] == ["subagent_coder"] + + +def test_a_tool_and_a_subagent_may_share_a_name() -> None: + """The reserved prefix keeps the two namespaces apart, so one name may serve + both. + """ + store: Dict[str, Resource] = {"reviewer": _StubTool.of("reviewer")} + agent_store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ) + } + setup, _ = _build(store, agent_store, tools=["reviewer"], subagents=["reviewer"]) + + setup.open() + + assert [tool.metadata.name for tool in setup.tools] == [ + "reviewer", + "subagent_reviewer", + ] + + +def test_a_repeated_tool_name_is_rejected() -> None: + """The same tool declared twice would be declared twice to the model.""" + store: Dict[str, Resource] = {"lookup": _StubTool.of("lookup")} + setup, _ = _build(store, tools=["lookup", "lookup"]) + + with pytest.raises(ValueError, match="Duplicate callable name: lookup"): + setup.open() + + +def test_a_repeated_subagent_name_is_rejected() -> None: + """The same sub-agent declared twice would be declared twice to the model.""" + store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ) + } + setup, _ = _build(store, subagents=["reviewer", "reviewer"]) + + with pytest.raises(ValueError, match="Duplicate callable name: subagent_reviewer"): + setup.open() + + +def test_a_subagent_without_a_setup_is_rejected() -> None: + """An AGENT resource that is not a SubagentSetup carries no schema to declare. + + A sub-agent owned by the other language resolves to a bridge handle here, so + declaring it must fail loudly rather than leave the model with a callable it + cannot call. + """ + store: Dict[str, Resource] = {"reviewer": _StubTool.of("reviewer")} + setup, _ = _build(store, subagents=["reviewer"]) + + with pytest.raises(TypeError, match="must resolve to a SubagentSetup"): + setup.open() + + +def test_reopening_does_not_duplicate_the_callables() -> None: + """A second open() rebuilds the callables instead of appending to them.""" + store: Dict[str, Resource] = { + "lookup": _StubTool.of("lookup"), + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ), + } + setup, _ = _build(store, tools=["lookup"], subagents=["reviewer"]) + + setup.open() + setup.open() + + assert len(setup.tools) == 2 + + +def test_declared_subagents_inject_no_listing_message() -> None: + """Delegation is offered through the callables alone: no listing message is + injected. + """ + store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ), + "coder": _StubSubagentSetup( + description="Writes a patch.", input_schema=CUSTOM_SCHEMA + ), + } + setup, connection = _build(store, subagents=["reviewer", "coder"]) + setup.open() + + setup.chat( + [ + ChatMessage(role=MessageRole.SYSTEM, content="You are helpful."), + ChatMessage(role=MessageRole.USER, content="review it"), + ] + ) + + assert len(connection.captured_messages) == 2 + assert connection.captured_messages[0].content == "You are helpful." + assert len(connection.captured_tools) == 2 + + +def test_only_the_skill_listing_is_injected_when_subagents_are_declared() -> None: + """With sub-agents declared, the skill listing is still the only injection.""" + store: Dict[str, Resource] = { + "reviewer": _StubSubagentSetup( + description="Reviews a file.", input_schema=CUSTOM_SCHEMA + ), + "load_skill": _StubTool.of("load_skill"), + "bash": _StubTool.of("bash"), + } + setup, connection = _build(store, subagents=["reviewer"], skills=["github"]) + setup.open() + + setup.chat([ChatMessage(role=MessageRole.USER, content="review it")]) + + assert len(connection.captured_messages) == 2 + assert connection.captured_messages[0].content.startswith("") + assert connection.captured_messages[1].content == "review it" + + +def test_a_setup_without_subagents_injects_nothing() -> None: + """Without declared sub-agents, neither the tools nor the messages change.""" + setup, connection = _build({}) + setup.open() + + setup.chat([ChatMessage(role=MessageRole.USER, content="hi")]) + + assert len(connection.captured_messages) == 1 + assert connection.captured_tools == [] diff --git a/python/flink_agents/api/subagent.py b/python/flink_agents/api/subagent.py index 0b15a0557..89b5229cd 100644 --- a/python/flink_agents/api/subagent.py +++ b/python/flink_agents/api/subagent.py @@ -15,11 +15,14 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +import json import logging from abc import ABC, abstractmethod from dataclasses import dataclass from typing import TYPE_CHECKING, Any +from pydantic import BaseModel, field_validator, model_validator + from flink_agents.api.resource import ResourceType, SerializableResource if TYPE_CHECKING: @@ -27,6 +30,53 @@ _LOG = logging.getLogger(__name__) +# Prefix of the callable name a sub-agent is exposed to a chat model under. +# Tools are forbidden to register under this prefix, so a prefixed callable +# name unambiguously addresses a sub-agent and the executing side routes it to +# the AGENT namespace. +CALLABLE_NAME_PREFIX = "subagent_" + + +def _input_schema_from(input_type: type) -> str | None: + """Render the schema of ``input_type``, or None when it states no shape. + + ``object`` is what a sub-agent declares when it declares none, and it has no + fields for a model to build a call from. A type that is not a pydantic model + is refused rather than dropped: it is a declaration mistake, and the caller + has two ways to fix it. + + Args: + input_type: The declared type of the arguments. + + Returns: + The schema as a JSON string, the form the cross-language plan carries, + or None when the type states no shape. + + Raises: + TypeError: If ``input_type`` is not a pydantic model, or has no JSON + Schema. + """ + if input_type is object: + return None + if not (isinstance(input_type, type) and issubclass(input_type, BaseModel)): + msg = ( + f"Sub-agent input type {input_type} is not supported. Declare a" + " pydantic BaseModel, or declare an input schema explicitly." + ) + raise TypeError(msg) + try: + schema = input_type.model_json_schema() + except Exception as e: + msg = ( + f"Sub-agent input type {input_type.__module__}" + f".{input_type.__qualname__} cannot be rendered as a JSON Schema, so" + " it cannot be declared to a chat model. Declare an input schema" + " explicitly, or use an input type whose fields are all" + f" JSON-Schema-renderable. Rendering it reported: {e}" + ) + raise TypeError(msg) from e + return json.dumps(schema) + @dataclass class SubagentResult: @@ -152,7 +202,61 @@ def __await__(self) -> Any: class SubagentSetup(SerializableResource): - """Caller-facing definition of a sub-agent, registered as an AGENT resource.""" + """Caller-facing definition of a sub-agent, registered as an AGENT resource. + + Attributes: + ---------- + description : str + Tells a caller what this sub-agent is for, so that it can decide + whether to delegate to it. This is routing information for the caller, + not an instruction for the sub-agent itself. + input_schema : str | None + JSON Schema of the arguments this sub-agent accepts, as declared + explicitly. Derived from :meth:`input_type` when it is not, and left + ``None`` when neither says anything a model could build a call from. + """ + + description: str = "" + input_schema: str | None = None + + @field_validator("input_schema") + @classmethod + def __check_input_schema(cls, input_schema: str | None) -> str | None: + if input_schema is not None and not input_schema.strip(): + msg = "Sub-agent input schema must not be blank." + raise ValueError(msg) + return input_schema + + @model_validator(mode="after") + def __derive_input_schema(self) -> "SubagentSetup": + """Fill an undeclared schema from the input type. + + Filled at construction rather than read on demand, so that what a caller + reads off this setup and what the plan carries are the same schema. + """ + if self.input_schema is None: + self.input_schema = _input_schema_from(type(self).input_type()) + return self + + @classmethod + def input_type(cls) -> type: + """The type of the arguments this sub-agent accepts. + + Override to type the arguments instead of spelling out their schema; + :class:`object`, the default, states no shape, which leaves the + sub-agent undeclared to a chat model unless an input schema is given. + """ + return object + + @classmethod + def result_type(cls) -> type: + """The type the result is converted to before it is reported back. + + Override to state the shape of a result that is not already + JSON-compatible, or to narrow a wider one; :class:`object`, the + default, reports the result as it arrived. + """ + return object @classmethod def resource_type(cls) -> ResourceType: diff --git a/python/flink_agents/api/tests/subagent_test_utils.py b/python/flink_agents/api/tests/subagent_test_utils.py index ee1b25897..3e40e1704 100644 --- a/python/flink_agents/api/tests/subagent_test_utils.py +++ b/python/flink_agents/api/tests/subagent_test_utils.py @@ -19,6 +19,8 @@ from typing import TYPE_CHECKING, Any +from pydantic import BaseModel + from flink_agents.api.subagent import SubagentSetup if TYPE_CHECKING: @@ -26,6 +28,20 @@ from flink_agents.api.subagent import SubagentFuture +class Review(BaseModel): + """Arguments of the typed double below.""" + + path: str + lines: int = 0 + + +class Verdict(BaseModel): + """Result of the typed double below.""" + + approved: bool + note: str = "" + + class TestSubagentSetup(SubagentSetup): """Shared ``SubagentSetup`` test double, constructible directly or from a resource descriptor (the YAML shape). @@ -47,3 +63,17 @@ def submit( """Descriptor-only double; invocation lives in the runtime layer.""" msg = "Descriptor-only sub-agent setup; invocation lives in the runtime layer." raise NotImplementedError(msg) + + +class TypedTestSubagentSetup(TestSubagentSetup): + """Types its arguments and its result instead of spelling out a schema.""" + + @classmethod + def input_type(cls) -> type: + """Return the declared argument type.""" + return Review + + @classmethod + def result_type(cls) -> type: + """Return the declared result type.""" + return Verdict diff --git a/python/flink_agents/api/tests/test_subagent.py b/python/flink_agents/api/tests/test_subagent.py index a6cebfa8b..1b67352f0 100644 --- a/python/flink_agents/api/tests/test_subagent.py +++ b/python/flink_agents/api/tests/test_subagent.py @@ -16,11 +16,20 @@ # limitations under the License. ################################################################################ """Tests registering sub-agents as AGENT resources.""" + +import json + import pytest +from pydantic import ValidationError from flink_agents.api.agents.agent import Agent from flink_agents.api.resource import ResourceType -from flink_agents.api.tests.subagent_test_utils import TestSubagentSetup +from flink_agents.api.tests.subagent_test_utils import ( + TestSubagentSetup, + TypedTestSubagentSetup, +) + +DECLARED_SCHEMA = '{"type":"object","properties":{"path":{"type":"string"}}}' def test_register_subagent_setup_as_resource() -> None: @@ -58,3 +67,85 @@ def test_multiple_subagents_registered() -> None: assert len(agent_resources) == 2 assert agent_resources["reviewer"] is reviewer assert agent_resources["coder"] is coder + + +def test_a_setup_that_declares_nothing_states_no_shape_for_its_arguments() -> None: + """Without declared metadata, a sub-agent states no shape for its arguments.""" + setup = TestSubagentSetup() + + assert setup.description == "" + assert setup.input_type() is object + assert setup.result_type() is object + assert setup.input_schema is None + + +def test_declared_metadata_is_kept_verbatim() -> None: + """A declared description and input schema reach the caller unchanged.""" + setup = TestSubagentSetup( + description="Reviews a file.", input_schema=DECLARED_SCHEMA + ) + + assert setup.description == "Reviews a file." + assert setup.input_schema == DECLARED_SCHEMA + + +def test_a_blank_input_schema_is_rejected() -> None: + """A blank input schema would leave the caller with nothing to fill in.""" + with pytest.raises(ValidationError, match="input schema must not be blank"): + TestSubagentSetup(input_schema=" ") + + +def test_an_input_type_is_rendered_as_the_input_schema() -> None: + """A declared argument type is rendered, so it need not be written out.""" + schema = json.loads(TypedTestSubagentSetup().input_schema) + + assert schema["type"] == "object" + assert schema["properties"]["path"]["type"] == "string" + assert schema["properties"]["lines"]["type"] == "integer" + + +def test_an_explicit_input_schema_wins_over_the_input_type() -> None: + """What is written out is what gets declared, even next to a type.""" + setup = TypedTestSubagentSetup(input_schema=DECLARED_SCHEMA) + + assert setup.input_schema == DECLARED_SCHEMA + + +def test_an_input_type_that_is_not_a_model_is_rejected() -> None: + """Only a pydantic model states a shape a schema can be rendered from.""" + + class NotAModel(TestSubagentSetup): + @classmethod + def input_type(cls) -> type: + return dict + + with pytest.raises(TypeError, match="is not supported"): + NotAModel() + + +def test_the_declared_types_stay_out_of_the_plan_json() -> None: + """The declared types drive behavior, so they must not leak into the plan.""" + dumped = json.loads(TypedTestSubagentSetup().model_dump_json()) + + assert "input_type" not in dumped + assert "result_type" not in dumped + + +def test_a_derived_input_schema_is_carried_by_the_plan_json() -> None: + """The plan carries the schema a caller reads off the setup, derived or not.""" + dumped = json.loads(TypedTestSubagentSetup().model_dump_json()) + + assert json.loads(dumped["input_schema"])["type"] == "object" + + +def test_the_metadata_serializes_under_the_cross_language_keys() -> None: + """The plan JSON keys are a contract with the Java side, so they are pinned.""" + dumped = json.loads( + TestSubagentSetup( + description="Reviews a file.", input_schema=DECLARED_SCHEMA + ).model_dump_json() + ) + + assert dumped["description"] == "Reviews a file." + assert dumped["input_schema"] == DECLARED_SCHEMA + assert "inputSchema" not in dumped diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index 4952bfd8d..a182ba0a6 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -25,6 +25,11 @@ from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType from flink_agents.api.runner_context import DurableCall, Outcome, RunnerContext +from flink_agents.api.subagent import ( + CALLABLE_NAME_PREFIX, + SubagentResult, + SubagentSetup, +) from flink_agents.api.tools import ToolExecutionMetadataProvider, ToolResponse from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, @@ -37,6 +42,10 @@ ToolExecutionMetadataKeys, ) from flink_agents.plan.actions.action import Action +from flink_agents.plan.actions.tool_result_utils import ( + normalize_agent_result, + to_chat_message_content, +) from flink_agents.plan.function import PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool @@ -82,8 +91,10 @@ def _tool_entity_metadata( class _ToolCallExecution: id: str name: str - durable_call: DurableCall + durable_call: DurableCall | None entity_metadata: dict[str, Any] + agent: SubagentSetup | None = None + agent_kwargs: dict[str, Any] | None = None async def process_tool_request(event: Event, ctx: RunnerContext) -> None: @@ -150,11 +161,20 @@ def _build_tool_call_executions( call_kwargs = dict(kwargs or {}) tool = None + agent = None preparation_error = None + # The reserved subagent_ prefix separates the two namespaces: tool registration + # rejects the prefix, so a prefixed callable name can only address a sub-agent. + # Matched once here and carried down, because resolving the AGENT resource + # raises when the name is absent and would otherwise have to be attempted for + # every plain tool call. + delegated = name.startswith(CALLABLE_NAME_PREFIX) try: - tool = ctx.get_resource(name, ResourceType.TOOL) + tool, agent = _resolve_callable(name, ctx, delegated=delegated) except Exception as e: preparation_error = e + # Injection is a tool-only contract, so a sub-agent call carries the model + # arguments unchanged. if tool is not None: try: # Framework-owned injected args must win over model-provided values so @@ -170,14 +190,13 @@ def _build_tool_call_executions( ctx, ExecutionEntityTypes.TOOL, name, entity_metadata ) - if not tool or preparation_error is not None: + unresolved = tool is None and agent is None + if unresolved or preparation_error is not None: failure = preparation_error or RuntimeError( f"Tool `{name}` does not exist." ) - responses[call_id] = ( - f"Tool `{name}` does not exist." - if not tool - else f"Tool `{name}` execute failed." + responses[call_id] = _preparation_failure_message( + name, failure, delegated=delegated, unresolved=unresolved ) success[call_id] = False error[call_id] = str(failure) @@ -191,17 +210,32 @@ def _build_tool_call_executions( ) continue - executions.append( - _ToolCallExecution( - id=call_id, - name=name, - durable_call=DurableCall( - func=tool.call, - kwargs=call_kwargs, - ), - entity_metadata=entity_metadata, + if agent is not None: + # A sub-agent call cannot join the durable batch: submit() and awaiting the + # handle already run through durable execution inside the setup, so wrapping + # it again here would nest durable cursors. + executions.append( + _ToolCallExecution( + id=call_id, + name=name, + durable_call=None, + entity_metadata=entity_metadata, + agent=agent, + agent_kwargs=call_kwargs, + ) + ) + else: + executions.append( + _ToolCallExecution( + id=call_id, + name=name, + durable_call=DurableCall( + func=tool.call, + kwargs=call_kwargs, + ), + entity_metadata=entity_metadata, + ) ) - ) return executions @@ -212,14 +246,20 @@ async def _execute_parallel( success: dict, error: dict, ) -> None: + tool_executions = [] + for execution in executions: + if execution.agent is not None: + await _dispatch_agent_execution(execution, ctx, responses, success, error) + else: + tool_executions.append(execution) try: outcomes = await ctx.durable_execute_all_async( - [execution.durable_call for execution in executions] + [execution.durable_call for execution in tool_executions] ) - for execution, outcome in zip(executions, outcomes, strict=True): + for execution, outcome in zip(tool_executions, outcomes, strict=True): _record_outcome(execution, outcome, ctx, responses, success, error) except Exception as e: - for execution in executions: + for execution in tool_executions: _record_execution_exception(execution, e, ctx, responses, success, error) @@ -233,6 +273,9 @@ async def _execute_sequentially( error: dict, ) -> None: for execution in executions: + if execution.agent is not None: + await _dispatch_agent_execution(execution, ctx, responses, success, error) + continue try: call = execution.durable_call if tool_call_async: @@ -248,7 +291,7 @@ async def _execute_sequentially( **(call.kwargs or {}), ) _record_tool_response(execution, response, ctx, responses, success, error) - except Exception as e: # noqa: PERF203 + except Exception as e: _record_execution_exception(execution, e, ctx, responses, success, error) @@ -302,6 +345,61 @@ def _record_tool_response( ) +async def _dispatch_agent_execution( + execution: _ToolCallExecution, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + try: + # submit() and awaiting the handle already run through durable execution inside + # the setup, so wrapping the call again here would nest durable cursors. + future = await execution.agent.submit(ctx, execution.agent_kwargs) + result = await future + _record_agent_result(execution, result, ctx, responses, success, error) + except Exception as e: + _record_execution_exception(execution, e, ctx, responses, success, error) + + +def _record_agent_result( + execution: _ToolCallExecution, + result: SubagentResult, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + if result.success: + responses[execution.id] = to_chat_message_content( + normalize_agent_result(result.result, execution.agent.result_type()) + ) + success[execution.id] = True + ExecutionReporters.succeeded( + ctx, + ExecutionEntityTypes.TOOL, + execution.name, + execution.entity_metadata, + ) + else: + # The model sees why the delegation failed, so it can correct the call + # instead of repeating it blindly; the error map keeps the same detail + # for observability. + responses[execution.id] = _with_reason( + f"Sub-agent `{execution.name}` execute failed", result.error_message + ) + success[execution.id] = False + error[execution.id] = result.error_message + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.TOOL, + execution.name, + execution.entity_metadata, + result.exception, + ExecutionProblemCategories.TOOL_CALL_FAILED, + ) + + def _record_execution_exception( execution: _ToolCallExecution, exception: BaseException, @@ -310,7 +408,11 @@ def _record_execution_exception( success: dict, error: dict, ) -> None: - responses[execution.id] = f"Tool `{execution.name}` execute failed." + responses[execution.id] = ( + _with_reason(f"Sub-agent `{execution.name}` execute failed", str(exception)) + if execution.agent is not None + else f"Tool `{execution.name}` execute failed." + ) success[execution.id] = False error[execution.id] = str(exception) ExecutionReporters.failed( @@ -323,6 +425,57 @@ def _record_execution_exception( ) +def _resolve_callable( + name: str, ctx: RunnerContext, *, delegated: bool +) -> tuple[Any | None, SubagentSetup | None]: + """Resolve a callable name the model emitted to either a tool or a sub-agent. + + ``delegated`` is the namespace the caller decided on, matched once there + rather than again here. + """ + if delegated: + return None, _resolve_subagent(name[len(CALLABLE_NAME_PREFIX) :], ctx) + return ctx.get_resource(name, ResourceType.TOOL), None + + +def _preparation_failure_message( + name: str, failure: BaseException, *, delegated: bool, unresolved: bool +) -> str: + """The message the model sees when a call could not even be prepared. + + A rejected sub-agent call carries the reason, so the model can correct the + call instead of repeating it blindly. ``delegated`` is the namespace the + caller decided on, matched once there rather than again here. + """ + if delegated: + return _with_reason(f"Sub-agent `{name}` execute failed", str(failure)) + return ( + f"Tool `{name}` does not exist." + if unresolved + else f"Tool `{name}` execute failed." + ) + + +def _with_reason(message: str, reason: str | None) -> str: + return f"{message}: {reason}" if reason else f"{message}." + + +def _resolve_subagent(name: str, ctx: RunnerContext) -> SubagentSetup: + """Resolve a sub-agent, in one lookup: the AGENT resource is fetched once and + checked once here, and the caller carries the setup from then on. + """ + setup = ctx.get_resource(name, ResourceType.AGENT) + if not isinstance(setup, SubagentSetup): + # A sub-agent owned by the other language resolves to a bridge handle here, + # which cannot be called through this path. + msg = ( + f"Sub-agent {name} must resolve to a SubagentSetup, " + f"but was {type(setup).__name__}." + ) + raise TypeError(msg) + return setup + + def _resolve_injected_arguments(tool: object, ctx: RunnerContext) -> dict: if not isinstance(tool, FunctionTool): return {} diff --git a/python/flink_agents/plan/actions/tool_result_utils.py b/python/flink_agents/plan/actions/tool_result_utils.py new file mode 100644 index 000000000..6ebe0ab77 --- /dev/null +++ b/python/flink_agents/plan/actions/tool_result_utils.py @@ -0,0 +1,117 @@ +################################################################################ +# 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. +################################################################################# +"""Turns a sub-agent result into something a chat model can be told. + +A sub-agent result reaches the caller as an opaque object, but from here on it is +carried in a tool message and re-bound after a failover, so it must hold nothing +that JSON cannot express. Such a payload is rejected with the path where it was +found instead of being dropped or silently stringified. + +A sub-agent that declares a result type is read through it: the type says how to +interpret a result the checks below would refuse, and what comes out is only what +the type declares. +""" + +import json +import math +from functools import cache +from typing import Any + +from pydantic import TypeAdapter + + +def normalize_agent_result(raw: Any, result_type: type = object) -> Any: + """Reject a payload JSON cannot express, and reduce the rest to plain containers. + + Parameters + ---------- + raw : Any + The result a sub-agent produced. + result_type : type = object + The type the sub-agent declares for its result, or :class:`object` when + it declares none, in which case ``raw`` is taken as it arrived. + + Returns: + ------- + Any + The result in generic form: dicts, lists and scalars only. + """ + if result_type is object: + _require_json_compatible(raw, "result") + return json.loads(json.dumps(raw, allow_nan=False)) + # Two conversions: the first reads the result as the declared type, which is + # what admits a payload the compatibility check below would refuse; the + # second reduces that back to plain containers, so what is reported and + # re-bound after a failover is the same generic form the undeclared path + # produces. + adapter = _adapter_for(result_type) + generic = adapter.dump_python(adapter.validate_python(raw), mode="json") + # Still required: a declared type can render a field JSON cannot express, and + # the result outlives this call in a tool message and in state. + _require_json_compatible(generic, "result") + return generic + + +@cache +def _adapter_for(result_type: type) -> TypeAdapter[Any]: + """The adapter of a declared result type, built once per type. + + Building one validates the type against pydantic, which is far more than a + single result report should cost, and the set of declared types is bounded by + the sub-agents in the plan. + """ + return TypeAdapter(result_type) + + +def to_chat_message_content(value: Any) -> str: + """Render a value as the content of the tool message handed back to the model.""" + if value is None: + return "null" + if isinstance(value, bool | dict | list | tuple): + return json.dumps(value, separators=(",", ":"), allow_nan=False) + return str(value) + + +def _require_json_compatible(value: Any, path: str) -> None: + if value is None or isinstance(value, str | bool): + return + if isinstance(value, int | float): + _require_finite_number(value, path) + return + if isinstance(value, dict): + for key, item in value.items(): + if not isinstance(key, str): + msg = f"Map keys in sub-agent result must be strings at {path}" + raise TypeError(msg) + _require_json_compatible(item, f"{path}.{key}") + return + if isinstance(value, list | tuple): + for index, item in enumerate(value): + _require_json_compatible(item, f"{path}[{index}]") + return + msg = ( + f"Sub-agent result must be JSON-compatible at {path}, " + f"found {type(value).__name__}" + ) + raise TypeError(msg) + + +def _require_finite_number(number: float, path: str) -> None: + if isinstance(number, float) and not math.isfinite(number): + msg = f"Non-finite number in sub-agent result at {path}" + raise ValueError(msg) diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py index 5a56c5dc5..ab5a193c7 100644 --- a/python/flink_agents/plan/agent_plan.py +++ b/python/flink_agents/plan/agent_plan.py @@ -33,7 +33,7 @@ LOAD_SKILL_TOOL, Skills, ) -from flink_agents.api.subagent import SubagentSetup +from flink_agents.api.subagent import CALLABLE_NAME_PREFIX, SubagentSetup from flink_agents.api.tools.function_tool import FunctionTool as ApiFunctionTool from flink_agents.api.tools.tool import Tool from flink_agents.plan.actions.action import Action @@ -293,6 +293,22 @@ def _to_plan_function(func: ApiFunction) -> PythonFunction | JavaFunction: raise TypeError(msg) +def _check_tool_name_not_reserved(name: str) -> None: + """Reject a tool name carrying the reserved sub-agent callable prefix. + + Sub-agent callables are exposed to the model under the ``subagent_`` prefix, + and dispatch routes any prefixed call to the AGENT namespace, so a tool + registered under the prefix could never be called. Fail clearly at + plan-construction time rather than at call time. + """ + if name.startswith(CALLABLE_NAME_PREFIX): + msg = ( + f"Tool name '{name}' must not start with the reserved prefix " + f"'{CALLABLE_NAME_PREFIX}', which identifies sub-agent callables." + ) + raise ValueError(msg) + + def _get_resource_providers( agent: Agent, config: AgentConfiguration ) -> List[ResourceProvider]: @@ -322,6 +338,7 @@ def _get_resource_providers( ) elif hasattr(value, "_is_tool"): + _check_tool_name_not_reserved(name) injected_args = getattr(value, "_injected_args", None) if isinstance(value, staticmethod): value = value.__func__ @@ -366,6 +383,7 @@ def _get_resource_providers( ) for name, tool in agent.resources[ResourceType.TOOL].items(): + _check_tool_name_not_reserved(name) resource_providers.append( PythonSerializableResourceProvider.from_resource( name=name, diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action_subagent.py b/python/flink_agents/plan/tests/actions/test_tool_call_action_subagent.py new file mode 100644 index 000000000..f8ce27212 --- /dev/null +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action_subagent.py @@ -0,0 +1,336 @@ +################################################################################ +# 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. +################################################################################# +import asyncio +from typing import Any + +from pydantic import BaseModel, PrivateAttr +from typing_extensions import override + +from flink_agents.api.core_options import AgentExecutionOptions +from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent +from flink_agents.api.resource import ResourceType +from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.subagent import SubagentFuture, SubagentResult, SubagentSetup +from flink_agents.plan.actions.tool_call_action import process_tool_request +from flink_agents.plan.configuration import AgentConfiguration +from flink_agents.plan.function import PythonFunction +from flink_agents.plan.tools.function_tool import FunctionTool + + +def query_order(order_id: str) -> str: + return f"queried {order_id}" + + +class _ResolvedSubagentFuture(SubagentFuture): + """Handle that is already resolved to a preset outcome.""" + + def __init__(self, session_id: str, call_id: str, outcome: SubagentResult) -> None: + super().__init__(session_id, call_id) + self._outcome = outcome + + @override + def done(self) -> bool: + return True + + @override + def combine(self, *others: SubagentFuture) -> Any: + raise NotImplementedError + + @override + def __await__(self) -> Any: + async def resolve() -> SubagentResult: + return self._outcome + + return resolve().__await__() + + +class _RecordingSubagentSetup(SubagentSetup): + """Captures every prompt it is handed and resolves to a preset outcome.""" + + _outcome: SubagentResult = PrivateAttr(default=None) + _submit_failure: Exception | None = PrivateAttr(default=None) + _prompts: list[Any] = PrivateAttr(default_factory=list) + + @classmethod + def of( + cls, outcome: SubagentResult, submit_failure: Exception | None = None + ) -> "_RecordingSubagentSetup": + setup = cls(description="Reviews a diff.") + setup._outcome = outcome + setup._submit_failure = submit_failure + return setup + + @property + def prompts(self) -> list[Any]: + """Every prompt handed to this sub-agent, in call order.""" + return self._prompts + + @override + async def submit( + self, + ctx: RunnerContext, + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> SubagentFuture: + if self._submit_failure is not None: + raise self._submit_failure + self._prompts.append(prompt) + return _ResolvedSubagentFuture( + session_id or "session", call_id or "call", self._outcome + ) + + +class _Verdict(BaseModel): + """The result the typed double below declares.""" + + approved: bool + note: str = "" + + +class _TypedRecordingSubagentSetup(_RecordingSubagentSetup): + """Declares a result type, so its result is read through it.""" + + @classmethod + def result_type(cls) -> type: + """Return the declared result type.""" + return _Verdict + + +class _Context: + def __init__(self) -> None: + self.config = AgentConfiguration({}) + self.config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, False) + self.sensory_memory = None + self.short_term_memory = None + self.sent_events = [] + self.tools = {} + self.agents = {} + self.durable_executions = 0 + + def with_tool(self, name: str, tool: Any) -> "_Context": + self.tools[name] = tool + return self + + def with_agent(self, name: str, agent: Any) -> "_Context": + self.agents[name] = agent + return self + + def get_resource(self, name: str, type: ResourceType) -> Any: + registry = self.agents if type == ResourceType.AGENT else self.tools + if name not in registry: + msg = f"Resource does not exist: {name}" + raise ValueError(msg) + return registry[name] + + def durable_execute(self, func: Any, **kwargs: Any) -> Any: + self.durable_executions += 1 + return func(**kwargs) + + async def durable_execute_async(self, func: Any, **kwargs: Any) -> Any: + self.durable_executions += 1 + return func(**kwargs) + + def send_event(self, event: Any) -> None: + self.sent_events.append(event) + + +def tool_request(callable_name: str) -> ToolRequestEvent: + return ToolRequestEvent( + model="model", + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": { + "name": callable_name, + "arguments": {"prompt": "review the diff"}, + }, + } + ], + ) + + +def order_tool() -> FunctionTool: + return FunctionTool(func=PythonFunction.from_callable(query_order)) + + +def test_delegates_to_the_subagent_and_reports_its_normalized_result() -> None: + agent = _RecordingSubagentSetup.of( + SubagentResult.ok({"verdict": "approved", "findings": ["style"]}) + ) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is True + assert response.responses["call-1"] == '{"verdict":"approved","findings":["style"]}' + assert "call-1" not in response.error + + +def test_hands_the_model_arguments_to_the_subagent_as_the_prompt() -> None: + agent = _RecordingSubagentSetup.of(SubagentResult.ok("done")) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + assert agent.prompts == [{"prompt": "review the diff"}] + # A sub-agent call resolves through the setup, which owns its own durable + # execution. + assert ctx.durable_executions == 0 + + +def test_reports_a_failed_subagent_result_with_the_detail_exposed() -> None: + agent = _RecordingSubagentSetup.of(SubagentResult.error("upstream refused")) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is False + assert response.responses["call-1"] == ( + "Sub-agent `subagent_reviewer` execute failed: upstream refused" + ) + assert response.error["call-1"] == "upstream refused" + + +def test_reports_a_failure_raised_while_submitting() -> None: + agent = _RecordingSubagentSetup.of( + SubagentResult.ok("unreachable"), RuntimeError("mailbox is full") + ) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is False + assert response.responses["call-1"] == ( + "Sub-agent `subagent_reviewer` execute failed: mailbox is full" + ) + assert response.error["call-1"] == "mailbox is full" + + +def test_rejects_a_result_json_cannot_express() -> None: + agent = _RecordingSubagentSetup.of(SubagentResult.ok({"handle": object()})) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is False + assert response.responses["call-1"].startswith( + "Sub-agent `subagent_reviewer` execute failed" + ) + assert "result.handle" in response.responses["call-1"] + assert "result.handle" in response.error["call-1"] + + +def test_reads_a_result_through_the_type_the_subagent_declares() -> None: + """A declared result type is what admits a result JSON cannot express on its + own. + """ + agent = _TypedRecordingSubagentSetup.of( + SubagentResult.ok(_Verdict(approved=True, note="clean")) + ) + ctx = _Context().with_agent("reviewer", agent) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is True + assert response.responses["call-1"] == '{"approved":true,"note":"clean"}' + + +def test_routes_a_tool_and_a_subagent_sharing_a_name_to_their_own_namespace() -> None: + """The reserved prefix routes each namespace on its own, even under one + shared name. + """ + ctx = ( + _Context() + .with_agent("reviewer", _RecordingSubagentSetup.of(SubagentResult.ok("done"))) + .with_tool("reviewer", order_tool()) + ) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + delegated = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert delegated.success["call-1"] is True + assert delegated.responses["call-1"] == "done" + # A sub-agent call resolves through the setup, which owns its own durable + # execution. + assert ctx.durable_executions == 0 + + tool_event = ToolRequestEvent( + model="model", + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "reviewer", "arguments": {"order_id": "order-1"}}, + } + ], + ) + asyncio.run(process_tool_request(tool_event, ctx)) + direct = ToolResponseEvent.from_event(ctx.sent_events[1]) + assert direct.success["call-1"] is True + assert direct.responses["call-1"] == "queried order-1" + assert ctx.durable_executions == 1 + + +def test_refuses_an_agent_resource_that_carries_no_callable_setup() -> None: + ctx = _Context().with_agent("reviewer", order_tool()) + + asyncio.run(process_tool_request(tool_request("subagent_reviewer"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is False + assert response.responses["call-1"] == ( + "Sub-agent `subagent_reviewer` execute failed: Sub-agent reviewer must" + " resolve to a SubagentSetup, but was FunctionTool." + ) + assert response.error["call-1"] == ( + "Sub-agent reviewer must resolve to a SubagentSetup, but was FunctionTool." + ) + + +def test_still_dispatches_a_tool_when_both_kinds_are_registered() -> None: + ctx = ( + _Context() + .with_agent("reviewer", _RecordingSubagentSetup.of(SubagentResult.ok("done"))) + .with_tool("query_order", order_tool()) + ) + event = ToolRequestEvent( + model="model", + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": { + "name": "query_order", + "arguments": {"order_id": "order-1"}, + }, + } + ], + ) + + asyncio.run(process_tool_request(event, ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is True + assert response.responses["call-1"] == "queried order-1" + assert ctx.durable_executions == 1 diff --git a/python/flink_agents/plan/tests/actions/test_tool_result_utils.py b/python/flink_agents/plan/tests/actions/test_tool_result_utils.py new file mode 100644 index 000000000..00a416bf1 --- /dev/null +++ b/python/flink_agents/plan/tests/actions/test_tool_result_utils.py @@ -0,0 +1,140 @@ +################################################################################ +# 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. +################################################################################# +import pytest +from pydantic import BaseModel + +from flink_agents.plan.actions.tool_result_utils import ( + normalize_agent_result, + to_chat_message_content, +) + + +class Verdict(BaseModel): + """The result a sub-agent that declares a result type produces.""" + + approved: bool + note: str = "" + score: float = 0.0 + + +def test_normalize_reduces_nested_containers_to_plain_ones() -> None: + raw = {"outer": {"inner": [1, "two"]}, "items": (3, 4)} + + assert normalize_agent_result(raw) == { + "outer": {"inner": [1, "two"]}, + "items": [3, 4], + } + + +def test_normalize_keeps_scalars_and_none() -> None: + assert normalize_agent_result("text") == "text" + assert normalize_agent_result(7) == 7 + assert normalize_agent_result(True) is True + assert normalize_agent_result(None) is None + + +def test_normalize_reports_where_an_unsupported_value_was_found() -> None: + with pytest.raises(TypeError) as error: + normalize_agent_result({"outer": [object()]}) + + assert "result.outer[0]" in str(error.value) + assert "found object" in str(error.value) + + +def test_normalize_rejects_a_non_string_mapping_key() -> None: + with pytest.raises(TypeError) as error: + normalize_agent_result({1: "a"}) + + assert str(error.value) == "Map keys in sub-agent result must be strings at result" + + +def test_normalize_rejects_a_non_finite_number() -> None: + with pytest.raises(ValueError) as error: + normalize_agent_result({"ratio": float("inf")}) + + assert str(error.value) == "Non-finite number in sub-agent result at result.ratio" + + with pytest.raises(ValueError) as error: + normalize_agent_result(float("nan")) + + assert str(error.value) == "Non-finite number in sub-agent result at result" + + +def test_normalize_walks_a_top_level_sequence_by_index() -> None: + assert normalize_agent_result([1, 2]) == [1, 2] + + with pytest.raises(TypeError) as error: + normalize_agent_result([object()]) + + assert "result[0]" in str(error.value) + + +def test_normalize_reads_a_declared_result_type_into_generic_form() -> None: + """Declaring a result type is what makes a result JSON cannot express + reportable: the type says how to read it. + """ + assert normalize_agent_result(Verdict(approved=True, note="clean"), Verdict) == { + "approved": True, + "note": "clean", + "score": 0.0, + } + + +def test_normalize_narrows_a_wider_result_to_what_the_type_declares() -> None: + raw = {"approved": True, "note": "clean", "score": 1.5, "extra": 1} + + assert normalize_agent_result(raw, Verdict) == { + "approved": True, + "note": "clean", + "score": 1.5, + } + + +def test_an_undeclared_result_type_still_rejects_a_model_instance() -> None: + """The undeclared path is unchanged: a model instance is still what it refuses.""" + with pytest.raises(TypeError, match="must be JSON-compatible at result"): + normalize_agent_result(Verdict(approved=True)) + + +def test_declaring_object_is_the_same_as_declaring_nothing() -> None: + raw = {"items": [1, 2]} + + assert normalize_agent_result(raw, object) == normalize_agent_result(raw) + + +def test_a_declared_result_type_still_rejects_what_json_cannot_express() -> None: + """A declared type can still render a field JSON cannot express.""" + raw = {"approved": True, "note": "clean", "score": float("inf")} + + with pytest.raises(ValueError) as error: + normalize_agent_result(raw, Verdict) + + assert str(error.value) == "Non-finite number in sub-agent result at result.score" + + +def test_chat_message_content_renders_a_missing_result_as_json_null() -> None: + assert to_chat_message_content(None) == "null" + + +def test_chat_message_content_renders_scalars_and_containers() -> None: + assert to_chat_message_content("text") == "text" + assert to_chat_message_content(7) == "7" + # A boolean has to read as JSON to the model rather than as Python's True. + assert to_chat_message_content(True) == "true" + assert to_chat_message_content({"a": 1, "b": [2]}) == '{"a":1,"b":[2]}' + assert to_chat_message_content([1, 2]) == "[1,2]" diff --git a/python/flink_agents/plan/tests/test_agent_plan.py b/python/flink_agents/plan/tests/test_agent_plan.py index 35226da00..de02f6469 100644 --- a/python/flink_agents/plan/tests/test_agent_plan.py +++ b/python/flink_agents/plan/tests/test_agent_plan.py @@ -514,6 +514,24 @@ def test_agent_plan_accepts_matching_decorated_python_tool_injected_args() -> No } +def test_tool_name_with_reserved_subagent_prefix_is_rejected() -> None: + """Sub-agent callables reach the model under the reserved ``subagent_`` + prefix, so a tool registered under that prefix could never be called and + is rejected at plan-construction time. + """ + agent = Agent() + agent.add_resource( + name="subagent_helper", + resource_type=ResourceType.TOOL, + instance=ApiFunctionTool( + func=ApiPythonFunction.from_callable(query_order), + ), + ) + + with pytest.raises(ValueError, match="reserved prefix 'subagent_'"): + AgentPlan.from_agent(agent, AgentConfiguration()) + + def test_agent_plan_rejects_conflicting_decorated_python_tool_injected_args() -> None: agent = Agent() agent.add_resource(