diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/AgentResult.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/AgentResult.java index a68fe7dcd..13529def3 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/AgentResult.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/AgentResult.java @@ -56,6 +56,9 @@ public record AgentResult(ChatMessage message, Status status, Usage usage, Objec /** Error code for a terminal message that only reports a model failure. */ public static final String ERROR_CODE_MODEL_FAILED = "MODEL_FAILED"; + /** Error code marking that the single bounded repair attempt still failed validation. */ + public static final String ERROR_CODE_REPAIRED_STILL_INVALID = "REPAIRED_STILL_INVALID"; + /** * Delivery status of the run. */ diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/ReactAgent.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/ReactAgent.java index 74c20258d..4cc9ebae7 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/ReactAgent.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/ReactAgent.java @@ -20,6 +20,8 @@ import io.github.agentic.ai.api.chat.ContextAuthorityKeys; import io.github.agentic.ai.api.chat.MessageProvenance; import io.github.agentic.ai.api.chat.Usage; +import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatResponse; import io.github.agentic.ai.graph.CompileConfig; import io.github.agentic.ai.graph.CompiledGraph; @@ -65,6 +67,8 @@ import io.github.agentic.ai.graph.internal.node.Node; import io.github.agentic.ai.graph.internal.node.ResumableSubGraphAction; import io.github.agentic.ai.graph.serializer.AgentInstructionMessage; +import io.github.agentic.ai.graph.streaming.OutputType; +import io.github.agentic.ai.graph.streaming.StreamingOutput; import io.github.agentic.ai.graph.serializer.StateSerializer; import io.github.agentic.ai.graph.serializer.plain_text.jackson.ArgiJacksonStateSerializer; import io.github.agentic.ai.graph.state.strategy.AppendStrategy; @@ -89,6 +93,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.function.Supplier; @@ -111,6 +116,10 @@ public class ReactAgent extends BaseAgent { private static final String MODEL_ERROR_MESSAGE_PREFIX = "Exception: "; + private static final String FINAL_RESPONSE_METADATA_KEY = "argiFinalResponse"; + + private final ChatClient chatClient; + private volatile StructuredOutputFinalizer structuredOutputFinalizer; private final ConcurrentMap> threadIdStateMap; @@ -151,6 +160,7 @@ public ReactAgent(AgentLlmNode llmNode, AgentToolNode toolNode, CompileConfig co this.inputType = builder.inputType; this.outputSchema = builder.outputSchema; this.outputType = builder.outputType; + this.chatClient = builder.chatClient; // Set state serializer from builder, or use default // Default to Jackson serializer for better compatibility and features @@ -338,6 +348,139 @@ public AgentResult callResult(Map inputs, RunnableConfig config) return doResultInvoke(inputs, config); } + /** + * Streams the agent run as provisional {@link AgentResult} chunks followed by the + * final validated delivery. Chunks with a {@code null} status are provisional and + * explicitly unvalidated; the last element of the stream carries the delivery + * outcome of the terminal message (one of the four terminal statuses), finalized + * through the same boundary as {@link #callResult(String)}. A run that ends without + * a terminal response is reported as INTERRUPTED, never as completed. + * @param message the user input message + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(String message) { + return doStreamResult(message, null); + } + + /** + * Streams the agent run with a runtime configuration, mirroring + * {@link #streamResult(String)}. + * @param message the user input message + * @param config runtime configuration controlling execution behavior + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(String message, RunnableConfig config) { + return doStreamResult(message, config); + } + + /** + * Streams the agent run with a {@link ChatMessage} input, mirroring + * {@link #streamResult(String)}. + * @param message the input message + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(ChatMessage message) { + return doStreamResult(message, null); + } + + /** + * Streams the agent run with a {@link ChatMessage} input and runtime configuration, + * mirroring {@link #streamResult(String)}. + * @param message the input message + * @param config runtime configuration controlling execution behavior + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(ChatMessage message, RunnableConfig config) { + return doStreamResult(message, config); + } + + /** + * Streams the agent run with a message list input, mirroring + * {@link #streamResult(String)}. + * @param messages the input messages + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(List messages) { + return doStreamResult(messages, null); + } + + /** + * Streams the agent run with a message list input and runtime configuration, + * mirroring {@link #streamResult(String)}. + * @param messages the input messages + * @param config runtime configuration controlling execution behavior + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(List messages, RunnableConfig config) { + return doStreamResult(messages, config); + } + + /** + * Streams the agent run with an inputs map, mirroring {@link #streamResult(String)}. + * @param inputs the input map (reserved: messages, input; other keys as state) + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(Map inputs) { + return doStreamResult(inputs, null); + } + + /** + * Streams the agent run with an inputs map and runtime configuration, mirroring + * {@link #streamResult(String)}. + * @param inputs the input map (reserved: messages, input; other keys as state) + * @param config runtime configuration controlling execution behavior + * @return the provisional chunk stream terminated by the validated result + */ + public Flux streamResult(Map inputs, RunnableConfig config) { + return doStreamResult(inputs, config); + } + + private Flux doStreamResult(Object message, RunnableConfig config) { + return doStream(buildMessageInput(message), config).transform(this::finalizeStream); + } + + private Flux doStreamResult(Map inputs, RunnableConfig config) { + return doStream(normalizeMessageInput(inputs), config).transform(this::finalizeStream); + } + + /** + * Consumes the node-output stream: text-delta chunks are re-published as provisional + * results (status {@code null}); the final aggregated response (marked by the + * existing {@code argiFinalResponse} protocol convention) is buffered and finalized + * once the stream completes, so unverified chunks are never reported as a delivery. + */ + private Flux finalizeStream(Flux outputs) { + AtomicReference finalResponse = new AtomicReference<>(); + return outputs.concatMap(output -> { + if (!(output instanceof StreamingOutput streamed)) { + return Flux.empty(); + } + if (streamed.getOriginData() instanceof ChatResponse response + && Boolean.TRUE.equals(response.metadata().get(FINAL_RESPONSE_METADATA_KEY))) { + finalResponse.set(response); + return Flux.empty(); + } + if (streamed.getOutputType() == OutputType.AGENT_MODEL_STREAMING) { + return Flux.fromIterable(streamed.messages()) + .map(message -> new AgentResult(message, null, Usage.unknown(), null, List.of())); + } + return Flux.empty(); + }).concatWith(Flux.defer(() -> { + ChatResponse response = finalResponse.get(); + if (response == null) { + return Flux.just(AgentResult.interrupted(Usage.unknown())); + } + ChatMessage terminal = response.messages().stream() + .filter(message -> message.role() == ChatRole.ASSISTANT) + .reduce((first, second) -> second) + .orElse(null); + if (terminal == null) { + return Flux.just(AgentResult.interrupted(response.usage())); + } + return Flux.just(deliver(terminal, response.usage())); + })); + } + private AgentResult doResultInvoke(Object message, RunnableConfig config) { try { return finalizeResult(doInvokeAndGetOutput(buildMessageInput(message), config)); @@ -373,6 +516,11 @@ private AgentResult finalizeResult(Optional output) { // interruption): report it as unfinished, never as completed. return AgentResult.interrupted(usage); } + return deliver(terminal, usage); + } + + /** Maps a terminal message onto the delivery boundary shared by call and stream paths. */ + private AgentResult deliver(ChatMessage terminal, Usage usage) { if (terminal.text() != null && terminal.text().startsWith(MODEL_ERROR_MESSAGE_PREFIX)) { // throwOnModelError=false surfaces model failures as assistant text. return AgentResult.error(terminal, usage, @@ -387,12 +535,15 @@ private AgentResult finalizeResult(Optional output) { * the JSON Schema validator classes, which are not guaranteed to be on every * consumer's classpath (the studio module, for one, resolves an older * json-schema-validator for its MCP client); agents that never finalize a - * schema-constrained output must not trigger that load. + * schema-constrained output must not trigger that load. When the agent has a chat + * client, the finalizer gains the bounded repair capability. */ private StructuredOutputFinalizer finalizer() { StructuredOutputFinalizer finalizer = this.structuredOutputFinalizer; if (finalizer == null) { - finalizer = new StructuredOutputFinalizer(); + finalizer = this.chatClient != null + ? new StructuredOutputFinalizer(this.chatClient) + : new StructuredOutputFinalizer(); this.structuredOutputFinalizer = finalizer; } return finalizer; diff --git a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizer.java b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizer.java index c16eba90d..ea2e0e70c 100644 --- a/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizer.java +++ b/argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizer.java @@ -17,6 +17,8 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; import com.networknt.schema.Error; import com.networknt.schema.Schema; @@ -24,7 +26,13 @@ import com.networknt.schema.SpecificationVersion; import org.springframework.util.StringUtils; +import io.github.agentic.ai.api.chat.ChatClient; import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ChatRequest; +import io.github.agentic.ai.api.chat.ChatResponse; +import io.github.agentic.ai.api.chat.ChatRole; +import io.github.agentic.ai.api.chat.GenerationOptions; +import io.github.agentic.ai.api.chat.ToolChoice; import io.github.agentic.ai.api.chat.Usage; import io.github.agentic.ai.graph.agent.AgentResult; import tools.jackson.core.StreamReadFeature; @@ -37,7 +45,7 @@ * and schema-validating the configured output schema. This is the single delivery * boundary shared by every termination path (normal final message, returnDirect * synthesis, hook {@code JumpTo.end}): the terminal message itself is carried through - * unchanged and never regenerated. + * unchanged and never regenerated by the validation itself. * *

* Prompt-level schema constraints do not count as validation: without this component @@ -52,6 +60,14 @@ * schema violations) becomes {@link AgentResult.Status#VALIDATION_FAILED} with issue * entries that do not echo the full output content. *

+ * + *

+ * When constructed with a {@link ChatClient}, a failed validation gets exactly one + * bounded repair attempt: a single-turn model re-ask that carries no tools (and never + * touches the graph), after which the repaired output is re-validated without a second + * repair. A failed repair is marked with {@link #ERROR_CODE_REPAIRED_STILL_INVALID}; + * a failing repair client leaves the original failure untouched. + *

*/ public final class StructuredOutputFinalizer { @@ -64,6 +80,23 @@ public final class StructuredOutputFinalizer { private static final String ROOT_PATH = "$"; + private final ChatClient repairClient; + + /** Creates a pure validating finalizer without repair capability. */ + public StructuredOutputFinalizer() { + this(null); + } + + /** + * Creates a finalizer that performs one bounded repair re-ask through the given + * client when validation fails. + * @param repairClient the client used for the single repair re-ask; the request + * carries no tools and never touches the graph + */ + public StructuredOutputFinalizer(ChatClient repairClient) { + this.repairClient = repairClient; + } + /** * Finalizes the terminal message of a run against the configured output schema. * @param terminal the terminal assistant message; must not be {@code null} @@ -87,6 +120,28 @@ public AgentResult finalizeOutput(ChatMessage terminal, String outputSchema, Usa ROOT_PATH, "Configured output schema is not a valid JSON Schema"))); } + AgentResult first = validateAgainst(terminal, schema, usage); + if (first.status() != AgentResult.Status.VALIDATION_FAILED || this.repairClient == null) { + return first; + } + + ChatMessage repaired = attemptRepair(terminal, outputSchema, first.issues()); + if (repaired == null) { + // The repair client failed: keep the original failure untouched. + return first; + } + AgentResult second = validateAgainst(repaired, schema, usage); + if (second.status() != AgentResult.Status.VALIDATION_FAILED) { + return second; + } + List marked = new ArrayList<>(); + marked.add(new AgentResult.OutputValidationIssue(AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID, ROOT_PATH, + "One bounded repair attempt did not produce a schema-valid output")); + marked.addAll(second.issues()); + return AgentResult.validationFailed(repaired, usage, List.copyOf(marked)); + } + + private AgentResult validateAgainst(ChatMessage terminal, Schema schema, Usage usage) { JsonNode root; try { root = MAPPER.readTree(terminal.text()); @@ -110,4 +165,33 @@ public AgentResult finalizeOutput(ChatMessage terminal, String outputSchema, Usa return AgentResult.completed(terminal, usage, root); } + /** + * One bounded repair re-ask: single-turn, no tools, no graph. Returns the repaired + * assistant message, or {@code null} when the repair client itself failed. + */ + private ChatMessage attemptRepair(ChatMessage terminal, String outputSchema, + List issues) { + try { + String issueSummary = issues.stream() + .map(issue -> issue.errorCode() + + (StringUtils.hasLength(issue.fieldPath()) && !ROOT_PATH.equals(issue.fieldPath()) + ? " at " + issue.fieldPath() : "")) + .collect(Collectors.joining("; ")); + ChatMessage system = ChatMessage + .system("Return only valid JSON matching this JSON Schema without any surrounding text: " + + outputSchema); + ChatMessage user = ChatMessage.user("The previous response failed schema validation (" + issueSummary + + "). Return the corrected JSON only.\nPrevious response:\n" + terminal.text()); + ChatResponse response = this.repairClient.callBlocking(new ChatRequest(List.of(system, user), List.of(), + GenerationOptions.builder().toolChoice(ToolChoice.none()).build(), Map.of())); + return response.messages().stream() + .filter(message -> message.role() == ChatRole.ASSISTANT) + .findFirst() + .orElse(null); + } + catch (RuntimeException ex) { + return null; + } + } + } diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentCallResultTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentCallResultTest.java index 519a1dcd0..4ce305678 100644 --- a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentCallResultTest.java +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentCallResultTest.java @@ -73,8 +73,11 @@ void invalidOutputFailsValidation() { AgentResult result = agent.callResult("give me the status"); assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED); - assertThat(result.issues()).hasSize(1); - assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON); + // The stub answers the bounded repair with the same invalid text, so the + // REPAIRED_STILL_INVALID marker is prepended to the original issue. + assertThat(result.issues()).hasSize(2); + assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID); + assertThat(result.issues().get(1).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON); } @Test @@ -85,8 +88,9 @@ void schemaViolationReported() { AgentResult result = agent.callResult("give me the status"); assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED); - assertThat(result.issues()).isNotEmpty(); - assertThat(result.issues()).allMatch(issue -> AgentResult.ERROR_CODE_SCHEMA_VIOLATION.equals(issue.errorCode())); + // Repair re-asks with the same payload: the marker precedes the actual violations. + assertThat(result.issues()).anyMatch(issue -> AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID.equals(issue.errorCode())); + assertThat(result.issues()).anyMatch(issue -> AgentResult.ERROR_CODE_SCHEMA_VIOLATION.equals(issue.errorCode())); } @Test diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentStreamResultTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentStreamResultTest.java new file mode 100644 index 000000000..5f4943f6c --- /dev/null +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentStreamResultTest.java @@ -0,0 +1,189 @@ +/* + * 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 + * + * https://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.github.agentic.ai.graph.agent; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicInteger; + +import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatClientCapabilities; +import io.github.agentic.ai.api.chat.ChatEvent; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ChatRequest; +import io.github.agentic.ai.api.chat.ChatResponse; +import io.github.agentic.ai.api.chat.FinishReason; +import io.github.agentic.ai.api.chat.Usage; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end tests for the streaming delivery contract {@link ReactAgent#streamResult} + * with a scripted streaming client: provisional chunks carry a {@code null} status, the + * final element is the validated delivery of the terminal message, and a validation + * failure gets exactly one bounded repair re-ask. + */ +@DisplayName("ReactAgent.streamResult streaming delivery contract") +class ReactAgentStreamResultTest { + + private static final String SCHEMA = """ + { + "type": "object", + "required": ["status"], + "additionalProperties": false, + "properties": { + "status": {"type": "string"} + } + }"""; + + @Test + @DisplayName("provisional chunks are unvalidated; the final element carries the validated delivery") + void validStreamDeliversValidatedResult() { + ScriptedStreamingClient client = new ScriptedStreamingClient( + ChatEvent.textDelta("{\"status\":"), + ChatEvent.textDelta("\"ok\"}"), + ChatEvent.response(finalResponse("{\"status\":\"ok\"}", new Usage(2L, 3L, 5L)))); + ReactAgent agent = agent(client, SCHEMA); + + List results = agent.streamResult("give me the status").collectList().block(); + + assertThat(results).isNotNull(); + assertThat(results).hasSize(3); + assertThat(results.get(0).status()).isNull(); + assertThat(results.get(0).message().text()).isEqualTo("{\"status\":"); + assertThat(results.get(1).status()).isNull(); + AgentResult delivery = results.get(2); + assertThat(delivery.status()).isEqualTo(AgentResult.Status.COMPLETED); + assertThat(delivery.data()).isNotNull(); + assertThat(delivery.usage().totalTokens()).isEqualTo(5L); + } + + @Test + @DisplayName("an invalid stream output is repaired once and delivered as COMPLETED") + void invalidStreamOutputIsRepaired() { + ScriptedStreamingClient client = new ScriptedStreamingClient( + ChatEvent.textDelta("plain prose"), + ChatEvent.response(finalResponse("plain prose", new Usage(2L, 3L, 5L)))); + client.repairedText = "{\"status\":\"ok\"}"; + ReactAgent agent = agent(client, SCHEMA); + + List results = agent.streamResult("give me the status").collectList().block(); + + assertThat(results).isNotNull(); + AgentResult delivery = results.get(results.size() - 1); + assertThat(delivery.status()).isEqualTo(AgentResult.Status.COMPLETED); + assertThat(delivery.message().text()).isEqualTo("{\"status\":\"ok\"}"); + assertThat(client.repairCalls.get()).isEqualTo(1); + } + + @Test + @DisplayName("a stream output that still fails validation is reported, not completed") + void invalidStreamOutputWithoutRepairFails() { + ScriptedStreamingClient client = new ScriptedStreamingClient( + ChatEvent.textDelta("plain prose"), + ChatEvent.response(finalResponse("plain prose", new Usage(2L, 3L, 5L)))); + ReactAgent agent = agent(client, SCHEMA); + + List results = agent.streamResult("give me the status").collectList().block(); + + assertThat(results).isNotNull(); + AgentResult delivery = results.get(results.size() - 1); + assertThat(delivery.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED); + assertThat(delivery.issues()).isNotEmpty(); + } + + @Test + @DisplayName("the legacy streamMessages behavior stays untouched") + void legacyStreamMessagesUnchanged() throws Exception { + ScriptedStreamingClient client = new ScriptedStreamingClient( + ChatEvent.textDelta("{\"status\":"), + ChatEvent.textDelta("\"ok\"}"), + ChatEvent.response(finalResponse("{\"status\":\"ok\"}", new Usage(2L, 3L, 5L)))); + ReactAgent agent = agent(client, SCHEMA); + + List messages = agent.streamMessages("give me the status").collectList().block(); + + assertThat(messages).isNotNull(); + // streamMessages publishes the text-delta chunks; the aggregated final block is + // consumed by streamResult, which is exactly what the delivery contract adds. + String combined = messages.stream().map(ChatMessage::text).reduce("", String::concat); + assertThat(combined).isEqualTo("{\"status\":\"ok\"}"); + } + + private static ChatResponse finalResponse(String text, Usage usage) { + return new ChatResponse("response-id", "test-model", List.of(ChatMessage.assistant(text)), List.of(), + FinishReason.STOP, usage, Map.of("argiFinalResponse", true)); + } + + private static ReactAgent agent(ChatClient client, String outputSchema) { + Builder builder = ReactAgent.builder() + .name("stream-result-test") + .chatClient(client); + if (outputSchema != null) { + builder.outputSchema(outputSchema); + } + return builder.build(); + } + + /** + * Streams the scripted events and answers the bounded repair re-ask (a synchronous + * {@code call}) with {@link #repairedText}, counting repair invocations. + */ + private static final class ScriptedStreamingClient implements ChatClient { + + private final ChatEvent[] script; + + private final AtomicInteger repairCalls = new AtomicInteger(); + + private volatile String repairedText; + + private volatile RuntimeException repairFailure; + + ScriptedStreamingClient(ChatEvent... script) { + this.script = script; + } + + @Override + public CompletionStage call(ChatRequest request) { + this.repairCalls.incrementAndGet(); + if (this.repairFailure != null) { + throw this.repairFailure; + } + return CompletableFuture.completedFuture(new ChatResponse("repair-id", "test-model", + List.of(ChatMessage.assistant(this.repairedText)), List.of(), null, + new Usage(1L, 1L, 2L), Map.of())); + } + + @Override + public Flow.Publisher stream(ChatRequest request) { + return JdkFlowAdapter.publisherToFlowPublisher(Flux.fromArray(this.script)); + } + + @Override + public ChatClientCapabilities capabilities() { + return new ChatClientCapabilities(true, true, false, false, true, true); + } + + } + +} diff --git a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizerTest.java b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizerTest.java index 294130b3d..eaefa306a 100644 --- a/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizerTest.java +++ b/argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizerTest.java @@ -21,6 +21,16 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicInteger; + +import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatEvent; +import io.github.agentic.ai.api.chat.ChatClientCapabilities; +import io.github.agentic.ai.api.chat.ChatRequest; +import io.github.agentic.ai.api.chat.ChatResponse; + import static org.assertj.core.api.Assertions.assertThat; /** @@ -121,4 +131,81 @@ void trailingTokensFailValidation() { assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON); } + @Test + @DisplayName("one bounded repair attempt turns an invalid output into a completed delivery") + void repairProducesValidOutput() { + RepairStub client = new RepairStub("{\"status\":\"ok\"}", null); + StructuredOutputFinalizer repairing = new StructuredOutputFinalizer(client); + + AgentResult result = repairing.finalizeOutput(ChatMessage.assistant("not json"), SCHEMA, Usage.unknown()); + + assertThat(result.status()).isEqualTo(AgentResult.Status.COMPLETED); + assertThat(result.data()).isNotNull(); + assertThat(result.message().text()).isEqualTo("{\"status\":\"ok\"}"); + assertThat(client.calls.get()).isEqualTo(1); + } + + @Test + @DisplayName("a repair that still fails is marked and never retried a second time") + void repairStillInvalidIsMarkedAndNotRetried() { + RepairStub client = new RepairStub("still not json", null); + StructuredOutputFinalizer repairing = new StructuredOutputFinalizer(client); + + AgentResult result = repairing.finalizeOutput(ChatMessage.assistant("not json"), SCHEMA, Usage.unknown()); + + assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED); + assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID); + assertThat(client.calls.get()).isEqualTo(1); + } + + @Test + @DisplayName("a failing repair client keeps the original validation failure") + void repairClientFailureKeepsOriginalFailure() { + RepairStub client = new RepairStub(null, new IllegalStateException("repair backend down")); + StructuredOutputFinalizer repairing = new StructuredOutputFinalizer(client); + + AgentResult result = repairing.finalizeOutput(ChatMessage.assistant("not json"), SCHEMA, Usage.unknown()); + + assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED); + assertThat(result.issues()).hasSize(1); + assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON); + assertThat(client.calls.get()).isEqualTo(1); + } + + private static final class RepairStub implements ChatClient { + + private final AtomicInteger calls = new AtomicInteger(); + + private final String repairedText; + + private final RuntimeException failure; + + RepairStub(String repairedText, RuntimeException failure) { + this.repairedText = repairedText; + this.failure = failure; + } + + @Override + public java.util.concurrent.CompletionStage call(ChatRequest request) { + this.calls.incrementAndGet(); + if (this.failure != null) { + throw this.failure; + } + return CompletableFuture.completedFuture(new ChatResponse("repair-id", "test-model", + java.util.List.of(ChatMessage.assistant(this.repairedText)), java.util.List.of(), null, + new Usage(1L, 1L, 2L), java.util.Map.of())); + } + + @Override + public Flow.Publisher stream(ChatRequest request) { + throw new UnsupportedOperationException("Streaming is unavailable"); + } + + @Override + public ChatClientCapabilities capabilities() { + return ChatClientCapabilities.basic(); + } + + } + }