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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 59 additions & 12 deletions api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,17 @@
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.skills.SkillSourceSpec;
import org.apache.flink.agents.api.skills.Skills;
import org.apache.flink.agents.api.subagent.InternalSubagentResourceDescriptorBuilder;
import org.apache.flink.agents.api.tools.FunctionTool;
import org.apache.flink.agents.api.yaml.spec.ActionSpec;
import org.apache.flink.agents.api.yaml.spec.AgentActionRef;
import org.apache.flink.agents.api.yaml.spec.AgentSpec;
import org.apache.flink.agents.api.yaml.spec.DescriptorSpec;
import org.apache.flink.agents.api.yaml.spec.InternalSubagentSpec;
import org.apache.flink.agents.api.yaml.spec.PackageSkillSpec;
import org.apache.flink.agents.api.yaml.spec.PromptSpec;
import org.apache.flink.agents.api.yaml.spec.SkillsSpec;
import org.apache.flink.agents.api.yaml.spec.SubagentSpec;
import org.apache.flink.agents.api.yaml.spec.ToolSpec;
import org.apache.flink.agents.api.yaml.spec.UrlSkillSpec;
import org.apache.flink.agents.api.yaml.spec.YamlAgentsDocument;
Expand Down Expand Up @@ -159,7 +162,7 @@ public static FunctionTool buildTool(ToolSpec spec) {

/** Build a {@link Prompt} from a parsed {@link PromptSpec}. */
public static Prompt buildPrompt(PromptSpec spec) {
if (spec.getText() != null) {
if (spec.getText() != null && !spec.getText().isEmpty()) {
return Prompt.fromText(spec.getText());
}
List<ChatMessage> messages =
Expand All @@ -169,6 +172,37 @@ public static Prompt buildPrompt(PromptSpec spec) {
return Prompt.fromMessages(messages);
}

private static ResourceDescriptor buildSubagent(SubagentSpec spec, Map<String, Agent> agents) {
if (spec.getExternal() != null) {
return buildDescriptor(spec.getExternal(), ResourceType.AGENT);
}
InternalSubagentSpec internal = spec.getInternal();
Agent child = agents.get(internal.getAgent());
if (child == null) {
throw new IllegalArgumentException(
"Sub-agent '"
+ internal.getName()
+ "' references unknown agent '"
+ internal.getAgent()
+ "'; the child must be declared in the same YAML file or already"
+ " registered in this environment.");
}
InternalSubagentResourceDescriptorBuilder builder =
new InternalSubagentResourceDescriptorBuilder()
.agent(child)
.description(internal.getDescription());
if (internal.getInputSchema() != null) {
try {
builder.inputSchema(
new ObjectMapper().writeValueAsString(internal.getInputSchema()));
} catch (IOException e) {
throw new IllegalArgumentException(
"Invalid input schema for sub-agent '" + internal.getName() + "'", e);
}
}
return builder.build();
}

/** Build a {@link Skills} resource from a parsed {@link SkillsSpec}. */
public static Skills buildSkills(SkillsSpec spec) {
List<SkillSourceSpec> sources = new ArrayList<>();
Expand Down Expand Up @@ -250,6 +284,10 @@ Map<String, AgentSpec> getAgentSpecs() {

/** Parse one YAML file and build the agents it declares. */
public static LoadedFile buildAgents(Path path) {
return buildAgents(path, Collections.emptyMap());
}

private static LoadedFile buildAgents(Path path, Map<String, Agent> registeredAgents) {
YamlAgentsDocument doc;
try {
doc = YAML_MAPPER.readValue(Files.readString(path), YamlAgentsDocument.class);
Expand All @@ -267,10 +305,27 @@ public static LoadedFile buildAgents(Path path) {
throw new IllegalArgumentException(
"Duplicate agent name '" + spec.getName() + "' in " + path);
}
if (registeredAgents.containsKey(spec.getName())) {
throw new IllegalArgumentException(
"Duplicate agent name '" + spec.getName() + "' (loading " + path + ")");
}
agentSpecs.put(spec.getName(), spec);
agents.put(spec.getName(), buildAgent(spec));
}

Map<String, Agent> availableAgents = new LinkedHashMap<>(registeredAgents);
availableAgents.putAll(agents);
// Bind after all in-file Agents exist, including references to previously loaded Agents.
for (AgentSpec spec : doc.getAgents()) {
for (SubagentSpec subagent : spec.getSubagents()) {
agents.get(spec.getName())
.addResource(
subagent.getName(),
ResourceType.AGENT,
buildSubagent(subagent, availableAgents));
}
}

Map<ResourceType, Map<String, Object>> sharedResources = new EnumMap<>(ResourceType.class);
for (ResourceType t : ResourceType.values()) {
sharedResources.put(t, new LinkedHashMap<>());
Expand Down Expand Up @@ -334,11 +389,12 @@ public static void loadYaml(AgentsExecutionEnvironment env, Path path) {
/**
* Load multiple YAML files. Multiple calls accumulate. Duplicate names — both within a single
* file (caught by {@link #buildAgents(Path)}) and across the current environment — raise {@link
* IllegalArgumentException}.
* IllegalArgumentException}. Files are loaded in order; child references may resolve Agents
* registered by earlier files or calls.
*/
public static void loadYaml(AgentsExecutionEnvironment env, List<Path> paths) {
for (Path path : paths) {
LoadedFile loaded = buildAgents(path);
LoadedFile loaded = buildAgents(path, env.getAgents());

for (Map.Entry<String, Agent> entry : loaded.getAgents().entrySet()) {
AgentSpec spec = loaded.getAgentSpecs().get(entry.getKey());
Expand Down Expand Up @@ -368,14 +424,6 @@ public static void loadYaml(AgentsExecutionEnvironment env, List<Path> paths) {
}
}

// Cross-environment agent name uniqueness check.
for (String name : loaded.getAgents().keySet()) {
if (env.getAgents().containsKey(name)) {
throw new IllegalArgumentException(
"Duplicate agent name '" + name + "' (loading " + path + ")");
}
}

// Commit shared resources — env.addResource enforces dedup with its own message.
for (Map.Entry<ResourceType, Map<String, Object>> e :
loaded.getSharedResources().entrySet()) {
Expand Down Expand Up @@ -413,7 +461,6 @@ private static Agent buildAgent(AgentSpec spec) {
addAgentDescriptors(agent, ResourceType.EMBEDDING_MODEL, spec.getEmbeddingModelSetups());
addAgentDescriptors(agent, ResourceType.VECTOR_STORE, spec.getVectorStores());
addAgentDescriptors(agent, ResourceType.MCP_SERVER, spec.getMcpServers());
addAgentDescriptors(agent, ResourceType.AGENT, spec.getSubagents());

for (ToolSpec t : spec.getTools()) {
agent.addResource(t.getName(), ResourceType.TOOL, buildTool(t));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ public final class AgentSpec {
private final List<DescriptorSpec> embeddingModelSetups;
private final List<DescriptorSpec> vectorStores;
private final List<DescriptorSpec> mcpServers;
private final List<DescriptorSpec> subagents;
private final List<SubagentSpec> subagents;

@JsonCreator
public AgentSpec(
Expand All @@ -57,7 +57,7 @@ public AgentSpec(
@JsonProperty("embedding_model_setups") List<DescriptorSpec> embeddingModelSetups,
@JsonProperty("vector_stores") List<DescriptorSpec> vectorStores,
@JsonProperty("mcp_servers") List<DescriptorSpec> mcpServers,
@JsonProperty("subagents") List<DescriptorSpec> subagents) {
@JsonProperty("subagents") List<SubagentSpec> subagents) {
this.name = name;
this.description = description;
this.prompts = orEmpty(prompts);
Expand Down Expand Up @@ -125,7 +125,7 @@ public List<DescriptorSpec> getMcpServers() {
return mcpServers;
}

public List<DescriptorSpec> getSubagents() {
public List<SubagentSpec> getSubagents() {
return subagents;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
/*
* 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.yaml.spec;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;

import java.util.Map;

/** A reference to a child Agent in this file or already registered in the environment. */
public final class InternalSubagentSpec {
private final String name;
private final String agent;
private final String description;
private final Map<String, Object> inputSchema;

@JsonCreator
public InternalSubagentSpec(
@JsonProperty(value = "name", required = true) String name,
@JsonProperty(value = "agent", required = true) String agent,
@JsonProperty("description") String description,
@JsonProperty("input_schema") Map<String, Object> inputSchema) {
if (agent == null || agent.isBlank()) {
throw new IllegalArgumentException(
"Internal sub-agent 'agent' must be a non-empty name");
}
this.name = name;
this.agent = agent;
this.description = description == null ? "" : description;
this.inputSchema = inputSchema;
}

public String getName() {
return name;
}

public String getAgent() {
return agent;
}

public String getDescription() {
return description;
}

public Map<String, Object> getInputSchema() {
return inputSchema;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
/*
* 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.yaml.spec;

import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;

import java.io.IOException;

/** Either an internal Agent reference or an external sub-agent descriptor. */
@JsonDeserialize(using = SubagentSpec.Deserializer.class)
public final class SubagentSpec {
private final InternalSubagentSpec internal;
private final DescriptorSpec external;

private SubagentSpec(InternalSubagentSpec internal, DescriptorSpec external) {
this.internal = internal;
this.external = external;
}

public String getName() {
return internal == null ? external.getName() : internal.getName();
}

public InternalSubagentSpec getInternal() {
return internal;
}

public DescriptorSpec getExternal() {
return external;
}

static final class Deserializer extends JsonDeserializer<SubagentSpec> {
@Override
public SubagentSpec deserialize(JsonParser parser, DeserializationContext context)
throws IOException {
JsonNode node = parser.getCodec().readTree(parser);
if (!node.has("clazz")) {
for (String field : new String[] {"name", "agent"}) {
if (!node.path(field).isTextual()) {
context.reportInputMismatch(
InternalSubagentSpec.class,
"Internal sub-agent '%s' must be a string",
field);
}
}
if (node.hasNonNull("description") && !node.get("description").isTextual()) {
context.reportInputMismatch(
InternalSubagentSpec.class,
"Internal sub-agent 'description' must be a string or null");
}
return new SubagentSpec(
parser.getCodec().treeToValue(node, InternalSubagentSpec.class), null);
}
return new SubagentSpec(
null, parser.getCodec().treeToValue(node, DescriptorSpec.class));
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.function.JavaFunction;
import org.apache.flink.agents.api.function.PythonFunction;
import org.apache.flink.agents.api.prompt.Prompt;
Expand All @@ -35,7 +37,10 @@
import org.apache.flink.agents.api.yaml.spec.SkillsSpec;
import org.apache.flink.agents.api.yaml.spec.ToolSpec;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;
Expand Down Expand Up @@ -178,15 +183,27 @@ void buildPromptText() throws Exception {
assertThat(prompt.formatString(Map.of())).isEqualTo("hi");
}

@Test
void buildPromptMessages() throws Exception {
@ParameterizedTest
@ValueSource(strings = {"", "text: null\n", "text: \"\"\n"})
void buildPromptMessagesWithUnsetText(String textField) throws Exception {
PromptSpec spec =
M.readValue(
"name: p\nmessages:\n - {role: system, content: hi}\n", PromptSpec.class);
"name: p\n"
+ textField
+ "messages:\n"
+ " - {role: system, content: 'Hello {name}'}\n"
+ " - {role: user, content: '{question}'}\n",
PromptSpec.class);
Prompt prompt = YamlLoader.buildPrompt(spec);
// Both message-templated and text-templated prompts produce LocalPrompt instances;
// confirm the constructed prompt formats without error.
assertThat(prompt).isNotNull();
List<ChatMessage> messages =
prompt.formatMessages(
MessageRole.SYSTEM, Map.of("name", "Alice", "question", "How are you?"));
assertThat(messages)
.extracting(ChatMessage::getRole)
.containsExactly(MessageRole.SYSTEM, MessageRole.USER);
assertThat(messages)
.extracting(ChatMessage::getText)
.containsExactly("Hello Alice", "How are you?");
}

@Test
Expand Down
Loading
Loading