diff --git a/sdk/ai/azure-ai-agents/CHANGELOG.md b/sdk/ai/azure-ai-agents/CHANGELOG.md index 5e275d86ca04..f510e7d90b5d 100644 --- a/sdk/ai/azure-ai-agents/CHANGELOG.md +++ b/sdk/ai/azure-ai-agents/CHANGELOG.md @@ -8,6 +8,11 @@ - Added new `OpenApiSync`/`OpenApiAsync` samples demonstrating end-to-end OpenAPI tool integration: loading a spec file, creating an agent with an `OpenApiTool`, and invoking an external API via conversation. - Added new tool samples for parity with the Python SDK: `AzureFunctionSync`/`AzureFunctionAsync`, `BingCustomSearchSync`/`BingCustomSearchAsync`, `MemorySearchSync`/`MemorySearchAsync`, `McpWithConnectionSync`/`McpWithConnectionAsync`, and `OpenApiWithConnectionSync`/`OpenApiWithConnectionAsync`. - Added `setComparisonFilter(ComparisonFilter)` and `setCompoundFilter(CompoundFilter)` convenience methods on `FileSearchTool`, accepting the openai-java filter types directly. +- Added streaming response methods to `ResponsesClient` and `ResponsesAsyncClient`: + - `createStreamingWithAgent` and `createStreamingWithAgentConversation` on `ResponsesClient` return `IterableStream` for synchronous streaming. + - `createStreamingWithAgent` and `createStreamingWithAgentConversation` on `ResponsesAsyncClient` return `Flux` for asynchronous streaming. +- Added `StreamingUtils` implementation helper that bridges OpenAI `StreamResponse` to `IterableStream` and `AsyncStreamResponse` to `Flux`. +- Added streaming samples: `SimpleStreamingSync`/`SimpleStreamingAsync`, `FunctionCallStreamingSync`/`FunctionCallStreamingAsync`, and `CodeInterpreterStreamingSync`/`CodeInterpreterStreamingAsync`. ### Breaking Changes @@ -21,6 +26,7 @@ ### Other Changes - Added `ToolsTests` and `ToolsAsyncTests` with recorded end-to-end test coverage for OpenAPI, Code Interpreter, Function Call, Web Search, MCP, and File Search tools. +- Added `StreamingTests` and `StreamingAsyncTests` with recorded test coverage for streaming responses (simple prompt, function calling, and Code Interpreter scenarios). ## 2.0.0-beta.2 (2026-03-04) diff --git a/sdk/ai/azure-ai-agents/README.md b/sdk/ai/azure-ai-agents/README.md index 79421f811241..1220f6ee0834 100644 --- a/sdk/ai/azure-ai-agents/README.md +++ b/sdk/ai/azure-ai-agents/README.md @@ -562,6 +562,67 @@ See the full sample in [OpenApiWithConnectionSync.java](https://github.com/Azure --- +### Streaming responses + +The `ResponsesClient` and `ResponsesAsyncClient` support streaming, which allows you to process response events as they arrive rather than waiting for the full response. This is useful for displaying text to users in real time and observing tool execution progress. + +#### Synchronous streaming + +The synchronous streaming methods return `IterableStream`, which can be consumed with a standard for-each loop. Use the `ResponseAccumulator` from the OpenAI SDK to collect events into a final `Response`: + +```java com.azure.ai.agents.streaming.simple_sync +// Use ResponseAccumulator to collect streamed events into a final Response +ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + +// Stream response - text is printed as it arrives +IterableStream events = + responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Tell me a short story about a brave explorer.")); + +for (ResponseStreamEvent event : events) { + responseAccumulator.accumulate(event); + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); +} +System.out.println(); // newline after streamed text + +// Access the complete accumulated response +Response response = responseAccumulator.response(); +System.out.println("\nResponse ID: " + response.id()); +``` + +See the full samples in [SimpleStreamingSync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingSync.java), [FunctionCallStreamingSync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingSync.java), and [CodeInterpreterStreamingSync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingSync.java). + +#### Asynchronous streaming + +The asynchronous streaming methods return `Flux`, integrating naturally with Reactor pipelines: + +```java com.azure.ai.agents.streaming.simple_async +// Use ResponseAccumulator to collect streamed events into a final Response +ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + +// Stream response asynchronously - text is printed as each chunk arrives +return responsesAsyncClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Tell me a short story about a brave explorer.")) + .doOnNext(event -> { + responseAccumulator.accumulate(event); + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); + }) + .then(Mono.fromCallable(() -> { + System.out.println(); // newline after streamed text + + // Access the complete accumulated response + Response response = responseAccumulator.response(); + System.out.println("\nResponse ID: " + response.id()); +``` + +See the full samples in [SimpleStreamingAsync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingAsync.java), [FunctionCallStreamingAsync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingAsync.java), and [CodeInterpreterStreamingAsync.java](https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingAsync.java). + +--- + ### Service API versions The client library targets the latest service API version by default. diff --git a/sdk/ai/azure-ai-agents/assets.json b/sdk/ai/azure-ai-agents/assets.json index 321243e18338..f30fca7a76ff 100644 --- a/sdk/ai/azure-ai-agents/assets.json +++ b/sdk/ai/azure-ai-agents/assets.json @@ -2,5 +2,5 @@ "AssetsRepo": "Azure/azure-sdk-assets", "AssetsRepoPrefixPath": "java", "TagPrefix": "java/ai/azure-ai-agents", - "Tag": "java/ai/azure-ai-agents_25e4920008" + "Tag": "java/ai/azure-ai-agents_4b3a190a4f" } diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesAsyncClient.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesAsyncClient.java index fe3258235e15..46847f198d95 100644 --- a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesAsyncClient.java +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesAsyncClient.java @@ -5,13 +5,16 @@ package com.azure.ai.agents; import com.azure.ai.agents.implementation.OpenAIJsonHelper; +import com.azure.ai.agents.implementation.StreamingUtils; import com.azure.ai.agents.models.AgentReference; import com.azure.core.annotation.ServiceClient; import com.openai.client.OpenAIClientAsync; import com.openai.core.JsonValue; import com.openai.models.responses.Response; import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseStreamEvent; import com.openai.services.async.ResponseServiceAsync; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.util.HashMap; @@ -107,4 +110,49 @@ public Mono createWithAgent(AgentReference agentReference, ResponseCre public Mono createWithAgent(AgentReference agentReference) { return createWithAgent(agentReference, new ResponseCreateParams.Builder()); } + + /** + * Creates a streaming response with an agent. + * + * @param agentReference The agent reference. + * @param params The parameters to create the response. + * @return A Flux of ResponseStreamEvent. + */ + public Flux createStreamingWithAgent(AgentReference agentReference, + ResponseCreateParams.Builder params) { + Objects.requireNonNull(agentReference, "agentReference cannot be null"); + Objects.requireNonNull(params, "params cannot be null"); + + JsonValue agentRefJsonValue = OpenAIJsonHelper.toJsonValue(agentReference); + + Map additionalBodyProperties = new HashMap<>(); + additionalBodyProperties.put("agent_reference", agentRefJsonValue); + + params.additionalBodyProperties(additionalBodyProperties); + return StreamingUtils.toFlux(this.responseServiceAsync.createStreaming(params.build())); + } + + /** + * Creates a streaming response with an agent conversation. + * + * @param agentReference The agent reference. + * @param conversationId The conversation ID. + * @param params The parameters to create the response. + * @return A Flux of ResponseStreamEvent. + */ + public Flux createStreamingWithAgentConversation(AgentReference agentReference, + String conversationId, ResponseCreateParams.Builder params) { + Objects.requireNonNull(agentReference, "agentReference cannot be null"); + Objects.requireNonNull(conversationId, "conversationId cannot be null"); + Objects.requireNonNull(params, "params cannot be null"); + + JsonValue agentRefJsonValue = OpenAIJsonHelper.toJsonValue(agentReference); + + Map additionalBodyProperties = new HashMap<>(); + params.conversation(conversationId); + additionalBodyProperties.put("agent_reference", agentRefJsonValue); + + params.additionalBodyProperties(additionalBodyProperties); + return StreamingUtils.toFlux(this.responseServiceAsync.createStreaming(params.build())); + } } diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesClient.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesClient.java index 08ead01962a7..a93aa57e7109 100644 --- a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesClient.java +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/ResponsesClient.java @@ -5,12 +5,15 @@ package com.azure.ai.agents; import com.azure.ai.agents.implementation.OpenAIJsonHelper; +import com.azure.ai.agents.implementation.StreamingUtils; import com.azure.ai.agents.models.AgentReference; import com.azure.core.annotation.ServiceClient; +import com.azure.core.util.IterableStream; import com.openai.client.OpenAIClient; import com.openai.core.JsonValue; import com.openai.models.responses.Response; import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseStreamEvent; import com.openai.services.blocking.ResponseService; import java.util.HashMap; @@ -106,4 +109,49 @@ public Response createWithAgent(AgentReference agentReference, ResponseCreatePar public Response createWithAgent(AgentReference agentReference) { return createWithAgent(agentReference, new ResponseCreateParams.Builder()); } + + /** + * Creates a streaming response with an agent. + * + * @param agentReference The agent reference. + * @param params The parameters to create the response. + * @return An IterableStream of ResponseStreamEvent. + */ + public IterableStream createStreamingWithAgent(AgentReference agentReference, + ResponseCreateParams.Builder params) { + Objects.requireNonNull(agentReference, "agentReference cannot be null"); + Objects.requireNonNull(params, "params cannot be null"); + + JsonValue agentRefJsonValue = OpenAIJsonHelper.toJsonValue(agentReference); + + Map additionalBodyProperties = new HashMap<>(); + additionalBodyProperties.put("agent_reference", agentRefJsonValue); + + params.additionalBodyProperties(additionalBodyProperties); + return StreamingUtils.toIterableStream(this.responseService.createStreaming(params.build())); + } + + /** + * Creates a streaming response with an agent conversation. + * + * @param agentReference The agent reference. + * @param conversationId The conversation ID. + * @param params The parameters to create the response. + * @return An IterableStream of ResponseStreamEvent. + */ + public IterableStream createStreamingWithAgentConversation(AgentReference agentReference, + String conversationId, ResponseCreateParams.Builder params) { + Objects.requireNonNull(agentReference, "agentReference cannot be null"); + Objects.requireNonNull(conversationId, "conversationId cannot be null"); + Objects.requireNonNull(params, "params cannot be null"); + + JsonValue agentRefJsonValue = OpenAIJsonHelper.toJsonValue(agentReference); + + Map additionalBodyProperties = new HashMap<>(); + params.conversation(conversationId); + additionalBodyProperties.put("agent_reference", agentRefJsonValue); + + params.additionalBodyProperties(additionalBodyProperties); + return StreamingUtils.toIterableStream(this.responseService.createStreaming(params.build())); + } } diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/StreamingUtils.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/StreamingUtils.java new file mode 100644 index 000000000000..83e7b28d7d7d --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/StreamingUtils.java @@ -0,0 +1,105 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.implementation; + +import com.azure.core.util.IterableStream; +import com.azure.core.util.logging.ClientLogger; +import com.openai.core.http.AsyncStreamResponse; +import com.openai.core.http.StreamResponse; +import reactor.core.publisher.Flux; + +import java.util.Iterator; +import java.util.NoSuchElementException; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Utility methods for bridging OpenAI streaming types to Azure SDK / Reactor types. + */ +public final class StreamingUtils { + + private static final ClientLogger LOGGER = new ClientLogger(StreamingUtils.class); + + private StreamingUtils() { + } + + /** + * Converts an OpenAI {@link StreamResponse} into an {@link IterableStream} that automatically + * closes the underlying stream resource when all events have been consumed. + * + *

