From 6a721227a4edb9983bdfe1f35b14887508d10af9 Mon Sep 17 00:00:00 2001
From: Alexxigang <37231458+Alexxigang@users.noreply.github.com>
Date: Thu, 9 Apr 2026 22:10:14 +0800
Subject: [PATCH] fix(agent): hoist structured output defs to root
---
.../agent/StructuredOutputCapableAgent.java | 714 +++++++++---------
.../StructuredOutputSchemaWrappingTest.java | 202 +++++
2 files changed, 571 insertions(+), 345 deletions(-)
create mode 100644 agentscope-core/src/test/java/io/agentscope/core/agent/StructuredOutputSchemaWrappingTest.java
diff --git a/agentscope-core/src/main/java/io/agentscope/core/agent/StructuredOutputCapableAgent.java b/agentscope-core/src/main/java/io/agentscope/core/agent/StructuredOutputCapableAgent.java
index 06a5eaeea..3a8b2313d 100644
--- a/agentscope-core/src/main/java/io/agentscope/core/agent/StructuredOutputCapableAgent.java
+++ b/agentscope-core/src/main/java/io/agentscope/core/agent/StructuredOutputCapableAgent.java
@@ -1,345 +1,369 @@
-/*
- * Copyright 2024-2026 the original author or authors.
- *
- * Licensed 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 io.agentscope.core.agent;
-
-import com.fasterxml.jackson.databind.JsonNode;
-import io.agentscope.core.hook.Hook;
-import io.agentscope.core.memory.Memory;
-import io.agentscope.core.message.ContentBlock;
-import io.agentscope.core.message.MessageMetadataKeys;
-import io.agentscope.core.message.Msg;
-import io.agentscope.core.message.MsgRole;
-import io.agentscope.core.message.TextBlock;
-import io.agentscope.core.message.ThinkingBlock;
-import io.agentscope.core.message.ToolResultBlock;
-import io.agentscope.core.model.ChatUsage;
-import io.agentscope.core.model.GenerateOptions;
-import io.agentscope.core.model.StructuredOutputReminder;
-import io.agentscope.core.tool.AgentTool;
-import io.agentscope.core.tool.ToolCallParam;
-import io.agentscope.core.tool.Toolkit;
-import io.agentscope.core.util.JsonSchemaUtils;
-import io.agentscope.core.util.JsonUtils;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import reactor.core.publisher.Mono;
-
-/**
- * Abstract base class for agents that support structured output generation.
- *
- * This class provides the infrastructure for generating structured output using the
- * {@code generate_response} tool pattern combined with StructuredOutputHook for flow control.
- *
- *
Key Features:
- *
- * - Automatic tool registration for structured output
- * - Schema validation before tool execution
- * - Memory compression after structured output completion
- * - Configurable reminder mode (TOOL_CHOICE or PROMPT)
- *
- *
- * Subclass Requirements:
- *
- * - Provide Toolkit via constructor
- * - Implement {@link #getMemory()} for memory access
- * - Implement {@link #buildGenerateOptions()} for model options
- *
- */
-public abstract class StructuredOutputCapableAgent extends AgentBase {
-
- private static final Logger log = LoggerFactory.getLogger(StructuredOutputCapableAgent.class);
-
- /** The tool name for structured output generation. */
- public static final String STRUCTURED_OUTPUT_TOOL_NAME = "generate_response";
-
- protected final Toolkit toolkit;
- protected final StructuredOutputReminder structuredOutputReminder;
-
- /**
- * Constructor with default reminder mode (TOOL_CHOICE).
- */
- protected StructuredOutputCapableAgent(
- String name,
- String description,
- boolean checkRunning,
- List hooks,
- Toolkit toolkit) {
- this(name, description, checkRunning, hooks, toolkit, StructuredOutputReminder.TOOL_CHOICE);
- }
-
- /**
- * Constructor with custom reminder mode.
- */
- protected StructuredOutputCapableAgent(
- String name,
- String description,
- boolean checkRunning,
- List hooks,
- Toolkit toolkit,
- StructuredOutputReminder structuredOutputReminder) {
- super(name, description, checkRunning, hooks);
- this.toolkit = toolkit != null ? toolkit : new Toolkit();
- this.structuredOutputReminder =
- structuredOutputReminder != null
- ? structuredOutputReminder
- : StructuredOutputReminder.TOOL_CHOICE;
- }
-
- /**
- * Get the toolkit for tool operations.
- */
- public Toolkit getToolkit() {
- return toolkit;
- }
-
- /**
- * Get the memory for structured output hook.
- * Subclasses must implement this.
- */
- public abstract Memory getMemory();
-
- /**
- * Build generate options for model calls.
- * Subclasses must implement this.
- */
- protected abstract GenerateOptions buildGenerateOptions();
-
- // ==================== Structured Output Implementation ====================
-
- @Override
- protected final Mono doCall(List msgs, Class structuredOutputClass) {
- return executeWithStructuredOutput(msgs, structuredOutputClass, null);
- }
-
- @Override
- protected final Mono doCall(List msgs, JsonNode outputSchema) {
- return executeWithStructuredOutput(msgs, null, outputSchema);
- }
-
- /**
- * Execute with structured output using StructuredOutputHook.
- */
- private Mono executeWithStructuredOutput(
- List msgs, Class targetClass, JsonNode schemaDesc) {
-
- // Validate parameters
- if (targetClass == null && schemaDesc == null) {
- return Mono.error(
- new IllegalArgumentException(
- "Either targetClass or schemaDesc must be provided"));
- }
- if (targetClass != null && schemaDesc != null) {
- return Mono.error(
- new IllegalArgumentException("Cannot provide both targetClass and schemaDesc"));
- }
-
- return Mono.defer(
- () -> {
- // Create and register temporary tool
- Map jsonSchema =
- targetClass != null
- ? JsonSchemaUtils.generateSchemaFromClass(targetClass)
- : JsonSchemaUtils.generateSchemaFromJsonNode(schemaDesc);
- AgentTool structuredOutputTool =
- createStructuredOutputTool(jsonSchema, targetClass, schemaDesc);
- toolkit.registerAgentTool(structuredOutputTool);
-
- // Create hook for flow control
- StructuredOutputHook hook =
- new StructuredOutputHook(
- structuredOutputReminder, buildGenerateOptions(), getMemory());
-
- addHook(hook);
-
- return doCall(msgs)
- .flatMap(
- result -> {
- // Extract result from hook's output
- Msg hookResult = hook.getResultMsg();
- if (hookResult != null) {
- Msg extracted = extractStructuredResult(hookResult);
- // Merge aggregated metadata from reasoning rounds
- if (extracted != null) {
- extracted =
- mergeCollectedMetadata(
- extracted,
- hook.getAggregatedUsage(),
- hook.getAggregatedThinking());
- }
- return Mono.just(extracted);
- }
- return Mono.just(result);
- })
- .doFinally(
- signal -> {
- removeHook(hook);
- toolkit.removeToolIfSame(
- STRUCTURED_OUTPUT_TOOL_NAME, structuredOutputTool);
- });
- });
- }
-
- /**
- * Create the structured output tool with validation.
- */
- private AgentTool createStructuredOutputTool(
- Map schema, Class targetClass, JsonNode schemaDesc) {
- return new AgentTool() {
- @Override
- public String getName() {
- return STRUCTURED_OUTPUT_TOOL_NAME;
- }
-
- @Override
- public String getDescription() {
- return "Generate the final structured response. Call this function when"
- + " you have all the information needed to provide a complete answer.";
- }
-
- @Override
- public Map getParameters() {
- Map params = new HashMap<>();
- params.put("type", "object");
- params.put("properties", Map.of("response", schema));
- params.put("required", List.of("response"));
- return params;
- }
-
- @Override
- public Mono callAsync(ToolCallParam param) {
- return Mono.fromCallable(
- () -> {
- Object responseData = param.getInput().get("response");
-
- // The tool simply stores the raw data
- // Validation is done by ToolExecutor before calling this
- String contentText = "";
- if (responseData != null) {
- try {
- contentText = JsonUtils.getJsonCodec().toJson(responseData);
- } catch (Exception e) {
- contentText = responseData.toString();
- }
- }
-
- log.debug("Structured output generated: {}", contentText);
-
- // Create response message
- Msg responseMsg =
- Msg.builder()
- .name(getName())
- .role(MsgRole.ASSISTANT)
- .content(TextBlock.builder().text(contentText).build())
- .metadata(
- responseData != null
- ? Map.of("response", responseData)
- : Map.of())
- .build();
-
- Map toolMetadata = new HashMap<>();
- toolMetadata.put("success", true);
- toolMetadata.put("response_msg", responseMsg);
-
- return ToolResultBlock.of(
- List.of(
- TextBlock.builder()
- .text("Successfully generated response.")
- .build()),
- toolMetadata);
- });
- }
- };
- }
-
- /**
- * Extract structured result from tool result message.
- */
- private Msg extractStructuredResult(Msg hookResultMsg) {
- if (hookResultMsg == null) {
- return null;
- }
-
- List toolResults = hookResultMsg.getContentBlocks(ToolResultBlock.class);
- for (ToolResultBlock result : toolResults) {
- if (result.getMetadata() != null
- && Boolean.TRUE.equals(result.getMetadata().get("success"))
- && result.getMetadata().containsKey("response_msg")) {
- Object responseMsgObj = result.getMetadata().get("response_msg");
- if (responseMsgObj instanceof Msg responseMsg) {
- return extractResponseData(responseMsg);
- }
- }
- }
-
- return hookResultMsg;
- }
-
- @SuppressWarnings("unchecked")
- private Msg extractResponseData(Msg responseMsg) {
- if (responseMsg.getMetadata() != null
- && responseMsg.getMetadata().containsKey("response")) {
- Object responseData = responseMsg.getMetadata().get("response");
- // Preserve all original metadata and add structured output under dedicated key
- Map metadata = new HashMap<>(responseMsg.getMetadata());
- metadata.put(MessageMetadataKeys.STRUCTURED_OUTPUT, responseData);
- metadata.remove("response"); // Remove temp key, use standard key
- return Msg.builder()
- .name(responseMsg.getName())
- .role(responseMsg.getRole())
- .content(responseMsg.getContent())
- .metadata(metadata)
- .build();
- }
- return responseMsg;
- }
-
- /**
- * Merge collected metadata (ChatUsage and ThinkingBlock) into the message.
- */
- private Msg mergeCollectedMetadata(Msg msg, ChatUsage chatUsage, ThinkingBlock thinking) {
- // Merge ChatUsage into metadata
- Map metadata =
- new HashMap<>(msg.getMetadata() != null ? msg.getMetadata() : Map.of());
- if (chatUsage != null) {
- metadata.put(MessageMetadataKeys.CHAT_USAGE, chatUsage);
- }
-
- // Merge ThinkingBlock into content
- List newContent;
- if (thinking != null) {
- newContent = new ArrayList<>();
- newContent.add(thinking); // ThinkingBlock first
- if (msg.getContent() != null) {
- newContent.addAll(msg.getContent());
- }
- } else {
- newContent = msg.getContent();
- }
-
- return Msg.builder()
- .id(msg.getId())
- .name(msg.getName())
- .role(msg.getRole())
- .content(newContent)
- .metadata(metadata)
- .timestamp(msg.getTimestamp())
- .build();
- }
-}
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed 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 io.agentscope.core.agent;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import io.agentscope.core.hook.Hook;
+import io.agentscope.core.memory.Memory;
+import io.agentscope.core.message.ContentBlock;
+import io.agentscope.core.message.MessageMetadataKeys;
+import io.agentscope.core.message.Msg;
+import io.agentscope.core.message.MsgRole;
+import io.agentscope.core.message.TextBlock;
+import io.agentscope.core.message.ThinkingBlock;
+import io.agentscope.core.message.ToolResultBlock;
+import io.agentscope.core.model.ChatUsage;
+import io.agentscope.core.model.GenerateOptions;
+import io.agentscope.core.model.StructuredOutputReminder;
+import io.agentscope.core.tool.AgentTool;
+import io.agentscope.core.tool.ToolCallParam;
+import io.agentscope.core.tool.Toolkit;
+import io.agentscope.core.util.JsonSchemaUtils;
+import io.agentscope.core.util.JsonUtils;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
+
+/**
+ * Abstract base class for agents that support structured output generation.
+ *
+ * This class provides the infrastructure for generating structured output using the
+ * {@code generate_response} tool pattern combined with StructuredOutputHook for flow control.
+ *
+ *
Key Features:
+ *
+ * - Automatic tool registration for structured output
+ * - Schema validation before tool execution
+ * - Memory compression after structured output completion
+ * - Configurable reminder mode (TOOL_CHOICE or PROMPT)
+ *
+ *
+ * Subclass Requirements:
+ *
+ * - Provide Toolkit via constructor
+ * - Implement {@link #getMemory()} for memory access
+ * - Implement {@link #buildGenerateOptions()} for model options
+ *
+ */
+public abstract class StructuredOutputCapableAgent extends AgentBase {
+
+ private static final Logger log = LoggerFactory.getLogger(StructuredOutputCapableAgent.class);
+
+ /** The tool name for structured output generation. */
+ public static final String STRUCTURED_OUTPUT_TOOL_NAME = "generate_response";
+
+ protected final Toolkit toolkit;
+ protected final StructuredOutputReminder structuredOutputReminder;
+
+ /**
+ * Constructor with default reminder mode (TOOL_CHOICE).
+ */
+ protected StructuredOutputCapableAgent(
+ String name,
+ String description,
+ boolean checkRunning,
+ List hooks,
+ Toolkit toolkit) {
+ this(name, description, checkRunning, hooks, toolkit, StructuredOutputReminder.TOOL_CHOICE);
+ }
+
+ /**
+ * Constructor with custom reminder mode.
+ */
+ protected StructuredOutputCapableAgent(
+ String name,
+ String description,
+ boolean checkRunning,
+ List hooks,
+ Toolkit toolkit,
+ StructuredOutputReminder structuredOutputReminder) {
+ super(name, description, checkRunning, hooks);
+ this.toolkit = toolkit != null ? toolkit : new Toolkit();
+ this.structuredOutputReminder =
+ structuredOutputReminder != null
+ ? structuredOutputReminder
+ : StructuredOutputReminder.TOOL_CHOICE;
+ }
+
+ /**
+ * Get the toolkit for tool operations.
+ */
+ public Toolkit getToolkit() {
+ return toolkit;
+ }
+
+ /**
+ * Get the memory for structured output hook.
+ * Subclasses must implement this.
+ */
+ public abstract Memory getMemory();
+
+ /**
+ * Build generate options for model calls.
+ * Subclasses must implement this.
+ */
+ protected abstract GenerateOptions buildGenerateOptions();
+
+ // ==================== Structured Output Implementation ====================
+
+ @Override
+ protected final Mono doCall(List msgs, Class structuredOutputClass) {
+ return executeWithStructuredOutput(msgs, structuredOutputClass, null);
+ }
+
+ @Override
+ protected final Mono doCall(List msgs, JsonNode outputSchema) {
+ return executeWithStructuredOutput(msgs, null, outputSchema);
+ }
+
+ /**
+ * Execute with structured output using StructuredOutputHook.
+ */
+ private Mono executeWithStructuredOutput(
+ List msgs, Class targetClass, JsonNode schemaDesc) {
+
+ // Validate parameters
+ if (targetClass == null && schemaDesc == null) {
+ return Mono.error(
+ new IllegalArgumentException(
+ "Either targetClass or schemaDesc must be provided"));
+ }
+ if (targetClass != null && schemaDesc != null) {
+ return Mono.error(
+ new IllegalArgumentException("Cannot provide both targetClass and schemaDesc"));
+ }
+
+ return Mono.defer(
+ () -> {
+ // Create and register temporary tool
+ Map jsonSchema =
+ targetClass != null
+ ? JsonSchemaUtils.generateSchemaFromClass(targetClass)
+ : JsonSchemaUtils.generateSchemaFromJsonNode(schemaDesc);
+ AgentTool structuredOutputTool =
+ createStructuredOutputTool(jsonSchema, targetClass, schemaDesc);
+ toolkit.registerAgentTool(structuredOutputTool);
+
+ // Create hook for flow control
+ StructuredOutputHook hook =
+ new StructuredOutputHook(
+ structuredOutputReminder, buildGenerateOptions(), getMemory());
+
+ addHook(hook);
+
+ return doCall(msgs)
+ .flatMap(
+ result -> {
+ // Extract result from hook's output
+ Msg hookResult = hook.getResultMsg();
+ if (hookResult != null) {
+ Msg extracted = extractStructuredResult(hookResult);
+ // Merge aggregated metadata from reasoning rounds
+ if (extracted != null) {
+ extracted =
+ mergeCollectedMetadata(
+ extracted,
+ hook.getAggregatedUsage(),
+ hook.getAggregatedThinking());
+ }
+ return Mono.just(extracted);
+ }
+ return Mono.just(result);
+ })
+ .doFinally(
+ signal -> {
+ removeHook(hook);
+ toolkit.removeToolIfSame(
+ STRUCTURED_OUTPUT_TOOL_NAME, structuredOutputTool);
+ });
+ });
+ }
+
+ /**
+ * Create the structured output tool with validation.
+ */
+ private AgentTool createStructuredOutputTool(
+ Map schema, Class targetClass, JsonNode schemaDesc) {
+ return new AgentTool() {
+ @Override
+ public String getName() {
+ return STRUCTURED_OUTPUT_TOOL_NAME;
+ }
+
+ @Override
+ public String getDescription() {
+ return "Generate the final structured response. Call this function when"
+ + " you have all the information needed to provide a complete answer.";
+ }
+
+ @Override
+ public Map getParameters() {
+ return wrapStructuredOutputSchema(schema);
+ }
+
+ @Override
+ public Mono callAsync(ToolCallParam param) {
+ return Mono.fromCallable(
+ () -> {
+ Object responseData = param.getInput().get("response");
+
+ // The tool simply stores the raw data
+ // Validation is done by ToolExecutor before calling this
+ String contentText = "";
+ if (responseData != null) {
+ try {
+ contentText = JsonUtils.getJsonCodec().toJson(responseData);
+ } catch (Exception e) {
+ contentText = responseData.toString();
+ }
+ }
+
+ log.debug("Structured output generated: {}", contentText);
+
+ // Create response message
+ Msg responseMsg =
+ Msg.builder()
+ .name(getName())
+ .role(MsgRole.ASSISTANT)
+ .content(TextBlock.builder().text(contentText).build())
+ .metadata(
+ responseData != null
+ ? Map.of("response", responseData)
+ : Map.of())
+ .build();
+
+ Map toolMetadata = new HashMap<>();
+ toolMetadata.put("success", true);
+ toolMetadata.put("response_msg", responseMsg);
+
+ return ToolResultBlock.of(
+ List.of(
+ TextBlock.builder()
+ .text("Successfully generated response.")
+ .build()),
+ toolMetadata);
+ });
+ }
+ };
+ }
+
+ static Map wrapStructuredOutputSchema(Map schema) {
+ Map responseSchema = new HashMap<>(schema);
+ Map params = new HashMap<>();
+ Map properties = new HashMap<>();
+ Map rootDefs = new HashMap<>();
+
+ hoistDefs(responseSchema, "$defs", rootDefs);
+ hoistDefs(responseSchema, "definitions", rootDefs);
+
+ properties.put("response", responseSchema);
+ params.put("type", "object");
+ params.put("properties", properties);
+ params.put("required", List.of("response"));
+ if (!rootDefs.isEmpty()) {
+ params.put("$defs", rootDefs);
+ }
+ return params;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static void hoistDefs(
+ Map schema, String key, Map rootDefs) {
+ Object rawDefs = schema.remove(key);
+ if (rawDefs instanceof Map defs && !defs.isEmpty()) {
+ rootDefs.putAll((Map) defs);
+ }
+ }
+
+ /**
+ * Extract structured result from tool result message.
+ */
+ private Msg extractStructuredResult(Msg hookResultMsg) {
+ if (hookResultMsg == null) {
+ return null;
+ }
+
+ List toolResults = hookResultMsg.getContentBlocks(ToolResultBlock.class);
+ for (ToolResultBlock result : toolResults) {
+ if (result.getMetadata() != null
+ && Boolean.TRUE.equals(result.getMetadata().get("success"))
+ && result.getMetadata().containsKey("response_msg")) {
+ Object responseMsgObj = result.getMetadata().get("response_msg");
+ if (responseMsgObj instanceof Msg responseMsg) {
+ return extractResponseData(responseMsg);
+ }
+ }
+ }
+
+ return hookResultMsg;
+ }
+
+ @SuppressWarnings("unchecked")
+ private Msg extractResponseData(Msg responseMsg) {
+ if (responseMsg.getMetadata() != null
+ && responseMsg.getMetadata().containsKey("response")) {
+ Object responseData = responseMsg.getMetadata().get("response");
+ // Preserve all original metadata and add structured output under dedicated key
+ Map metadata = new HashMap<>(responseMsg.getMetadata());
+ metadata.put(MessageMetadataKeys.STRUCTURED_OUTPUT, responseData);
+ metadata.remove("response"); // Remove temp key, use standard key
+ return Msg.builder()
+ .name(responseMsg.getName())
+ .role(responseMsg.getRole())
+ .content(responseMsg.getContent())
+ .metadata(metadata)
+ .build();
+ }
+ return responseMsg;
+ }
+
+ /**
+ * Merge collected metadata (ChatUsage and ThinkingBlock) into the message.
+ */
+ private Msg mergeCollectedMetadata(Msg msg, ChatUsage chatUsage, ThinkingBlock thinking) {
+ // Merge ChatUsage into metadata
+ Map metadata =
+ new HashMap<>(msg.getMetadata() != null ? msg.getMetadata() : Map.of());
+ if (chatUsage != null) {
+ metadata.put(MessageMetadataKeys.CHAT_USAGE, chatUsage);
+ }
+
+ // Merge ThinkingBlock into content
+ List newContent;
+ if (thinking != null) {
+ newContent = new ArrayList<>();
+ newContent.add(thinking); // ThinkingBlock first
+ if (msg.getContent() != null) {
+ newContent.addAll(msg.getContent());
+ }
+ } else {
+ newContent = msg.getContent();
+ }
+
+ return Msg.builder()
+ .id(msg.getId())
+ .name(msg.getName())
+ .role(msg.getRole())
+ .content(newContent)
+ .metadata(metadata)
+ .timestamp(msg.getTimestamp())
+ .build();
+ }
+}
diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/StructuredOutputSchemaWrappingTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/StructuredOutputSchemaWrappingTest.java
new file mode 100644
index 000000000..b7dd3c277
--- /dev/null
+++ b/agentscope-core/src/test/java/io/agentscope/core/agent/StructuredOutputSchemaWrappingTest.java
@@ -0,0 +1,202 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * Licensed 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 io.agentscope.core.agent;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.agentscope.core.ReActAgent;
+import io.agentscope.core.agent.test.MockModel;
+import io.agentscope.core.memory.InMemoryMemory;
+import io.agentscope.core.memory.Memory;
+import io.agentscope.core.message.Msg;
+import io.agentscope.core.message.MsgRole;
+import io.agentscope.core.message.TextBlock;
+import io.agentscope.core.message.ToolUseBlock;
+import io.agentscope.core.model.ChatResponse;
+import io.agentscope.core.model.ChatUsage;
+import io.agentscope.core.tool.ToolValidator;
+import io.agentscope.core.tool.Toolkit;
+import io.agentscope.core.util.JsonSchemaUtils;
+import io.agentscope.core.util.JsonUtils;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+class StructuredOutputSchemaWrappingTest {
+
+ static class OrderInfo {
+ public String orderId;
+ public Address billingAddress;
+ public Address shippingAddress;
+ public List items;
+ }
+
+ static class Address {
+ public String city;
+ public String street;
+ public String recipientName;
+ }
+
+ static class OrderItem {
+ public String productName;
+ public Integer quantity;
+ }
+
+ @Test
+ void shouldHoistDefsToRootWhenWrappingStructuredOutputSchema() {
+ Map responseSchema =
+ JsonSchemaUtils.generateSchemaFromClass(OrderInfo.class);
+ Map wrappedSchema =
+ StructuredOutputCapableAgent.wrapStructuredOutputSchema(responseSchema);
+
+ @SuppressWarnings("unchecked")
+ Map properties = (Map) wrappedSchema.get("properties");
+ @SuppressWarnings("unchecked")
+ Map wrappedResponse = (Map) properties.get("response");
+ @SuppressWarnings("unchecked")
+ Map defs = (Map) wrappedSchema.get("$defs");
+
+ assertNotNull(defs, "Wrapped schema should expose shared defs at root level");
+ assertTrue(defs.containsKey("Address"), "Address should be hoisted to root $defs");
+ assertFalse(
+ wrappedResponse.containsKey("$defs"),
+ "Nested response schema should not keep $defs");
+
+ Map validInput =
+ Map.of(
+ "response",
+ Map.of(
+ "orderId",
+ "ORD-001",
+ "billingAddress",
+ Map.of(
+ "city", "Beijing",
+ "street", "Jianguo Road",
+ "recipientName", "Zhang San"),
+ "shippingAddress",
+ Map.of(
+ "city", "Shanghai",
+ "street", "Century Avenue",
+ "recipientName", "Li Si"),
+ "items",
+ List.of(Map.of("productName", "Phone", "quantity", 2))));
+
+ String validationError =
+ ToolValidator.validateInput(
+ JsonUtils.getJsonCodec().toJson(validInput), wrappedSchema);
+ assertNull(
+ validationError, "Wrapped structured-output schema should validate reused objects");
+ }
+
+ @Test
+ void shouldSupportStructuredOutputWithReusedNestedObjects() {
+ Toolkit toolkit = new Toolkit();
+ Memory memory = new InMemoryMemory();
+ Map toolInput =
+ Map.of(
+ "response",
+ Map.of(
+ "orderId",
+ "ORD-001",
+ "billingAddress",
+ Map.of(
+ "city", "Beijing",
+ "street", "Jianguo Road",
+ "recipientName", "Zhang San"),
+ "shippingAddress",
+ Map.of(
+ "city", "Shanghai",
+ "street", "Century Avenue",
+ "recipientName", "Li Si"),
+ "items",
+ List.of(
+ Map.of("productName", "Phone", "quantity", 2),
+ Map.of("productName", "Headset", "quantity", 1))));
+
+ MockModel mockModel =
+ new MockModel(
+ msgs -> {
+ boolean hasToolResults =
+ msgs.stream().anyMatch(m -> m.getRole() == MsgRole.TOOL);
+ if (!hasToolResults) {
+ return List.of(
+ ChatResponse.builder()
+ .id("msg_1")
+ .content(
+ List.of(
+ ToolUseBlock.builder()
+ .id("call_structured_1")
+ .name("generate_response")
+ .input(toolInput)
+ .content(
+ JsonUtils
+ .getJsonCodec()
+ .toJson(
+ toolInput))
+ .build()))
+ .usage(new ChatUsage(10, 20, 30))
+ .build());
+ }
+ return List.of(
+ ChatResponse.builder()
+ .id("msg_2")
+ .content(
+ List.of(
+ TextBlock.builder()
+ .text(
+ "Structured response"
+ + " generated")
+ .build()))
+ .usage(new ChatUsage(5, 10, 15))
+ .build());
+ });
+
+ ReActAgent agent =
+ ReActAgent.builder()
+ .name("order-agent")
+ .sysPrompt("You extract order information")
+ .model(mockModel)
+ .toolkit(toolkit)
+ .memory(memory)
+ .build();
+
+ Msg inputMsg =
+ Msg.builder()
+ .name("user")
+ .role(MsgRole.USER)
+ .content(TextBlock.builder().text("Extract the order information.").build())
+ .build();
+
+ Msg responseMsg = agent.call(inputMsg, OrderInfo.class).block();
+ assertNotNull(responseMsg);
+ assertNotNull(responseMsg.getMetadata());
+
+ OrderInfo result = responseMsg.getStructuredData(OrderInfo.class);
+ assertNotNull(result);
+ assertEquals("ORD-001", result.orderId);
+ assertNotNull(result.billingAddress);
+ assertNotNull(result.shippingAddress);
+ assertEquals("Beijing", result.billingAddress.city);
+ assertEquals("Shanghai", result.shippingAddress.city);
+ assertNotNull(result.items);
+ assertEquals(2, result.items.size());
+ assertEquals("Phone", result.items.get(0).productName);
+ }
+}