diff --git a/sdk/search/azure-search-documents/customizations/src/main/java/SearchCustomizations.java b/sdk/search/azure-search-documents/customizations/src/main/java/SearchCustomizations.java index 8296bc71630dc..e3b05500e0704 100644 --- a/sdk/search/azure-search-documents/customizations/src/main/java/SearchCustomizations.java +++ b/sdk/search/azure-search-documents/customizations/src/main/java/SearchCustomizations.java @@ -293,32 +293,65 @@ private static String streamEventDocumentation() { return "Events are polymorphic: see {@link KnowledgeBaseRetrievalStreamEvent} for the known subtypes. " + "Use a known subtype's {@code getValue()} to access its typed payload, for example " + "{@link com.azure.search.documents.knowledgebases.models.KnowledgeBaseAnswerCompletedStreamEvent" - + "#getValue()}. " - + "Unrecognized events expose only their names through " + + "#getValue()}. " + "Unrecognized events expose only their names through " + "{@link KnowledgeBaseRetrievalStreamEvent#getEventName()}; their payloads are not available.\n\n"; } private static String baseEventSource() { return header("com.azure.search.documents.knowledgebases.models") - + "import com.azure.core.annotation.Generated;\n\n" + "/**\n" - + " * Abstract base for polymorphic events emitted by a streaming knowledge base retrieval.\n" + + "import com.azure.core.annotation.Generated;\n" + + "import com.azure.search.documents.knowledgebases.implementation." + + "KnowledgeBaseRetrievalStreamEventConverter;\n" + "import java.util.List;\n\n" + "/**\n" + + " * Base for polymorphic events emitted by a streaming knowledge base retrieval.\n" + " * Known events expose typed payloads through their subtype's {@code getValue()} method.\n" - + " * Unrecognized events expose only their names through {@link #getEventName()}; " - + "their payloads are not available.\n" + + " * Unrecognized events expose only their names through\n" + + " * {@link #getEventName()}; their payloads are not available and they are not terminal.\n" + " * Stream events are envelopes, not JSON payload models; serialize the typed payload instead.\n" + " *\n" + " * @see KnowledgeBaseRetrievalStartedStreamEvent\n" + " * @see KnowledgeBaseActivityStartedStreamEvent\n" + " * @see KnowledgeBaseActivityCompletedStreamEvent\n" + " * @see KnowledgeBaseAnswerCompletedStreamEvent\n" + " * @see KnowledgeBaseReferencesCompletedStreamEvent\n" + " * @see KnowledgeBaseErrorStreamEvent\n" + " * @see KnowledgeBaseResponseCompletedStreamEvent\n" + " */\n" - + "public abstract class KnowledgeBaseRetrievalStreamEvent {\n" + " @Generated\n" + + "public abstract class KnowledgeBaseRetrievalStreamEvent {\n" + " static {\n" + + " KnowledgeBaseRetrievalStreamEventConverter.setConstructorAccessor(\n" + + " new KnowledgeBaseRetrievalStreamEventConverter.ConstructorAccessor() {\n" + + " @Override\n" + + " public KnowledgeBaseRetrievalStreamEvent createEvent(String eventName) {\n" + + " return new KnowledgeBaseRetrievalStreamEvent(eventName) { };\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseRetrievalStartedStreamEvent createRetrievalStarted(\n" + + " KnowledgeBaseRetrievalStartedEvent value) {\n" + + " return new KnowledgeBaseRetrievalStartedStreamEvent(value);\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseActivityStartedStreamEvent createActivityStarted(\n" + + " KnowledgeBaseActivityStartedEvent value) {\n" + + " return new KnowledgeBaseActivityStartedStreamEvent(value);\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseActivityCompletedStreamEvent createActivityCompleted(\n" + + " KnowledgeBaseActivityRecord value) {\n" + + " return new KnowledgeBaseActivityCompletedStreamEvent(value);\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseAnswerCompletedStreamEvent createAnswerCompleted(\n" + + " KnowledgeBaseAnswerCompletedEvent value) {\n" + + " return new KnowledgeBaseAnswerCompletedStreamEvent(value);\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseReferencesCompletedStreamEvent createReferencesCompleted(\n" + + " List value) {\n" + + " return new KnowledgeBaseReferencesCompletedStreamEvent(value);\n" + + " }\n\n" + " @Override\n" + + " public KnowledgeBaseErrorStreamEvent createError(KnowledgeBaseStreamErrorEvent value) {\n" + + " return new KnowledgeBaseErrorStreamEvent(value);\n" + " }\n\n" + + " @Override\n" + + " public KnowledgeBaseResponseCompletedStreamEvent createResponseCompleted(\n" + + " KnowledgeBaseResponseCompletedEvent value) {\n" + + " return new KnowledgeBaseResponseCompletedStreamEvent(value);\n" + + " }\n" + " });\n" + " }\n\n" + " @Generated\n" + " private final String eventName;\n\n" + " /**\n" + " * Creates a stream event.\n" + " *\n" + " * @param eventName The server-sent event name.\n" + " */\n" + " @Generated\n" - + " protected KnowledgeBaseRetrievalStreamEvent(String eventName) {\n" - + " this.eventName = eventName;\n" + " }\n\n" + " /**\n" - + " * Gets the server-sent event name.\n" + " *\n" + " * @return The event name.\n" - + " */\n" + " @Generated\n" + " public final String getEventName() {\n" - + " return eventName;\n" + " }\n\n" + " /**\n" + + " KnowledgeBaseRetrievalStreamEvent(String eventName) {\n" + " this.eventName = eventName;\n" + + " }\n\n" + " /**\n" + " * Gets the server-sent event name.\n" + " *\n" + + " * @return The event name.\n" + " */\n" + " @Generated\n" + + " public final String getEventName() {\n" + " return eventName;\n" + " }\n\n" + " /**\n" + " * Gets whether this event terminates the retrieval stream.\n" + " *\n" + " * @return {@code true} if this is a terminal event; otherwise {@code false}.\n" + " */\n" + " @Generated\n" + " public boolean isTerminal() {\n" + " return false;\n" + " }\n" @@ -341,11 +374,11 @@ private static String wrapperSource(String className, String payloadType, String + " *\n" + " * @see KnowledgeBaseRetrievalStreamEvent\n" + " */\n" + "@Immutable\n" + "public final class " + className + " extends KnowledgeBaseRetrievalStreamEvent {\n" + " @Generated\n" + " private final " + valueType + " value;\n\n" + " /**\n" + " * Creates an event wrapper.\n" + " *\n" - + " * @param value The event payload.\n" + " */\n" + " @Generated\n" + " public " + className - + "(" + valueType + " value) {\n" + " super(\"" + eventName + "\");\n" - + " this.value = value;\n" + " }\n\n" + " /**\n" + " * Gets the event payload.\n" - + " *\n" + " * @return The event payload.\n" + " */\n" + " @Generated\n" + " public " - + valueType + " getValue() {\n" + " return value;\n" + " }\n" + terminalOverride + "}\n"; + + " * @param value The event payload.\n" + " */\n" + " @Generated\n" + " " + className + "(" + + valueType + " value) {\n" + " super(\"" + eventName + "\");\n" + " this.value = value;\n" + + " }\n\n" + " /**\n" + " * Gets the event payload.\n" + " *\n" + + " * @return The event payload.\n" + " */\n" + " @Generated\n" + " public " + valueType + + " getValue() {\n" + " return value;\n" + " }\n" + terminalOverride + "}\n"; } private static String header(String packageName) { diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/indexes/models/KnowledgeSourceFileCapacity.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/indexes/models/KnowledgeSourceFileCapacity.java index f96f8b9bcee3f..ad77e34d6d469 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/indexes/models/KnowledgeSourceFileCapacity.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/indexes/models/KnowledgeSourceFileCapacity.java @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. // Code generated by Microsoft (R) TypeSpec Code Generator. - package com.azure.search.documents.indexes.models; import com.azure.core.annotation.Generated; @@ -17,6 +16,7 @@ */ @Immutable public final class KnowledgeSourceFileCapacity implements JsonSerializable { + /* * The maximum number of files allowed in the knowledge source. */ @@ -44,7 +44,7 @@ private KnowledgeSourceFileCapacity() { /** * Get the maxFileCount property: The maximum number of files allowed in the knowledge source. - * + * * @return the maxFileCount value. */ @Generated @@ -55,7 +55,7 @@ public int getMaxFileCount() { /** * Get the remainingFileCount property: The number of additional files that can be uploaded at the time of the * request. - * + * * @return the remainingFileCount value. */ @Generated @@ -65,7 +65,7 @@ public int getRemainingFileCount() { /** * Get the maxFileSizeBytes property: The maximum size in bytes of an individual uploaded file. - * + * * @return the maxFileSizeBytes value. */ @Generated @@ -85,7 +85,7 @@ public JsonWriter toJson(JsonWriter jsonWriter) throws IOException { /** * Reads an instance of KnowledgeSourceFileCapacity from the JsonReader. - * + * * @param jsonReader The JsonReader being read. * @return An instance of KnowledgeSourceFileCapacity if the JsonReader was pointing to an instance of it, or null * if it was pointing to JSON null. @@ -99,7 +99,6 @@ public static KnowledgeSourceFileCapacity fromJson(JsonReader jsonReader) throws while (reader.nextToken() != JsonToken.END_OBJECT) { String fieldName = reader.getFieldName(); reader.nextToken(); - if ("maxFileCount".equals(fieldName)) { deserializedKnowledgeSourceFileCapacity.maxFileCount = reader.getInt(); } else if ("remainingFileCount".equals(fieldName)) { @@ -110,7 +109,6 @@ public static KnowledgeSourceFileCapacity fromJson(JsonReader jsonReader) throws reader.skipChildren(); } } - return deserializedKnowledgeSourceFileCapacity; }); } diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/implementation/KnowledgeBaseRetrievalStreamEventConverter.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/implementation/KnowledgeBaseRetrievalStreamEventConverter.java index 83b1898930dbf..f657be60a4448 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/implementation/KnowledgeBaseRetrievalStreamEventConverter.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/implementation/KnowledgeBaseRetrievalStreamEventConverter.java @@ -3,6 +3,7 @@ package com.azure.search.documents.knowledgebases.implementation; +import com.azure.core.util.logging.ClientLogger; import com.azure.json.JsonProviders; import com.azure.json.JsonReader; import com.azure.json.ReadValueCallback; @@ -23,14 +24,116 @@ import com.azure.search.documents.knowledgebases.models.KnowledgeBaseStreamErrorEvent; import java.io.IOException; import java.io.UncheckedIOException; +import java.util.List; +import java.util.Objects; /** * Converts knowledge base retrieval stream event payloads to typed event models. */ public final class KnowledgeBaseRetrievalStreamEventConverter { + private static final ClientLogger LOGGER = new ClientLogger(KnowledgeBaseRetrievalStreamEventConverter.class); + private static volatile ConstructorAccessor constructorAccessor; + private KnowledgeBaseRetrievalStreamEventConverter() { } + /** + * Provides access to the SDK-only event constructors. + */ + public interface ConstructorAccessor { + /** + * Creates a name-only event. + * + * @param eventName The event name. + * @return The event. + */ + KnowledgeBaseRetrievalStreamEvent createEvent(String eventName); + + /** + * Creates a retrieval-started event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseRetrievalStartedStreamEvent createRetrievalStarted(KnowledgeBaseRetrievalStartedEvent value); + + /** + * Creates an activity-started event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseActivityStartedStreamEvent createActivityStarted(KnowledgeBaseActivityStartedEvent value); + + /** + * Creates an activity-completed event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseActivityCompletedStreamEvent createActivityCompleted(KnowledgeBaseActivityRecord value); + + /** + * Creates an answer-completed event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseAnswerCompletedStreamEvent createAnswerCompleted(KnowledgeBaseAnswerCompletedEvent value); + + /** + * Creates a references-completed event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseReferencesCompletedStreamEvent createReferencesCompleted(List value); + + /** + * Creates an error event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseErrorStreamEvent createError(KnowledgeBaseStreamErrorEvent value); + + /** + * Creates a response-completed event. + * + * @param value The decoded payload. + * @return The event. + */ + KnowledgeBaseResponseCompletedStreamEvent createResponseCompleted(KnowledgeBaseResponseCompletedEvent value); + } + + /** + * Registers the event constructors from the models package. + * + * @param accessor The constructor accessor. + */ + public static void setConstructorAccessor(ConstructorAccessor accessor) { + constructorAccessor = Objects.requireNonNull(accessor, "'accessor' cannot be null."); + } + + private static ConstructorAccessor getConstructorAccessor() { + ConstructorAccessor accessor = constructorAccessor; + if (accessor == null) { + try { + Class.forName(KnowledgeBaseRetrievalStreamEvent.class.getName(), true, + KnowledgeBaseRetrievalStreamEvent.class.getClassLoader()); + } catch (ClassNotFoundException | LinkageError exception) { + throw LOGGER.logExceptionAsError(new IllegalStateException( + "Failed to initialize knowledge base retrieval stream event constructors.", exception)); + } + accessor = constructorAccessor; + if (accessor == null) { + throw LOGGER.logExceptionAsError(new IllegalStateException( + "Knowledge base retrieval stream event constructors were not registered.")); + } + } + return accessor; + } + /** * Converts a stream event payload. * @@ -39,37 +142,36 @@ private KnowledgeBaseRetrievalStreamEventConverter() { * @return The typed stream event. */ public static KnowledgeBaseRetrievalStreamEvent convert(String eventName, String data) { + ConstructorAccessor accessor = getConstructorAccessor(); switch (eventName) { case "retrieval.started": - return new KnowledgeBaseRetrievalStartedStreamEvent( - read(eventName, data, KnowledgeBaseRetrievalStartedEvent::fromJson)); + return accessor + .createRetrievalStarted(read(eventName, data, KnowledgeBaseRetrievalStartedEvent::fromJson)); case "activity.started": - return new KnowledgeBaseActivityStartedStreamEvent( - read(eventName, data, KnowledgeBaseActivityStartedEvent::fromJson)); + return accessor + .createActivityStarted(read(eventName, data, KnowledgeBaseActivityStartedEvent::fromJson)); case "activity.completed": - return new KnowledgeBaseActivityCompletedStreamEvent( - read(eventName, data, KnowledgeBaseActivityRecord::fromJson)); + return accessor.createActivityCompleted(read(eventName, data, KnowledgeBaseActivityRecord::fromJson)); case "answer.completed": - return new KnowledgeBaseAnswerCompletedStreamEvent( - read(eventName, data, KnowledgeBaseAnswerCompletedEvent::fromJson)); + return accessor + .createAnswerCompleted(read(eventName, data, KnowledgeBaseAnswerCompletedEvent::fromJson)); case "references.completed": - return new KnowledgeBaseReferencesCompletedStreamEvent( + return accessor.createReferencesCompleted( read(eventName, data, reader -> reader.readArray(KnowledgeBaseReference::fromJson))); case "error": - return new KnowledgeBaseErrorStreamEvent( - read(eventName, data, KnowledgeBaseStreamErrorEvent::fromJson)); + return accessor.createError(read(eventName, data, KnowledgeBaseStreamErrorEvent::fromJson)); case "response.completed": - return new KnowledgeBaseResponseCompletedStreamEvent( - read(eventName, data, KnowledgeBaseResponseCompletedEvent::fromJson)); + return accessor + .createResponseCompleted(read(eventName, data, KnowledgeBaseResponseCompletedEvent::fromJson)); default: - return new UnrecognizedStreamEvent(eventName); + return accessor.createEvent(eventName); } } @@ -80,12 +182,6 @@ private static T read(String eventName, String data, ReadValueCallback value) { + KnowledgeBaseReferencesCompletedStreamEvent(List value) { super("references.completed"); this.value = value; } diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseResponseCompletedStreamEvent.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseResponseCompletedStreamEvent.java index 7e8f0cc42f2c6..438c3bc039694 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseResponseCompletedStreamEvent.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseResponseCompletedStreamEvent.java @@ -23,7 +23,7 @@ public final class KnowledgeBaseResponseCompletedStreamEvent extends KnowledgeBa * @param value The event payload. */ @Generated - public KnowledgeBaseResponseCompletedStreamEvent(KnowledgeBaseResponseCompletedEvent value) { + KnowledgeBaseResponseCompletedStreamEvent(KnowledgeBaseResponseCompletedEvent value) { super("response.completed"); this.value = value; } diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStartedStreamEvent.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStartedStreamEvent.java index 9c2be7f80e0cc..a47db8de65ebf 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStartedStreamEvent.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStartedStreamEvent.java @@ -23,7 +23,7 @@ public final class KnowledgeBaseRetrievalStartedStreamEvent extends KnowledgeBas * @param value The event payload. */ @Generated - public KnowledgeBaseRetrievalStartedStreamEvent(KnowledgeBaseRetrievalStartedEvent value) { + KnowledgeBaseRetrievalStartedStreamEvent(KnowledgeBaseRetrievalStartedEvent value) { super("retrieval.started"); this.value = value; } diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStreamEvent.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStreamEvent.java index 12fc8ba8dfd01..41a8838398b71 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStreamEvent.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/knowledgebases/models/KnowledgeBaseRetrievalStreamEvent.java @@ -3,11 +3,14 @@ package com.azure.search.documents.knowledgebases.models; import com.azure.core.annotation.Generated; +import com.azure.search.documents.knowledgebases.implementation.KnowledgeBaseRetrievalStreamEventConverter; +import java.util.List; /** - * Abstract base for polymorphic events emitted by a streaming knowledge base retrieval. + * Base for polymorphic events emitted by a streaming knowledge base retrieval. * Known events expose typed payloads through their subtype's {@code getValue()} method. - * Unrecognized events expose only their names through {@link #getEventName()}; their payloads are not available. + * Unrecognized events expose only their names through + * {@link #getEventName()}; their payloads are not available and they are not terminal. * Stream events are envelopes, not JSON payload models; serialize the typed payload instead. * * @see KnowledgeBaseRetrievalStartedStreamEvent @@ -29,7 +32,7 @@ public abstract class KnowledgeBaseRetrievalStreamEvent { * @param eventName The server-sent event name. */ @Generated - protected KnowledgeBaseRetrievalStreamEvent(String eventName) { + KnowledgeBaseRetrievalStreamEvent(String eventName) { this.eventName = eventName; } @@ -52,4 +55,57 @@ public final String getEventName() { public boolean isTerminal() { return false; } + + static { + KnowledgeBaseRetrievalStreamEventConverter + .setConstructorAccessor(new KnowledgeBaseRetrievalStreamEventConverter.ConstructorAccessor() { + + @Override + public KnowledgeBaseRetrievalStreamEvent createEvent(String eventName) { + return new KnowledgeBaseRetrievalStreamEvent(eventName) { + }; + } + + @Override + public KnowledgeBaseRetrievalStartedStreamEvent + createRetrievalStarted(KnowledgeBaseRetrievalStartedEvent value) { + return new KnowledgeBaseRetrievalStartedStreamEvent(value); + } + + @Override + public KnowledgeBaseActivityStartedStreamEvent + createActivityStarted(KnowledgeBaseActivityStartedEvent value) { + return new KnowledgeBaseActivityStartedStreamEvent(value); + } + + @Override + public KnowledgeBaseActivityCompletedStreamEvent + createActivityCompleted(KnowledgeBaseActivityRecord value) { + return new KnowledgeBaseActivityCompletedStreamEvent(value); + } + + @Override + public KnowledgeBaseAnswerCompletedStreamEvent + createAnswerCompleted(KnowledgeBaseAnswerCompletedEvent value) { + return new KnowledgeBaseAnswerCompletedStreamEvent(value); + } + + @Override + public KnowledgeBaseReferencesCompletedStreamEvent + createReferencesCompleted(List value) { + return new KnowledgeBaseReferencesCompletedStreamEvent(value); + } + + @Override + public KnowledgeBaseErrorStreamEvent createError(KnowledgeBaseStreamErrorEvent value) { + return new KnowledgeBaseErrorStreamEvent(value); + } + + @Override + public KnowledgeBaseResponseCompletedStreamEvent + createResponseCompleted(KnowledgeBaseResponseCompletedEvent value) { + return new KnowledgeBaseResponseCompletedStreamEvent(value); + } + }); + } } diff --git a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/models/implementation/sse/ServerSentEventStream.java b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/models/implementation/sse/ServerSentEventStream.java index 6aff6eb7a3a09..b96df5a7953e6 100644 --- a/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/models/implementation/sse/ServerSentEventStream.java +++ b/sdk/search/azure-search-documents/src/main/java/com/azure/search/documents/models/implementation/sse/ServerSentEventStream.java @@ -7,6 +7,7 @@ import com.azure.core.util.BinaryData; import com.azure.core.util.CloseableIterableStream; import com.azure.core.util.FluxUtil; +import com.azure.core.util.logging.ClientLogger; import java.io.Closeable; import java.io.IOException; import java.io.InputStream; @@ -30,6 +31,7 @@ * Parses one server-sent event response. */ final class ServerSentEventStream { + private static final ClientLogger LOGGER = new ClientLogger(ServerSentEventStream.class); private static final String DEFAULT_EVENT = "message"; private ServerSentEventStream() { @@ -161,7 +163,7 @@ public boolean hasNext() { @Override public ServerSentEventFrame next() { if (!hasNext()) { - throw new NoSuchElementException(); + throw LOGGER.logExceptionAsError(new NoSuchElementException()); } ServerSentEventFrame frame = next; next = null; diff --git a/sdk/search/azure-search-documents/src/test/java/com/azure/search/documents/knowledgebases/KnowledgeBaseRetrievalStreamTests.java b/sdk/search/azure-search-documents/src/test/java/com/azure/search/documents/knowledgebases/KnowledgeBaseRetrievalStreamTests.java index 75a0266ecfda8..44ce071e03dd0 100644 --- a/sdk/search/azure-search-documents/src/test/java/com/azure/search/documents/knowledgebases/KnowledgeBaseRetrievalStreamTests.java +++ b/sdk/search/azure-search-documents/src/test/java/com/azure/search/documents/knowledgebases/KnowledgeBaseRetrievalStreamTests.java @@ -10,7 +10,6 @@ import com.azure.core.test.http.MockHttpResponse; import com.azure.core.util.CloseableIterableStream; import com.azure.json.JsonSerializable; -import com.azure.json.JsonWriter; import com.azure.search.documents.knowledgebases.implementation.KnowledgeBaseRetrievalStreamEventConverter; import com.azure.search.documents.knowledgebases.models.KnowledgeBaseActivityCompletedStreamEvent; import com.azure.search.documents.knowledgebases.models.KnowledgeBaseActivityRecord; @@ -48,7 +47,6 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; @@ -59,7 +57,6 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -74,6 +71,11 @@ public class KnowledgeBaseRetrievalStreamTests { + "\"reasoningEffort\":{\"kind\":\"low\"}}"; private static final String RESPONSE_COMPLETED_JSON = "{\"statusCode\":200,\"response\":{}}"; + @Test + public void baseEventRemainsAbstract() { + assertTrue(Modifier.isAbstract(KnowledgeBaseRetrievalStreamEvent.class.getModifiers())); + } + @Test public void convertsAllStreamEventVariants() { KnowledgeBaseRetrievalStreamEvent retrievalStarted @@ -132,18 +134,10 @@ public void preservesUnknownEventsAndRejectsMalformedKnownEvents() { KnowledgeBaseRetrievalStreamEvent event = KnowledgeBaseRetrievalStreamEventConverter.convert("future.event", "not json\nsecond line"); - assertEvent(event, "future.event", false); - assertTrue(Modifier.isPrivate(event.getClass().getModifiers())); - assertTrue(Arrays.stream(event.getClass().getDeclaredFields()).allMatch(field -> field.isSynthetic())); - assertEquals(1, - Arrays.stream(event.getClass().getDeclaredConstructors()) - .filter(constructor -> !constructor.isSynthetic()) - .count()); - assertTrue(Arrays.stream(event.getClass().getDeclaredConstructors()) - .filter(constructor -> !constructor.isSynthetic()) - .allMatch(constructor -> constructor.getParameterCount() == 1 - && constructor.getParameterTypes()[0] == String.class)); - assertEvent(KnowledgeBaseRetrievalStreamEventConverter.convert("future.event", null), "future.event", false); + assertUnknownEvent(event, "future.event"); + KnowledgeBaseRetrievalStreamEvent nullPayload + = KnowledgeBaseRetrievalStreamEventConverter.convert("future.event", null); + assertUnknownEvent(nullPayload, "future.event"); assertThrows(RuntimeException.class, () -> KnowledgeBaseRetrievalStreamEventConverter.convert("response.completed", "{")); } @@ -170,6 +164,13 @@ public void bothClientsDecodeTypedPayloadsForEveryKnownVariant(String eventName, KnowledgeBaseRetrievalStreamEvent converted = KnowledgeBaseRetrievalStreamEventConverter.convert(eventName, data); assertEquals("error".equals(eventName) || "response.completed".equals(eventName), converted.isTerminal()); + if ("activity.completed".equals(eventName)) { + assertInstanceOf(KnowledgeBaseSearchIndexActivityRecord.class, payload(converted)); + } + if ("references.completed".equals(eventName)) { + assertInstanceOf(KnowledgeBaseSearchIndexReference.class, + ((KnowledgeBaseReferencesCompletedStreamEvent) converted).getValue().get(0)); + } try (CloseableIterableStream stream = createBuilder(body, null).buildClient().retrieveStream(new KnowledgeBaseRetrievalOptions())) { @@ -231,106 +232,25 @@ public void bothClientsPreserveAgenticReasoningInProgressAndCompletedResponse() @ParameterizedTest @ValueSource(strings = { "", " \t ", "not json", "not json\n second line ", " {\"future\":true}\n " }) public void unknownPayloadsAreDiscardedWithoutDisplacingLaterTerminalEvents(String data) throws IOException { - assertEvent(KnowledgeBaseRetrievalStreamEventConverter.convert("future.event", data), "future.event", false); + KnowledgeBaseRetrievalStreamEvent converted + = KnowledgeBaseRetrievalStreamEventConverter.convert("future.event", data); + assertUnknownEvent(converted, "future.event"); String body = frame("future.event", data) + frame("response.completed", RESPONSE_COMPLETED_JSON); try (CloseableIterableStream stream = createBuilder(body, null).buildClient().retrieveStream(new KnowledgeBaseRetrievalOptions())) { Iterator iterator = stream.iterator(); KnowledgeBaseRetrievalStreamEvent event = iterator.next(); - assertEvent(event, "future.event", false); + assertUnknownEvent(event, "future.event"); assertEvent(iterator.next(), "response.completed", true); assertFalse(iterator.hasNext()); } StepVerifier .create(createBuilder(body, null).buildAsyncClient().retrieveStream(new KnowledgeBaseRetrievalOptions())) - .assertNext(event -> assertEvent(event, "future.event", false)) + .assertNext(event -> assertUnknownEvent(event, "future.event")) .assertNext(event -> assertEvent(event, "response.completed", true)) .verifyComplete(); } - @ParameterizedTest - @MethodSource("knownEvents") - public void payloadOnlyConstructorsRetainTypedPayloadsWithoutRawOrJsonApis(String eventName, String json) - throws ReflectiveOperationException { - KnowledgeBaseRetrievalStreamEvent converted - = KnowledgeBaseRetrievalStreamEventConverter.convert(eventName, json); - KnowledgeBaseRetrievalStreamEvent event = withPayload(converted); - Object value = payload(event); - assertSame(value, payload(converted)); - assertEvent(event, eventName, converted.isTerminal()); - assertTrue(Modifier.isAbstract(KnowledgeBaseRetrievalStreamEvent.class.getModifiers())); - assertFalse(JsonSerializable.class.isAssignableFrom(event.getClass())); - assertFalse(Arrays.stream(event.getClass().getMethods()) - .anyMatch(method -> "fromJson".equals(method.getName()) || method.getName().startsWith("toJson"))); - assertThrows(NoSuchMethodException.class, () -> event.getClass().getMethod("getRawValue")); - assertThrows(NoSuchMethodException.class, - () -> KnowledgeBaseRetrievalStreamEvent.class.getMethod("getRawValue")); - assertEquals(1, event.getClass().getDeclaredConstructors().length); - assertEquals(1, event.getClass().getDeclaredConstructors()[0].getParameterCount()); - assertEquals(1, KnowledgeBaseRetrievalStreamEvent.class.getDeclaredConstructors().length); - assertNotNull(KnowledgeBaseRetrievalStreamEvent.class.getDeclaredConstructor(String.class)); - assertThrows(NoSuchMethodException.class, - () -> KnowledgeBaseRetrievalStreamEvent.class.getDeclaredConstructor(String.class, String.class)); - assertEquals(1, - Arrays.stream(KnowledgeBaseRetrievalStreamEvent.class.getDeclaredFields()) - .filter(field -> !field.isSynthetic()) - .count()); - assertNotNull(KnowledgeBaseRetrievalStreamEvent.class.getDeclaredField("eventName")); - assertThrows(NoSuchFieldException.class, - () -> KnowledgeBaseRetrievalStreamEvent.class.getDeclaredField("rawValue")); - if ("activity.completed".equals(eventName)) { - assertInstanceOf(KnowledgeBaseSearchIndexActivityRecord.class, value); - } - if ("references.completed".equals(eventName)) { - assertInstanceOf(KnowledgeBaseSearchIndexReference.class, - ((KnowledgeBaseReferencesCompletedStreamEvent) event).getValue().get(0)); - } - } - - @Test - public void nullPayloadsAndEmptyReferenceListsRemainAccessible() { - List absent = Arrays.asList( - new KnowledgeBaseRetrievalStartedStreamEvent(null), new KnowledgeBaseActivityStartedStreamEvent(null), - new KnowledgeBaseActivityCompletedStreamEvent(null), new KnowledgeBaseAnswerCompletedStreamEvent(null), - new KnowledgeBaseReferencesCompletedStreamEvent(null), new KnowledgeBaseErrorStreamEvent(null), - new KnowledgeBaseResponseCompletedStreamEvent(null)); - absent.forEach(event -> assertNull(payload(event))); - List empty = Collections.emptyList(); - assertSame(empty, new KnowledgeBaseReferencesCompletedStreamEvent(empty).getValue()); - assertEvent(new KnowledgeBaseRetrievalStreamEvent("custom") { - }, "custom", false); - } - - @Test - public void constructorsAndPayloadAccessNeverSerialize() { - AtomicInteger writes = new AtomicInteger(); - KnowledgeBaseActivityRecord activity = new KnowledgeBaseActivityRecord(0) { - @Override - public JsonWriter toJson(JsonWriter writer) throws IOException { - writes.incrementAndGet(); - throw new IOException("The event wrapper must not serialize its payload."); - } - }; - KnowledgeBaseActivityCompletedStreamEvent event = new KnowledgeBaseActivityCompletedStreamEvent(activity); - assertSame(activity, event.getValue()); - assertSame(activity, event.getValue()); - assertEquals(0, writes.get()); - - AtomicInteger referenceWrites = new AtomicInteger(); - KnowledgeBaseReference reference = new KnowledgeBaseReference("reference", 0) { - @Override - public JsonWriter toJson(JsonWriter writer) throws IOException { - referenceWrites.incrementAndGet(); - throw new IOException("The event wrapper must not serialize its payload."); - } - }; - List references = Collections.singletonList(reference); - KnowledgeBaseReferencesCompletedStreamEvent array = new KnowledgeBaseReferencesCompletedStreamEvent(references); - assertSame(references, array.getValue()); - assertSame(reference, array.getValue().get(0)); - assertEquals(0, referenceWrites.get()); - } - @ParameterizedTest @MethodSource("knownEvents") public void everyKnownVariantRejectsMalformedData(String eventName, String json) { @@ -450,9 +370,10 @@ public void clientsEmitErrorEventInclusivelyWithoutParsingFollowingMalformedPayl .verifyComplete(); } - @Test - public void clientsPropagateMalformedKnownPayload() throws IOException { - String body = "event: response.completed\ndata: {\n\n"; + @ParameterizedTest + @MethodSource("knownEvents") + public void clientsPropagateMalformedKnownPayload(String eventName, String json) throws IOException { + String body = frame(eventName, "references.completed".equals(eventName) ? "[{" : "{"); try (CloseableIterableStream stream = createBuilder(body, null).buildClient().retrieveStream(new KnowledgeBaseRetrievalOptions())) { @@ -507,7 +428,7 @@ private static void assertStreamEvents(List e assertEquals(3, events.size()); assertInstanceOf(KnowledgeBaseRetrievalStartedStreamEvent.class, events.get(0)); assertEvent(events.get(0), "retrieval.started", false); - assertEvent(events.get(1), "future.event", false); + assertUnknownEvent(events.get(1), "future.event"); assertInstanceOf(KnowledgeBaseResponseCompletedStreamEvent.class, events.get(2)); assertEvent(events.get(2), "response.completed", true); } @@ -567,40 +488,6 @@ private static Object payload(KnowledgeBaseRetrievalStreamEvent event) { } } - private static KnowledgeBaseRetrievalStreamEvent withPayload(KnowledgeBaseRetrievalStreamEvent event) { - switch (event.getEventName()) { - case "retrieval.started": - return new KnowledgeBaseRetrievalStartedStreamEvent( - ((KnowledgeBaseRetrievalStartedStreamEvent) event).getValue()); - - case "activity.started": - return new KnowledgeBaseActivityStartedStreamEvent( - ((KnowledgeBaseActivityStartedStreamEvent) event).getValue()); - - case "activity.completed": - return new KnowledgeBaseActivityCompletedStreamEvent( - ((KnowledgeBaseActivityCompletedStreamEvent) event).getValue()); - - case "answer.completed": - return new KnowledgeBaseAnswerCompletedStreamEvent( - ((KnowledgeBaseAnswerCompletedStreamEvent) event).getValue()); - - case "references.completed": - return new KnowledgeBaseReferencesCompletedStreamEvent( - ((KnowledgeBaseReferencesCompletedStreamEvent) event).getValue()); - - case "error": - return new KnowledgeBaseErrorStreamEvent(((KnowledgeBaseErrorStreamEvent) event).getValue()); - - case "response.completed": - return new KnowledgeBaseResponseCompletedStreamEvent( - ((KnowledgeBaseResponseCompletedStreamEvent) event).getValue()); - - default: - throw new IllegalArgumentException("Expected a known event: " + event.getEventName()); - } - } - private static KnowledgeBaseRetrievalClientBuilder createBuilder(Flux responseBody) { HttpHeaders headers = new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, "text/event-stream"); return new KnowledgeBaseRetrievalClientBuilder().endpoint("https://example.search.windows.net") @@ -635,6 +522,13 @@ private static String streamWithUnknownEvent() { + "data: " + RESPONSE_COMPLETED_JSON + "\n\n" + "event: response.completed\ndata: {\n\n"; } + private static void assertUnknownEvent(KnowledgeBaseRetrievalStreamEvent event, String eventName) { + assertEquals(KnowledgeBaseRetrievalStreamEvent.class, event.getClass().getSuperclass()); + assertTrue(event.getClass().isAnonymousClass()); + assertFalse(Modifier.isPublic(event.getClass().getModifiers())); + assertEvent(event, eventName, false); + } + private static void assertEvent(KnowledgeBaseRetrievalStreamEvent event, String eventName, boolean terminal) { assertEquals(eventName, event.getEventName()); assertEquals(terminal, event.isTerminal());