The returned {@link IterableStream} is single-use: calling {@code iterator()} or + * {@code stream()} more than once will throw {@link IllegalStateException}, because the + * underlying {@link StreamResponse} can only be consumed once.

+ * + * @param streamResponse The OpenAI stream response. + * @param The type of the streamed items. + * @return An IterableStream wrapping the given StreamResponse. + */ + public static IterableStream toIterableStream(StreamResponse streamResponse) { + Iterator inner = streamResponse.stream().iterator(); + AtomicBoolean iteratorCreated = new AtomicBoolean(false); + + Iterable singleUseIterable = () -> { + if (!iteratorCreated.compareAndSet(false, true)) { + throw LOGGER + .logExceptionAsError(new IllegalStateException("This streaming response has already been consumed. " + + "A streaming IterableStream can only be iterated once.")); + } + return new Iterator() { + @Override + public boolean hasNext() { + boolean hasNext = inner.hasNext(); + if (!hasNext) { + streamResponse.close(); + } + return hasNext; + } + + @Override + public T next() { + if (!inner.hasNext()) { + throw LOGGER.logExceptionAsError(new NoSuchElementException()); + } + return inner.next(); + } + }; + }; + return new IterableStream<>(singleUseIterable); + } + + /** + * Bridges an OpenAI {@link AsyncStreamResponse} into a Reactor {@link Flux}, forwarding + * events, errors, and completion signals while closing the underlying stream on disposal. + * + *

Thread safety: {@code Flux.create} defaults to {@code PUSH_PULL} mode, which wraps the + * sink in a {@code SerializedFluxSink} that serializes {@code next/error/complete} calls. + * Additionally, the OpenAI {@code AsyncStreamResponse} implementation invokes the handler + * sequentially on a single executor thread (via {@code Stream.forEach}), so concurrent + * emission cannot occur in practice.

+ * + * @param asyncStream The OpenAI async stream response. + * @param The type of the streamed items. + * @return A Flux wrapping the given AsyncStreamResponse. + */ + public static Flux toFlux(AsyncStreamResponse asyncStream) { + return Flux.create(sink -> { + asyncStream.subscribe(new AsyncStreamResponse.Handler() { + @Override + public void onNext(T event) { + sink.next(event); + } + + @Override + public void onComplete(Optional error) { + if (error.isPresent()) { + sink.error(error.get()); + } else { + sink.complete(); + } + } + }); + sink.onDispose(asyncStream::close); + }); + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingAsync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingAsync.java new file mode 100644 index 000000000000..0a28716d75ea --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingAsync.java @@ -0,0 +1,105 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsAsyncClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesAsyncClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.CodeInterpreterTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.Configuration; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import reactor.core.publisher.Mono; + +import java.util.Collections; +import java.util.concurrent.atomic.AtomicReference; + +/** + * This sample demonstrates how to stream a response from an agent configured with the + * Azure-specific Code Interpreter tool using the asynchronous client. Code execution + * progress events and text output are printed as they arrive. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class CodeInterpreterStreamingAsync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsAsyncClient agentsAsyncClient = builder.buildAgentsAsyncClient(); + ResponsesAsyncClient responsesAsyncClient = builder.buildResponsesAsyncClient(); + + AtomicReference agentRef = new AtomicReference<>(); + + // Create a CodeInterpreterTool - an Azure-specific tool for executing Python code + CodeInterpreterTool tool = new CodeInterpreterTool(); + + // Create agent with Code Interpreter tool + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that can execute Python code to solve problems. " + + "When asked to perform calculations, use the code interpreter to run Python code.") + .setTools(Collections.singletonList(tool)); + + agentsAsyncClient.createAgentVersion("code-interpreter-streaming-async-agent", agentDefinition) + .flatMap(agent -> { + agentRef.set(agent); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.code_interpreter_async + // Stream response asynchronously with Code Interpreter + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + return responsesAsyncClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Calculate the first 10 prime numbers using Python.")) + .doOnNext(event -> { + responseAccumulator.accumulate(event); + // Print text deltas as they arrive + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); + // Observe code interpreter progress events + event.codeInterpreterCallInProgress() + .ifPresent(e -> System.out.println("\n[Code interpreter running...]")); + event.codeInterpreterCallCodeDelta() + .ifPresent(e -> System.out.print(e.delta())); + event.codeInterpreterCallCompleted() + .ifPresent(e -> System.out.println("\n[Code interpreter completed]")); + }) + .then(Mono.fromCallable(() -> { + System.out.println(); + + // Access the complete accumulated response + Response response = responseAccumulator.response(); + System.out.println("\nResponse ID: " + response.id()); + // END: com.azure.ai.agents.streaming.code_interpreter_async + return response; + })); + }) + .then(Mono.defer(() -> { + AgentVersionDetails agent = agentRef.get(); + if (agent != null) { + return agentsAsyncClient.deleteAgentVersion(agent.getName(), agent.getVersion()) + .doOnSuccess(v -> System.out.println("Agent deleted")); + } + return Mono.empty(); + })) + .block(); + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingSync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingSync.java new file mode 100644 index 000000000000..d50082703891 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/CodeInterpreterStreamingSync.java @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.CodeInterpreterTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.Configuration; +import com.azure.core.util.IterableStream; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseStreamEvent; + +import java.util.Collections; + +/** + * This sample demonstrates how to stream a response from an agent configured with the + * Azure-specific Code Interpreter tool. Code execution progress events and text output + * are printed as they arrive. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class CodeInterpreterStreamingSync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsClient agentsClient = builder.buildAgentsClient(); + ResponsesClient responsesClient = builder.buildResponsesClient(); + + AgentVersionDetails agent = null; + + try { + // Create a CodeInterpreterTool - an Azure-specific tool for executing Python code + CodeInterpreterTool tool = new CodeInterpreterTool(); + + // Create agent with Code Interpreter tool + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that can execute Python code to solve problems. " + + "When asked to perform calculations, use the code interpreter to run Python code.") + .setTools(Collections.singletonList(tool)); + + agent = agentsClient.createAgentVersion("code-interpreter-streaming-agent", agentDefinition); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.code_interpreter_sync + // Stream response with Code Interpreter - observe code execution events as they arrive + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + IterableStream events = + responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Calculate the first 10 prime numbers using Python.")); + + for (ResponseStreamEvent event : events) { + responseAccumulator.accumulate(event); + // Print text deltas as they stream in + event.outputTextDelta().ifPresent(textEvent -> + System.out.print(textEvent.delta())); + + // Observe code interpreter progress events + event.codeInterpreterCallInProgress().ifPresent(e -> + System.out.println("\n[Code interpreter running...]")); + event.codeInterpreterCallCodeDelta().ifPresent(e -> + System.out.print(e.delta())); + event.codeInterpreterCallCompleted().ifPresent(e -> + System.out.println("\n[Code interpreter completed]")); + } + System.out.println(); + + // Access the complete accumulated response + Response response = responseAccumulator.response(); + System.out.println("\nResponse ID: " + response.id()); + // END: com.azure.ai.agents.streaming.code_interpreter_sync + } finally { + if (agent != null) { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + System.out.println("Agent deleted"); + } + } + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingAsync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingAsync.java new file mode 100644 index 000000000000..a8f968042795 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingAsync.java @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsAsyncClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesAsyncClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.FunctionTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.BinaryData; +import com.azure.core.util.Configuration; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseOutputItem; +import reactor.core.publisher.Mono; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +/** + * This sample demonstrates how to stream a response from an agent configured with a + * Function Calling tool using the asynchronous client. Function call arguments and text + * output are printed as they arrive. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class FunctionCallStreamingAsync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsAsyncClient agentsAsyncClient = builder.buildAgentsAsyncClient(); + ResponsesAsyncClient responsesAsyncClient = builder.buildResponsesAsyncClient(); + + AtomicReference agentRef = new AtomicReference<>(); + + // Define a function tool with parameter schema + Map locationProp = new LinkedHashMap<>(); + locationProp.put("type", "string"); + locationProp.put("description", "The city and state, e.g. Seattle, WA"); + + Map unitProp = new LinkedHashMap<>(); + unitProp.put("type", "string"); + unitProp.put("enum", Arrays.asList("celsius", "fahrenheit")); + + Map properties = new LinkedHashMap<>(); + properties.put("location", locationProp); + properties.put("unit", unitProp); + + Map parameters = new HashMap<>(); + parameters.put("type", BinaryData.fromObject("object")); + parameters.put("properties", BinaryData.fromObject(properties)); + parameters.put("required", BinaryData.fromObject(Arrays.asList("location", "unit"))); + parameters.put("additionalProperties", BinaryData.fromObject(false)); + + FunctionTool tool = new FunctionTool("get_weather", parameters, true) + .setDescription("Get the current weather in a given location"); + + // Create agent with function tool + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that can get weather information. " + + "When asked about the weather, use the get_weather function.") + .setTools(Collections.singletonList(tool)); + + agentsAsyncClient.createAgentVersion("function-streaming-async-agent", agentDefinition) + .flatMap(agent -> { + agentRef.set(agent); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.function_call_async + // Stream response asynchronously with function tool + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + return responsesAsyncClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("What's the weather like in Seattle?")) + .doOnNext(event -> { + responseAccumulator.accumulate(event); + // Print text deltas as they arrive + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); + // Print function call argument deltas as they arrive + event.functionCallArgumentsDelta() + .ifPresent(argEvent -> System.out.print(argEvent.delta())); + }) + .then(Mono.fromCallable(() -> { + System.out.println(); + + // Access the final response and inspect function calls + Response response = responseAccumulator.response(); + for (ResponseOutputItem outputItem : response.output()) { + outputItem.functionCall().ifPresent(functionCall -> { + System.out.println("\n--- Function Tool Call ---"); + System.out.println("Call ID: " + functionCall.callId()); + System.out.println("Function Name: " + functionCall.name()); + System.out.println("Arguments: " + functionCall.arguments()); + System.out.println("Status: " + functionCall.status()); + }); + } + // END: com.azure.ai.agents.streaming.function_call_async + return response; + })); + }) + .then(Mono.defer(() -> { + AgentVersionDetails agent = agentRef.get(); + if (agent != null) { + return agentsAsyncClient.deleteAgentVersion(agent.getName(), agent.getVersion()) + .doOnSuccess(v -> System.out.println("Agent deleted")); + } + return Mono.empty(); + })) + .block(); + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingSync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingSync.java new file mode 100644 index 000000000000..b46d1eaf612c --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/FunctionCallStreamingSync.java @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.FunctionTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.BinaryData; +import com.azure.core.util.Configuration; +import com.azure.core.util.IterableStream; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseOutputItem; +import com.openai.models.responses.ResponseStreamEvent; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * This sample demonstrates how to stream a response from an agent configured with a + * Function Calling tool. Function call arguments are streamed as they arrive, followed + * by any text output. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class FunctionCallStreamingSync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsClient agentsClient = builder.buildAgentsClient(); + ResponsesClient responsesClient = builder.buildResponsesClient(); + + AgentVersionDetails agent = null; + + try { + // Define a function tool with parameter schema + Map locationProp = new LinkedHashMap<>(); + locationProp.put("type", "string"); + locationProp.put("description", "The city and state, e.g. Seattle, WA"); + + Map unitProp = new LinkedHashMap<>(); + unitProp.put("type", "string"); + unitProp.put("enum", Arrays.asList("celsius", "fahrenheit")); + + Map properties = new LinkedHashMap<>(); + properties.put("location", locationProp); + properties.put("unit", unitProp); + + Map parameters = new HashMap<>(); + parameters.put("type", BinaryData.fromObject("object")); + parameters.put("properties", BinaryData.fromObject(properties)); + parameters.put("required", BinaryData.fromObject(Arrays.asList("location", "unit"))); + parameters.put("additionalProperties", BinaryData.fromObject(false)); + + FunctionTool tool = new FunctionTool("get_weather", parameters, true) + .setDescription("Get the current weather in a given location"); + + // Create agent with function tool + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that can get weather information. " + + "When asked about the weather, use the get_weather function.") + .setTools(Collections.singletonList(tool)); + + agent = agentsClient.createAgentVersion("function-streaming-agent", agentDefinition); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.function_call_sync + // Stream response with function tool - observe function call arguments and text as they arrive + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + IterableStream events = + responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("What's the weather like in Seattle?")); + + for (ResponseStreamEvent event : events) { + responseAccumulator.accumulate(event); + // Print text deltas as they stream in + event.outputTextDelta().ifPresent(textEvent -> + System.out.print(textEvent.delta())); + // Print function call argument deltas as they stream in + event.functionCallArgumentsDelta().ifPresent(argEvent -> + System.out.print(argEvent.delta())); + } + System.out.println(); + + // Access the final response and inspect any function calls + Response response = responseAccumulator.response(); + for (ResponseOutputItem outputItem : response.output()) { + outputItem.functionCall().ifPresent(functionCall -> { + System.out.println("\n--- Function Tool Call ---"); + System.out.println("Call ID: " + functionCall.callId()); + System.out.println("Function Name: " + functionCall.name()); + System.out.println("Arguments: " + functionCall.arguments()); + System.out.println("Status: " + functionCall.status()); + }); + } + // END: com.azure.ai.agents.streaming.function_call_sync + } finally { + if (agent != null) { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + System.out.println("Agent deleted"); + } + } + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingAsync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingAsync.java new file mode 100644 index 000000000000..79bab5958718 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingAsync.java @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsAsyncClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesAsyncClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.Configuration; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import reactor.core.publisher.Mono; + +import java.util.concurrent.atomic.AtomicReference; + +/** + * This sample demonstrates how to create a streaming response using the asynchronous client. + * Text is printed as it arrives rather than waiting for the full response. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class SimpleStreamingAsync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsAsyncClient agentsAsyncClient = builder.buildAgentsAsyncClient(); + ResponsesAsyncClient responsesAsyncClient = builder.buildResponsesAsyncClient(); + + AtomicReference agentRef = new AtomicReference<>(); + + // Create an agent + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that tells short, engaging stories."); + + agentsAsyncClient.createAgentVersion("streaming-async-agent", agentDefinition) + .flatMap(agent -> { + agentRef.set(agent); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.simple_async + // Use ResponseAccumulator to collect streamed events into a final Response + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + // Stream response asynchronously - text is printed as each chunk arrives + return responsesAsyncClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Tell me a short story about a brave explorer.")) + .doOnNext(event -> { + responseAccumulator.accumulate(event); + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); + }) + .then(Mono.fromCallable(() -> { + System.out.println(); // newline after streamed text + + // Access the complete accumulated response + Response response = responseAccumulator.response(); + System.out.println("\nResponse ID: " + response.id()); + // END: com.azure.ai.agents.streaming.simple_async + return response; + })); + }) + .then(Mono.defer(() -> { + AgentVersionDetails agent = agentRef.get(); + if (agent != null) { + return agentsAsyncClient.deleteAgentVersion(agent.getName(), agent.getVersion()) + .doOnSuccess(v -> System.out.println("Agent deleted")); + } + return Mono.empty(); + })) + .block(); + } +} diff --git a/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingSync.java b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingSync.java new file mode 100644 index 000000000000..1dc444187ad3 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/samples/java/com/azure/ai/agents/streaming/SimpleStreamingSync.java @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.streaming; + +import com.azure.ai.agents.AgentsClient; +import com.azure.ai.agents.AgentsClientBuilder; +import com.azure.ai.agents.ResponsesClient; +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.util.Configuration; +import com.azure.core.util.IterableStream; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseStreamEvent; + +/** + * This sample demonstrates how to create a streaming response using the synchronous client. + * Text is printed as it arrives rather than waiting for the full response. + * + *

Before running the sample, set these environment variables:

+ *
    + *
  • FOUNDRY_PROJECT_ENDPOINT - The Azure AI Project endpoint.
    • + *
  • FOUNDRY_MODEL_DEPLOYMENT_NAME - The model deployment name.
    • + *
+ */ +public class SimpleStreamingSync { + public static void main(String[] args) { + String endpoint = Configuration.getGlobalConfiguration().get("FOUNDRY_PROJECT_ENDPOINT"); + String model = Configuration.getGlobalConfiguration().get("FOUNDRY_MODEL_DEPLOYMENT_NAME"); + + AgentsClientBuilder builder = new AgentsClientBuilder() + .credential(new DefaultAzureCredentialBuilder().build()) + .endpoint(endpoint); + + AgentsClient agentsClient = builder.buildAgentsClient(); + ResponsesClient responsesClient = builder.buildResponsesClient(); + + AgentVersionDetails agent = null; + + try { + // Create an agent + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(model) + .setInstructions("You are a helpful assistant that tells short, engaging stories."); + + agent = agentsClient.createAgentVersion("streaming-agent", agentDefinition); + System.out.printf("Agent created: %s (version %s)%n", agent.getName(), agent.getVersion()); + + AgentReference agentReference = new AgentReference(agent.getName()) + .setVersion(agent.getVersion()); + + // BEGIN: com.azure.ai.agents.streaming.simple_sync + // Use ResponseAccumulator to collect streamed events into a final Response + ResponseAccumulator responseAccumulator = ResponseAccumulator.create(); + + // Stream response - text is printed as it arrives + IterableStream events = + responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder() + .input("Tell me a short story about a brave explorer.")); + + for (ResponseStreamEvent event : events) { + responseAccumulator.accumulate(event); + event.outputTextDelta() + .ifPresent(textEvent -> System.out.print(textEvent.delta())); + } + System.out.println(); // newline after streamed text + + // Access the complete accumulated response + Response response = responseAccumulator.response(); + System.out.println("\nResponse ID: " + response.id()); + // END: com.azure.ai.agents.streaming.simple_sync + } finally { + if (agent != null) { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + System.out.println("Agent deleted"); + } + } + } +} diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingAsyncTests.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingAsyncTests.java new file mode 100644 index 000000000000..1b70df8fdb09 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingAsyncTests.java @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents; + +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.CodeInterpreterTool; +import com.azure.ai.agents.models.FunctionTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.http.HttpClient; +import com.azure.core.util.BinaryData; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseOutputItem; +import com.openai.models.responses.ResponseStatus; +import com.openai.models.responses.ResponseStreamEvent; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.ai.agents.TestUtils.DISPLAY_NAME_WITH_ARGUMENTS; +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.assertTrue; + +public class StreamingAsyncTests extends ClientTestBase { + + private static final String AGENT_MODEL = "gpt-4o"; + + // ======================================================================== + // Simple prompt streaming + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void simpleStreamingProducesTextDeltas(HttpClient httpClient, AgentsServiceVersion serviceVersion) { + AgentsAsyncClient agentsClient = getAgentsAsyncClient(httpClient, serviceVersion); + ResponsesAsyncClient responsesClient = getResponsesAsyncClient(httpClient, serviceVersion); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant. Reply in one sentence."); + + AtomicReference agentRef = new AtomicReference<>(); + + StepVerifier + .create(agentsClient.createAgentVersion("streaming-async-test-agent", agentDefinition).flatMap(agent -> { + agentRef.set(agent); + + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List textDeltas = new ArrayList<>(); + + Flux events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("Say hello.")); + + return events.doOnNext(event -> { + accumulator.accumulate(event); + event.outputTextDelta().ifPresent(textEvent -> textDeltas.add(textEvent.delta())); + }).then(Mono.fromCallable(() -> { + assertFalse(textDeltas.isEmpty(), "Should have received at least one text delta"); + + Response response = accumulator.response(); + assertNotNull(response.id()); + assertTrue(response.status().isPresent()); + assertEquals(ResponseStatus.COMPLETED, response.status().get()); + + String streamedText = String.join("", textDeltas); + assertFalse(streamedText.isEmpty()); + return response; + })); + }).flatMap(response -> { + AgentVersionDetails agent = agentRef.get(); + return agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()).thenReturn(response); + })) + .assertNext(response -> assertNotNull(response.id())) + .verifyComplete(); + } + + // ======================================================================== + // Function calling streaming + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void functionCallStreamingProducesFunctionCallEvents(HttpClient httpClient, + AgentsServiceVersion serviceVersion) { + AgentsAsyncClient agentsClient = getAgentsAsyncClient(httpClient, serviceVersion); + ResponsesAsyncClient responsesClient = getResponsesAsyncClient(httpClient, serviceVersion); + + FunctionTool tool = createWeatherFunctionTool(); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant. When asked about weather, use the get_weather function.") + .setTools(Collections.singletonList(tool)); + + AtomicReference agentRef = new AtomicReference<>(); + + StepVerifier.create( + agentsClient.createAgentVersion("function-streaming-async-test-agent", agentDefinition).flatMap(agent -> { + agentRef.set(agent); + + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List functionArgDeltas = new ArrayList<>(); + + Flux events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("What's the weather like in Seattle?")); + + return events.doOnNext(event -> { + accumulator.accumulate(event); + event.functionCallArgumentsDelta().ifPresent(argEvent -> functionArgDeltas.add(argEvent.delta())); + }).then(Mono.fromCallable(() -> { + assertFalse(functionArgDeltas.isEmpty(), "Should have received function call argument deltas"); + + Response response = accumulator.response(); + assertNotNull(response.id()); + + boolean hasFunctionCall = false; + for (ResponseOutputItem outputItem : response.output()) { + if (outputItem.functionCall().isPresent()) { + hasFunctionCall = true; + assertEquals("get_weather", outputItem.functionCall().get().name()); + assertNotNull(outputItem.functionCall().get().arguments()); + break; + } + } + assertTrue(hasFunctionCall, "Response should contain a function call output item"); + return response; + })); + }).flatMap(response -> { + AgentVersionDetails agent = agentRef.get(); + return agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()).thenReturn(response); + })).assertNext(response -> assertNotNull(response.id())).verifyComplete(); + } + + // ======================================================================== + // Code Interpreter streaming (Azure-specific tool) + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void codeInterpreterStreamingProducesCodeEvents(HttpClient httpClient, AgentsServiceVersion serviceVersion) { + AgentsAsyncClient agentsClient = getAgentsAsyncClient(httpClient, serviceVersion); + ResponsesAsyncClient responsesClient = getResponsesAsyncClient(httpClient, serviceVersion); + + CodeInterpreterTool tool = new CodeInterpreterTool(); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant that uses code interpreter for calculations.") + .setTools(Collections.singletonList(tool)); + + AtomicReference agentRef = new AtomicReference<>(); + + StepVerifier + .create(agentsClient.createAgentVersion("code-interpreter-streaming-async-test-agent", agentDefinition) + .flatMap(agent -> { + agentRef.set(agent); + + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List codeDeltas = new ArrayList<>(); + boolean[] codeInterpreterCompleted = { false }; + + Flux events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("What is 42 * 37? Use code to calculate.")); + + return events.doOnNext(event -> { + accumulator.accumulate(event); + event.codeInterpreterCallCodeDelta().ifPresent(e -> codeDeltas.add(e.delta())); + event.codeInterpreterCallCompleted().ifPresent(e -> codeInterpreterCompleted[0] = true); + }).then(Mono.fromCallable(() -> { + Response response = accumulator.response(); + assertNotNull(response.id()); + assertTrue(response.status().isPresent()); + assertEquals(ResponseStatus.COMPLETED, response.status().get()); + + assertFalse(codeDeltas.isEmpty(), "Should have received code interpreter code deltas"); + assertTrue(codeInterpreterCompleted[0], "Code interpreter should have completed"); + return response; + })); + }) + .flatMap(response -> { + AgentVersionDetails agent = agentRef.get(); + return agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()).thenReturn(response); + })) + .assertNext(response -> assertNotNull(response.id())) + .verifyComplete(); + } + + // ======================================================================== + // Helpers + // ======================================================================== + + private static FunctionTool createWeatherFunctionTool() { + Map locationProp = new LinkedHashMap<>(); + locationProp.put("type", "string"); + locationProp.put("description", "The city and state, e.g. Seattle, WA"); + + Map unitProp = new LinkedHashMap<>(); + unitProp.put("type", "string"); + unitProp.put("enum", Arrays.asList("celsius", "fahrenheit")); + + Map properties = new LinkedHashMap<>(); + properties.put("location", locationProp); + properties.put("unit", unitProp); + + Map parameters = new HashMap<>(); + parameters.put("type", BinaryData.fromObject("object")); + parameters.put("properties", BinaryData.fromObject(properties)); + parameters.put("required", BinaryData.fromObject(Arrays.asList("location", "unit"))); + parameters.put("additionalProperties", BinaryData.fromObject(false)); + + return new FunctionTool("get_weather", parameters, true) + .setDescription("Get the current weather in a given location"); + } +} diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingTests.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingTests.java new file mode 100644 index 000000000000..d1aab10da950 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/StreamingTests.java @@ -0,0 +1,213 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents; + +import com.azure.ai.agents.models.AgentReference; +import com.azure.ai.agents.models.AgentVersionDetails; +import com.azure.ai.agents.models.CodeInterpreterTool; +import com.azure.ai.agents.models.FunctionTool; +import com.azure.ai.agents.models.PromptAgentDefinition; +import com.azure.core.http.HttpClient; +import com.azure.core.util.BinaryData; +import com.azure.core.util.IterableStream; +import com.openai.helpers.ResponseAccumulator; +import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseOutputItem; +import com.openai.models.responses.ResponseStatus; +import com.openai.models.responses.ResponseStreamEvent; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static com.azure.ai.agents.TestUtils.DISPLAY_NAME_WITH_ARGUMENTS; +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.assertTrue; + +public class StreamingTests extends ClientTestBase { + + private static final String AGENT_MODEL = "gpt-4o"; + + // ======================================================================== + // Simple prompt streaming + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void simpleStreamingProducesTextDeltas(HttpClient httpClient, AgentsServiceVersion serviceVersion) { + AgentsClient agentsClient = getAgentsSyncClient(httpClient, serviceVersion); + ResponsesClient responsesClient = getResponsesSyncClient(httpClient, serviceVersion); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant. Reply in one sentence."); + + AgentVersionDetails agent = agentsClient.createAgentVersion("streaming-test-agent", agentDefinition); + + try { + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List textDeltas = new ArrayList<>(); + + IterableStream events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("Say hello.")); + + for (ResponseStreamEvent event : events) { + accumulator.accumulate(event); + event.outputTextDelta().ifPresent(textEvent -> textDeltas.add(textEvent.delta())); + } + + assertFalse(textDeltas.isEmpty(), "Should have received at least one text delta"); + + Response response = accumulator.response(); + assertNotNull(response.id()); + assertTrue(response.status().isPresent()); + assertEquals(ResponseStatus.COMPLETED, response.status().get()); + + // Concatenated deltas should match the final output text + String streamedText = String.join("", textDeltas); + assertFalse(streamedText.isEmpty()); + } finally { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + } + } + + // ======================================================================== + // Function calling streaming + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void functionCallStreamingProducesFunctionCallEvents(HttpClient httpClient, + AgentsServiceVersion serviceVersion) { + AgentsClient agentsClient = getAgentsSyncClient(httpClient, serviceVersion); + ResponsesClient responsesClient = getResponsesSyncClient(httpClient, serviceVersion); + + FunctionTool tool = createWeatherFunctionTool(); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant. When asked about weather, use the get_weather function.") + .setTools(Collections.singletonList(tool)); + + AgentVersionDetails agent = agentsClient.createAgentVersion("function-streaming-test-agent", agentDefinition); + + try { + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List functionArgDeltas = new ArrayList<>(); + + IterableStream events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("What's the weather like in Seattle?")); + + for (ResponseStreamEvent event : events) { + accumulator.accumulate(event); + event.functionCallArgumentsDelta().ifPresent(argEvent -> functionArgDeltas.add(argEvent.delta())); + } + + assertFalse(functionArgDeltas.isEmpty(), "Should have received function call argument deltas"); + + Response response = accumulator.response(); + assertNotNull(response.id()); + + // Verify that at least one output item is a function call + boolean hasFunctionCall = false; + for (ResponseOutputItem outputItem : response.output()) { + if (outputItem.functionCall().isPresent()) { + hasFunctionCall = true; + assertEquals("get_weather", outputItem.functionCall().get().name()); + assertNotNull(outputItem.functionCall().get().arguments()); + break; + } + } + assertTrue(hasFunctionCall, "Response should contain a function call output item"); + } finally { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + } + } + + // ======================================================================== + // Code Interpreter streaming (Azure-specific tool) + // ======================================================================== + + @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS) + @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters") + public void codeInterpreterStreamingProducesCodeEvents(HttpClient httpClient, AgentsServiceVersion serviceVersion) { + AgentsClient agentsClient = getAgentsSyncClient(httpClient, serviceVersion); + ResponsesClient responsesClient = getResponsesSyncClient(httpClient, serviceVersion); + + CodeInterpreterTool tool = new CodeInterpreterTool(); + + PromptAgentDefinition agentDefinition = new PromptAgentDefinition(AGENT_MODEL) + .setInstructions("You are a helpful assistant that uses code interpreter for calculations.") + .setTools(Collections.singletonList(tool)); + + AgentVersionDetails agent + = agentsClient.createAgentVersion("code-interpreter-streaming-test-agent", agentDefinition); + + try { + AgentReference agentReference = new AgentReference(agent.getName()).setVersion(agent.getVersion()); + + ResponseAccumulator accumulator = ResponseAccumulator.create(); + List codeDeltas = new ArrayList<>(); + boolean[] codeInterpreterCompleted = { false }; + + IterableStream events = responsesClient.createStreamingWithAgent(agentReference, + ResponseCreateParams.builder().input("What is 42 * 37? Use code to calculate.")); + + for (ResponseStreamEvent event : events) { + accumulator.accumulate(event); + event.codeInterpreterCallCodeDelta().ifPresent(e -> codeDeltas.add(e.delta())); + event.codeInterpreterCallCompleted().ifPresent(e -> codeInterpreterCompleted[0] = true); + } + + Response response = accumulator.response(); + assertNotNull(response.id()); + assertTrue(response.status().isPresent()); + assertEquals(ResponseStatus.COMPLETED, response.status().get()); + + // Code interpreter should have run and produced code + assertFalse(codeDeltas.isEmpty(), "Should have received code interpreter code deltas"); + assertTrue(codeInterpreterCompleted[0], "Code interpreter should have completed"); + } finally { + agentsClient.deleteAgentVersion(agent.getName(), agent.getVersion()); + } + } + + // ======================================================================== + // Helpers + // ======================================================================== + + private static FunctionTool createWeatherFunctionTool() { + Map locationProp = new LinkedHashMap<>(); + locationProp.put("type", "string"); + locationProp.put("description", "The city and state, e.g. Seattle, WA"); + + Map unitProp = new LinkedHashMap<>(); + unitProp.put("type", "string"); + unitProp.put("enum", Arrays.asList("celsius", "fahrenheit")); + + Map properties = new LinkedHashMap<>(); + properties.put("location", locationProp); + properties.put("unit", unitProp); + + Map parameters = new HashMap<>(); + parameters.put("type", BinaryData.fromObject("object")); + parameters.put("properties", BinaryData.fromObject(properties)); + parameters.put("required", BinaryData.fromObject(Arrays.asList("location", "unit"))); + parameters.put("additionalProperties", BinaryData.fromObject(false)); + + return new FunctionTool("get_weather", parameters, true) + .setDescription("Get the current weather in a given location"); + } +} diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/StreamingUtilsTest.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/StreamingUtilsTest.java new file mode 100644 index 000000000000..8af0eb8c4d6b --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/StreamingUtilsTest.java @@ -0,0 +1,347 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.implementation; + +import com.azure.core.util.IterableStream; +import com.openai.core.http.AsyncStreamResponse; +import com.openai.core.http.StreamResponse; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.test.StepVerifier; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link StreamingUtils}. + */ +public class StreamingUtilsTest { + + // ======================================================================== + // toIterableStream tests + // ======================================================================== + + @Test + public void toIterableStreamSingleItem() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("hello"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + List items = collect(result); + assertEquals(1, items.size()); + assertEquals("hello", items.get(0)); + assertTrue(closed.get(), "StreamResponse should be closed after full consumption"); + } + + @Test + public void toIterableStreamMultipleItems() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("event1", "event2", "event3"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + List items = collect(result); + assertEquals(3, items.size()); + assertEquals("event1", items.get(0)); + assertEquals("event2", items.get(1)); + assertEquals("event3", items.get(2)); + assertTrue(closed.get()); + } + + @Test + public void toIterableStreamEmptyClosesResource() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.empty(), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + List items = collect(result); + assertTrue(items.isEmpty()); + assertTrue(closed.get(), "StreamResponse should be closed even when stream is empty"); + } + + @Test + public void toIterableStreamPreservesOrder() { + AtomicBoolean closed = new AtomicBoolean(false); + List expected = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + expected.add("item-" + i); + } + StreamResponse streamResponse = testStreamResponse(expected.stream(), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + List items = collect(result); + assertEquals(expected, items); + assertTrue(closed.get()); + } + + @Test + public void toIterableStreamNotClosedBeforeFullConsumption() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("a", "b", "c"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + Iterator iterator = result.iterator(); + + // Consume only one item + assertTrue(iterator.hasNext()); + assertEquals("a", iterator.next()); + assertFalse(closed.get(), "StreamResponse should not be closed before full consumption"); + + // Consume the rest + assertEquals("b", iterator.next()); + assertEquals("c", iterator.next()); + assertFalse(iterator.hasNext()); + assertTrue(closed.get()); + } + + @Test + public void toIterableStreamIteratorThrowsWhenExhausted() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("only"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + Iterator iterator = result.iterator(); + + assertEquals("only", iterator.next()); + assertFalse(iterator.hasNext()); + assertThrows(NoSuchElementException.class, iterator::next); + } + + @Test + public void toIterableStreamThrowsOnSecondIteratorCall() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("a", "b"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + // First iterator() call should succeed + Iterator first = result.iterator(); + assertEquals("a", first.next()); + + // Second iterator() call should throw + assertThrows(IllegalStateException.class, result::iterator); + } + + @Test + public void toIterableStreamSecondIteratorCallAfterFullConsumptionThrows() { + AtomicBoolean closed = new AtomicBoolean(false); + StreamResponse streamResponse = testStreamResponse(Stream.of("a", "b"), closed); + + IterableStream result = StreamingUtils.toIterableStream(streamResponse); + + // Fully consume the first iterator + List items = collect(result); + assertEquals(Arrays.asList("a", "b"), items); + assertTrue(closed.get()); + + // Without the single-use guard, this would silently return an empty iterator + // instead of failing. The guard converts that silent bug into a clear error. + assertThrows(IllegalStateException.class, result::iterator); + } + + // ======================================================================== + // toFlux tests + // ======================================================================== + + @Test + public void toFluxSingleItem() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux).then(() -> { + asyncStream.emit("hello"); + asyncStream.complete(); + }).expectNext("hello").verifyComplete(); + } + + @Test + public void toFluxMultipleItems() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux).then(() -> { + asyncStream.emit("event1"); + asyncStream.emit("event2"); + asyncStream.emit("event3"); + asyncStream.complete(); + }).expectNext("event1").expectNext("event2").expectNext("event3").verifyComplete(); + } + + @Test + public void toFluxEmptyStream() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux).then(asyncStream::complete).verifyComplete(); + } + + @Test + public void toFluxPropagatesError() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + RuntimeException expectedError = new RuntimeException("stream failed"); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux).then(() -> { + asyncStream.emit("before-error"); + asyncStream.completeWithError(expectedError); + }) + .expectNext("before-error") + .expectErrorMatches(e -> e instanceof RuntimeException && "stream failed".equals(e.getMessage())) + .verify(); + } + + @Test + public void toFluxErrorWithNoEmissions() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + RuntimeException expectedError = new RuntimeException("immediate failure"); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux) + .then(() -> asyncStream.completeWithError(expectedError)) + .expectErrorMatches(e -> e instanceof RuntimeException && "immediate failure".equals(e.getMessage())) + .verify(); + } + + @Test + public void toFluxPreservesOrder() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + List expected = new ArrayList<>(); + for (int i = 0; i < 50; i++) { + expected.add("item-" + i); + } + + Flux flux = StreamingUtils.toFlux(asyncStream); + + StepVerifier.create(flux).then(() -> { + for (String item : expected) { + asyncStream.emit(item); + } + asyncStream.complete(); + }) + .recordWith(ArrayList::new) + .thenConsumeWhile(x -> true) + .consumeRecordedWith(recorded -> assertEquals(expected, new ArrayList<>(recorded))) + .verifyComplete(); + } + + @Test + public void toFluxClosesOnDispose() { + TestAsyncStreamResponse asyncStream = new TestAsyncStreamResponse<>(); + + Flux flux = StreamingUtils.toFlux(asyncStream); + + // Subscribe and immediately dispose + flux.subscribe().dispose(); + assertTrue(asyncStream.isClosed(), "AsyncStreamResponse should be closed on disposal"); + } + + // ======================================================================== + // Test helpers + // ======================================================================== + + /** + * Creates a simple StreamResponse backed by the given stream, tracking close state. + */ + private static StreamResponse testStreamResponse(Stream stream, AtomicBoolean closed) { + return new StreamResponse() { + @Override + public Stream stream() { + return stream; + } + + @Override + public void close() { + closed.set(true); + } + }; + } + + /** + * Collects all items from an IterableStream into a list. + */ + private static List collect(IterableStream iterableStream) { + List items = new ArrayList<>(); + for (T item : iterableStream) { + items.add(item); + } + return items; + } + + /** + * A test implementation of AsyncStreamResponse that allows manual emission of items + * and completion/error signals. Executes the handler synchronously for test simplicity. + */ + private static class TestAsyncStreamResponse implements AsyncStreamResponse { + private Handler handler; + private final AtomicBoolean closed = new AtomicBoolean(false); + private final CompletableFuture completeFuture = new CompletableFuture<>(); + + void emit(T item) { + if (handler != null) { + handler.onNext(item); + } + } + + void complete() { + if (handler != null) { + handler.onComplete(Optional.empty()); + } + completeFuture.complete(null); + } + + void completeWithError(Throwable error) { + if (handler != null) { + handler.onComplete(Optional.of(error)); + } + completeFuture.completeExceptionally(error); + } + + boolean isClosed() { + return closed.get(); + } + + @Override + public AsyncStreamResponse subscribe(Handler handler) { + this.handler = handler; + return this; + } + + @Override + public AsyncStreamResponse subscribe(Handler handler, Executor executor) { + this.handler = handler; + return this; + } + + @Override + public CompletableFuture onCompleteFuture() { + return completeFuture; + } + + @Override + public void close() { + closed.set(true); + } + } +}