diff --git a/codecov.yaml b/codecov.yaml index 1b7bcf19df9..1219d81f838 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -3,9 +3,9 @@ codecov: require_ci_to_pass: no # Wait for expected number of coverage uploads before sending notifications # This prevents premature comments with incomplete data. - # 54 = UnitTests 40 + PythonTests 7 + RTests 6 + WebsiteSamplesTests 1 + # 55 = UnitTests 41 + PythonTests 7 + RTests 6 + WebsiteSamplesTests 1 # (every leg of those jobs runs templates/codecov.yml on succeededOrFailed). - after_n_builds: 54 + after_n_builds: 55 comment: layout: "reach,diff,flags,files" @@ -14,7 +14,7 @@ comment: require_base: false require_head: true # Wait for all expected uploads before commenting (see note above) - after_n_builds: 54 + after_n_builds: 55 coverage: precision: 2 diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala index 293efeec448..6f8f486bfb1 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala @@ -329,7 +329,9 @@ object URLEncodingUtils { } private[ml] object ServiceAuthHeaders { - private[ml] case class Resolution(headers: Map[String, String], usesFabricFallback: Boolean) + private[ml] case class Resolution(headers: Map[String, String], + usesFabricFallback: Boolean, + usesSubscriptionKey: Boolean) private[ml] def nonBlank(value: String): Boolean = value != null && value.trim.nonEmpty @@ -377,20 +379,15 @@ private[ml] object ServiceAuthHeaders { def isAuthHeaderName(name: String): Boolean = name.equalsIgnoreCase(subscriptionKeyHeaderName) || name.equalsIgnoreCase(aadHeaderName) - // Resolve exactly one auth header across every credential source in a single, case-insensitive - // precedence step. Blank values are skipped so they cannot suppress a valid lower-priority - // credential, and the automatic Fabric fallback ranks below a credential embedded in - // customHeaders (matching how management-index requests, which have no fallback, resolve auth). - // fabricFallbackAuthHeader is by-name and lowest priority, so it is evaluated only when every - // higher-priority source above is absent. A fallback that acquires a token (and may throw) is - // therefore never run while a subscription key, AAD token, explicit custom-auth header, or an - // embedded customHeaders credential is present. - val explicitAuthHeader: Option[(String, String)] = subscriptionKey.filter(nonBlank) - .map(value => subscriptionKeyHeaderName -> value) - .orElse(aadToken.filter(nonBlank).map(value => aadHeaderName -> ("Bearer " + value))) - .orElse(customAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value)) - .orElse(embeddedCredential(providedCustomHeaders, subscriptionKeyHeaderName)) - .orElse(embeddedCredential(providedCustomHeaders, aadHeaderName)) + // Select one non-blank credential and retain its source even when the auth header names coincide. + // The by-name Fabric fallback is evaluated only when no explicit or embedded credential exists. + val explicitAuth: Option[((String, String), Boolean)] = subscriptionKey.filter(nonBlank) + .map(value => (subscriptionKeyHeaderName -> value) -> true) + .orElse(aadToken.filter(nonBlank).map(value => (aadHeaderName -> ("Bearer " + value)) -> false)) + .orElse(customAuthHeader.filter(nonBlank).map(value => (aadHeaderName -> value) -> false)) + .orElse(embeddedCredential(providedCustomHeaders, subscriptionKeyHeaderName).map(_ -> true)) + .orElse(embeddedCredential(providedCustomHeaders, aadHeaderName).map(_ -> false)) + val explicitAuthHeader = explicitAuth.map(_._1) val fabricAuthHeader = if (explicitAuthHeader.isDefined) { None } else { @@ -421,7 +418,8 @@ private[ml] object ServiceAuthHeaders { Resolution( new scala.collection.immutable.TreeMap[String, String]() ++ headers, - usesFabricFallback = fabricAuthHeader.isDefined) + usesFabricFallback = fabricAuthHeader.isDefined, + usesSubscriptionKey = explicitAuth.exists(_._2)) } def build(subscriptionKey: Option[String], diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIColumnPlan.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIColumnPlan.scala new file mode 100644 index 00000000000..e9ec7301122 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIColumnPlan.scala @@ -0,0 +1,83 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.sql.{Column, Row, functions => F} +import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructField, StructType} + +import java.util.{HashSet, Set => JavaSet, TreeSet} + +private[ai] final case class AIColumnPlan(input: StructType, outputs: Seq[StructField]) { + AIColumnPlan.distinct(input.fieldNames.toSeq, "input columns") + AIColumnPlan.distinct(outputs.map(_.name), "output roles") + require(outputs.forall(_.name.nonEmpty), "output roles must have non-empty names") + private val outputNames = outputs.map(_.name).toSet + private val replacements = input.fields.map(f => + outputs.find(o => AIColumnPlan.matches(f.name, o.name))) + val schema: StructType = StructType(input.fields.zip(replacements).map { + case (original, replacement) => replacement.getOrElse(original) + } ++ outputs.filterNot(o => input.fields.exists(f => AIColumnPlan.matches(f.name, o.name)))) + + private val indices = schema.fields.map(f => + if (outputNames(f.name)) -1 else input.fieldIndex(f.name)) + + def assemble(row: Row, values: Map[String, Any]): Row = + Row.fromSeq(schema.fields.zip(indices).map { case (field, index) => + if (index < 0) values.getOrElse(field.name, None.orNull) else row.get(index) + }) + + def scratch(prefix: String, additional: String*): String = + AIColumnPlan.unused(prefix, schema.fieldNames.toSeq ++ input.fieldNames ++ additional) +} + +private[ai] object AIColumnPlan { + def matches(left: String, right: String): Boolean = SQLConf.get.resolver(left, right) + + def distinct(names: Seq[String], description: String): Unit = { + require(names.forall(_ != null), s"$description must not have null names") + // This comparator matches Spark's equalsIgnoreCase, including Unicode variants lost by lowercasing. + val seen: JavaSet[String] = if (SQLConf.get.caseSensitiveAnalysis) new HashSet[String]() + else new TreeSet[String](String.CASE_INSENSITIVE_ORDER) + require(!names.exists(name => !seen.add(name)), + s"Ambiguous $description under Spark's column resolver") + } + + def resolve(schema: StructType, name: String): String = + schema.fieldNames.find(matches(_, name)).getOrElse( + throw new IllegalArgumentException(s"Input column '$name' does not exist")) + + def unused(prefix: String, names: Seq[String]): String = + (Iterator(prefix) ++ Iterator.from(1).map(i => s"${prefix}_$i")) + .find(candidate => !names.exists(matches(_, candidate))).get + + def column(name: String): Column = F.col("`" + name.replace("`", "``") + "`") + + // Spark can eliminate serialization between Dataset maps and retain external JVM arrays. + def arrayValues(value: Any): scala.collection.Seq[Any] = value match { + case values: scala.collection.Seq[_] => values + case values: Array[_] => values.toSeq + case _ => throw new IllegalArgumentException("Expected a Spark array value") + } + + // Fused Dataset maps can pass schema-less Rows directly to the protocol serializers. + def withSchema(row: Row, schema: StructType): Row = + new GenericRowWithSchema(schema.fields.zipWithIndex.map { case (field, index) => + nestedSchema(row.get(index), field.dataType) + }, schema) + + private def nestedSchema(value: Any, dataType: DataType): Any = Option(value).map { present => + dataType match { + case struct: StructType => withSchema(present.asInstanceOf[Row], struct) + case ArrayType(element, _) => + arrayValues(present).iterator.map(nestedSchema(_, element)).toVector + case MapType(key, element, _) => + present.asInstanceOf[scala.collection.Map[Any, Any]].iterator.map { case (k, v) => + nestedSchema(k, key) -> nestedSchema(v, element) + }.toMap + case _ => present + } + }.orNull +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIJson.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIJson.scala new file mode 100644 index 00000000000..89aa881d4d7 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIJson.scala @@ -0,0 +1,291 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.fasterxml.jackson.core.{JsonFactory, JsonParser, JsonProcessingException, JsonToken, StreamReadConstraints} +import com.fasterxml.jackson.core.exc.StreamConstraintsException +import com.microsoft.azure.synapse.ml.param.{RPrinter, ServiceParam} +import org.apache.spark.ml.param.Params +import spray.json._ + +import java.util.IdentityHashMap +import scala.collection.JavaConverters._ + +private[services] object AIJson { + private[ai] val MaxDepth = 100 + private val MaxNumberLength = 1000 + private val ResponsesEnvelopeDepth = 3 + private val Factory = new JsonFactory().enable(JsonParser.Feature.STRICT_DUPLICATE_DETECTION) + .setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(Int.MaxValue) + .maxNumberLength(MaxNumberLength).maxNestingDepth(MaxDepth + ResponsesEnvelopeDepth + 1).build()) + + private class OrderedMap[+V](entries: Vector[(String, V)], index: Map[String, Int]) + extends Map[String, V] with Serializable { + override def get(key: String): Option[V] = index.get(key).map(entries(_)._2) + override def iterator: Iterator[(String, V)] = entries.iterator + override def empty: OrderedMap[V] = new OrderedMap(Vector.empty, Map.empty) + // Scala's immutable Map contract requires these operator names. + //scalastyle:off method.name + override def +[V1 >: V](entry: (String, V1)): OrderedMap[V1] = index.get(entry._1) match { + case Some(position) => new OrderedMap(entries.updated(position, entry), index) + case None => new OrderedMap(entries :+ entry, index.updated(entry._1, entries.length)) + } + override def -(key: String): OrderedMap[V] = if (!index.contains(key)) this else { + val remaining = entries.filterNot(_._1 == key) + new OrderedMap(remaining, remaining.iterator.zipWithIndex.map { case ((name, _), position) => + name -> position + }.toMap) + } + //scalastyle:on method.name + } + + private[ai] def orderedMap[V](entries: Iterable[(String, V)]): Map[String, V] = + entries.foldLeft(new OrderedMap[V](Vector.empty, Map.empty))(_ + _) + + implicit val AnyFormat: RootJsonFormat[Any] = new RootJsonFormat[Any] { + override def write(value: Any): JsValue = fromAny(value) + + override def read(value: JsValue): Any = value match { + case JsObject(fields) => orderedMap(fields.toSeq.map { case (key, item) => key -> read(item) }) + case JsArray(items) => items.map(read) + case JsNumber(number) => number + case JsString(text) => text + case JsBoolean(boolean) => boolean + case JsNull => None.orNull + } + } + + private implicit val ObjectFormat: RootJsonFormat[Map[String, Any]] = new RootJsonFormat[Map[String, Any]] { + override def write(value: Map[String, Any]): JsValue = fromAny(value) + override def read(value: JsValue): Map[String, Any] = + orderedMap(obj(value).fields.toSeq.map { case (key, item) => key -> AnyFormat.read(item) }) + } + + class ObjectParam(parent: Params, name: String, doc: String, wireName: String) + extends ServiceParam[Map[String, Any]](parent, name, doc, binding => { + // Spark's validation error renderer dereferences an invalid null binding. + require(Option(binding).isDefined, "Object parameter binding cannot be null") + binding.fold(value => Option(value).isDefined, column => Option(column).isDefined) + }) { + override val payloadName: String = wireName + // Preserve the existing left/right encoding without the default parser's unordered objects. + override def jsonDecode(json: String): Either[Map[String, Any], String] = format.read(parseServiceParam(json)) + } + + def rJsonValue[T](value: Either[T, String]): String = + RPrinter(JsString(value.fold(v => fromAny(v).compactPrint, identity))) + + def parse(value: String): JsValue = parse(value, 0) + + // The persisted left/right envelope is not part of the provider payload's depth. + def parseServiceParam(value: String): JsValue = parse(value, -1) + + // Responses echoes a schema inside text.format.schema; these wrappers are not caller schema levels. + def parseResponsesResponse(value: String): JsValue = parse(value, -ResponsesEnvelopeDepth) + + private def parse(value: String, initialDepth: Int): JsValue = { + require(Option(value).isDefined, "Expected JSON, not null") + val parser = Factory.createParser(value) + try { + val result = read(parser, parser.nextToken(), initialDepth) + require(Option(parser.nextToken()).isEmpty, "JSON must contain exactly one value") + result + } catch { + case _: StreamConstraintsException => + throw new IllegalArgumentException("JSON exceeds supported parsing limits (nesting or numeric token size)") + case _: JsonProcessingException => + throw new IllegalArgumentException("Invalid JSON (including duplicate keys or non-finite numbers)") + case _: NumberFormatException => + throw new IllegalArgumentException("Invalid JSON number") + } finally { + parser.close() + } + } + + private def read(parser: JsonParser, token: JsonToken, depth: Int): JsValue = { + require(depth <= MaxDepth, "JSON nesting exceeds the supported depth") + token match { + case JsonToken.START_OBJECT => + val fields = Iterator.continually(parser.nextToken()).takeWhile(_ != JsonToken.END_OBJECT).map { next => + require(next == JsonToken.FIELD_NAME, "Expected a JSON object field") + val key = validatedString(parser.currentName()) + key -> read(parser, parser.nextToken(), depth + 1) + }.toVector + JsObject(orderedMap(fields)) + case JsonToken.START_ARRAY => + val values = Iterator.continually(parser.nextToken()).takeWhile(_ != JsonToken.END_ARRAY).map { next => + require(Option(next).isDefined, "Unterminated JSON array") + read(parser, next, depth + 1) + }.toVector + JsArray(values) + case JsonToken.VALUE_STRING => JsString(validatedString(parser.getText)) + case JsonToken.VALUE_NUMBER_INT | JsonToken.VALUE_NUMBER_FLOAT => + val number = BigDecimal(parser.getText) + validateRenderedExponent(number) + JsNumber(number) + case JsonToken.VALUE_TRUE => JsBoolean(true) + case JsonToken.VALUE_FALSE => JsBoolean(false) + case JsonToken.VALUE_NULL => JsNull + case _ => throw new IllegalArgumentException("Expected a JSON value") + } + } + + def obj(value: JsValue): JsObject = value match { + case result: JsObject => result + case _ => throw new IllegalArgumentException("Expected a JSON object") + } + + def objectString(value: String): JsObject = obj(parse(value)) + + def fromAny(value: Any): JsValue = new ValueConverter().convert(value, 0) + + private def hasPairedSurrogates(value: String): Boolean = + value.indices.forall { index => + val character = value.charAt(index) + if (Character.isHighSurrogate(character)) { + index + 1 < value.length && Character.isLowSurrogate(value.charAt(index + 1)) + } else if (Character.isLowSurrogate(character)) { + index > 0 && Character.isHighSurrogate(value.charAt(index - 1)) + } else true + } + + private[ai] def validatedString(value: String): String = + validatedText(value, "JSON strings and keys") + + private[ai] def validatedText(value: String, description: String): String = { + require(hasPairedSurrogates(value), s"$description must not contain unpaired surrogates") + value + } + + // Consume one value from the current token without replacing the caller's JSON syntax or limits. + def hasValidStringTokens(parser: JsonParser): Boolean = { + var depth = 0 + val tokens = Iterator.single(parser.currentToken()) ++ + Iterator.continually(()).takeWhile(_ => depth > 0).map(_ => parser.nextToken()) + tokens.forall { + case JsonToken.START_OBJECT | JsonToken.START_ARRAY => + depth += 1 + true + case JsonToken.END_OBJECT | JsonToken.END_ARRAY => + depth -= 1 + true + case JsonToken.VALUE_STRING | JsonToken.FIELD_NAME => hasPairedSurrogates(parser.getText) + case token => Option(token).isDefined + } + } + + private def validateRenderedExponent(number: BigDecimal): Unit = + require(number.bigDecimal.precision.toLong - number.bigDecimal.scale.toLong - 1 <= Int.MaxValue, + "JSON number exponent exceeds supported parsing limits") + + private def validatedNumber(value: JsNumber): JsNumber = { + validateRenderedExponent(value.value) + val encoded = value.compactPrint + // Let the parser handle sign, decimal and exponent accounting at the boundary. + if (encoded.length > MaxNumberLength) parse(encoded) + value + } + + private class ValueConverter { + private val ancestors = new IdentityHashMap[AnyRef, java.lang.Boolean]() + private def guarded(ref: AnyRef, depth: Int)(f: => JsValue): JsValue = { + require(depth <= MaxDepth, "JSON nesting exceeds the supported depth") + require(!ancestors.containsKey(ref), "Cyclic JSON values are not supported") + ancestors.put(ref, true) + try f finally ancestors.remove(ref) + } + private def objectFields(values: Iterable[(Any, Any)], depth: Int): JsObject = + JsObject(orderedMap(values.toSeq.map { + case (key: String, item) => validatedString(key) -> convert(item, depth + 1) + case _ => throw new IllegalArgumentException("JSON object keys must be strings") + })) + def convert(item: Any, depth: Int): JsValue = { + require(depth <= MaxDepth, "JSON nesting exceeds the supported depth") + Option(item).fold[JsValue](JsNull) { + case json: JsValue => jsonValue(json, depth) + case text: String => JsString(validatedString(text)) + case boolean: Boolean => JsBoolean(boolean) + case number: Double => + require(!number.isNaN && !number.isInfinity, "JSON numbers must be finite") + validatedNumber(JsNumber(number)) + case number: Float => + require(!number.isNaN && !number.isInfinity, "JSON numbers must be finite") + validatedNumber(JsNumber(BigDecimal(java.lang.Float.toString(number)))) + case number: java.lang.Number => validatedNumber(JsNumber(BigDecimal(number.toString))) + case other => collection(other, depth) + } + } + private def jsonValue(value: JsValue, depth: Int): JsValue = value match { + case JsObject(fields) => objectFields(fields.toSeq, depth) + case JsArray(values) => JsArray(values.map(convert(_, depth + 1))) + case number: JsNumber => validatedNumber(number) + case JsString(text) => JsString(validatedString(text)) + case scalar => scalar + } + private def collection(item: Any, depth: Int): JsValue = item match { + case map: scala.collection.Map[_, _] => + guarded(map, depth)(objectFields(map.toSeq, depth)) + case map: java.util.Map[_, _] => + guarded(map, depth)(objectFields(map.asScala.toSeq, depth)) + case items: scala.collection.Seq[_] => + guarded(items, depth)(JsArray(items.map(convert(_, depth + 1)).toVector)) + case items: java.util.List[_] => + guarded(items, depth)(JsArray(items.asScala.map(convert(_, depth + 1)).toVector)) + case items: Array[_] => + guarded(items, depth)(JsArray(items.map(convert(_, depth + 1)).toVector)) + case _ => throw new IllegalArgumentException("Unsupported JSON value type") + } + } +} + +private[services] object AIExtraBody { + private val Reserved = Set( + "model", "messages", "input", "state", "questions", "url", "endpoint", "base_url", "api_version", + "api_key", "subscription_key", "subscriptionKey", "AADToken", "CustomAuthHeader", "customHeaders", + "headers", "extra_headers", "extra_query", "query", "timeout", "max_retries", "retry_policy", + "outputPath", "outputBatchSize", "output_path", "output_batch_size", "customUrlRoot") + private val TokenAliases = Set("max_tokens", "max_completion_tokens", "max_output_tokens") + + def merge(base: JsObject, extension: JsObject, owned: Set[List[String]], api: String, + lossless: Boolean): JsObject = { + val reserved = extension.fields.keySet.intersect(Reserved) + require(reserved.isEmpty, s"extraBody contains reserved fields: ${reserved.toSeq.sorted.mkString(", ")}") + val extraTokens = extension.fields.keySet.intersect(TokenAliases) + require(extraTokens.size <= 1, "extraBody contains conflicting token aliases") + if (extraTokens.nonEmpty) { + require(!owned.exists(p => p.nonEmpty && TokenAliases(p.head)), + "extraBody conflicts with explicit typed option: max_tokens/max_completion_tokens/max_output_tokens") + } + val initial = if (extraTokens.nonEmpty) JsObject(base.fields -- TokenAliases) else base + def combine(left: JsObject, right: JsObject, path: List[String]): JsObject = { + JsObject(right.fields.foldLeft(AIJson.orderedMap(left.fields.toSeq)) { case (result, (key, value)) => + val current = path :+ key + require(!owned.exists(p => current.startsWith(p)), + s"extraBody conflicts with explicit typed option: ${current.mkString(".")}") + val merged = (result.get(key), value) match { + case (Some(a: JsObject), b: JsObject) => combine(a, b, current) + case (None, b: JsObject) => combine(JsObject(), b, current) + case _ => + require(!owned.exists(_.startsWith(current)), + s"extraBody conflicts with explicit typed option: ${current.mkString(".")}") + value + } + result.updated(key, merged) + }) + } + val result = combine(initial, extension, Nil) + Seq("stream", "background").foreach { key => + result.fields.get(key).foreach(v => + require(v == JsFalse, s"$key=true is not supported; only synchronous responses are supported")) + } + require(result.fields.keySet.intersect(TokenAliases).size <= 1, "Request contains conflicting token aliases") + if (api == "embeddings") { + result.fields.get("encoding_format").foreach(v => + require(v == JsString("float"), "embeddings only supports encoding_format='float'")) + } + require(lossless || !result.fields.keySet.exists(Set("tools", "tool_choice", "functions", "function_call")), + "Tool requests require rawResponseCol; AIPrompt does not execute tools") + result + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPrompt.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPrompt.scala new file mode 100644 index 00000000000..77825625688 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPrompt.scala @@ -0,0 +1,567 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.spark.PEP3101Parser +import com.microsoft.azure.synapse.ml.fabric.FabricClient +import com.microsoft.azure.synapse.ml.io.http._ +import com.microsoft.azure.synapse.ml.logging.common.PlatformDetails +import com.microsoft.azure.synapse.ml.param.{GlobalParams, ServiceParam} +import com.microsoft.azure.synapse.ml.services.ServiceAuthHeaders +import com.microsoft.azure.synapse.ml.services.openai.{JsonParser => OpenAIJsonParser, _} +import org.apache.hadoop.conf.Configuration +import org.apache.http.client.methods.HttpPost +import org.apache.http.entity.{ContentType, StringEntity} +import org.apache.spark.injections.UDFUtils +import org.apache.spark.ml.{ExplicitComplexParamsReadable, ExplicitComplexParamsWritable, ParamInjections, Serializer} +import org.apache.spark.ml.param.{Param, ParamMap} +import org.apache.spark.ml.util.Identifiable +import org.apache.spark.sql.{AnalysisException, Column, DataFrame, Dataset, Encoders, Row, SparkSession, functions => F} +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema +import org.apache.spark.sql.catalyst.plans.logical.Project +import org.apache.spark.sql.types._ +import org.apache.spark.util.SerializableConfiguration +import spray.json._ + +import java.nio.ByteBuffer +import java.nio.charset.{CharacterCodingException, CodingErrorAction, StandardCharsets} +import java.util.concurrent.CancellationException +import scala.language.existentials + +object AIPrompt extends ExplicitComplexParamsReadable[AIPrompt] { + override protected def expectedClass: Class[AIPrompt] = classOf[AIPrompt] + + private val HeaderNamePattern = "[!#$%&'*+.^_`|~0-9A-Za-z-]+".r.pattern + private val FirstLatin1HeaderCharacter = 0xa0 + private val LastLatin1HeaderCharacter = 0xff + + private def validateHeaderName(name: String): Unit = { + require(name != null && HeaderNamePattern.matcher(name).matches(), "Invalid HTTP header name") + require(!"Content-Length".equalsIgnoreCase(name) && !"Transfer-Encoding".equalsIgnoreCase(name), + "HTTP framing headers Content-Length and Transfer-Encoding are managed by the client; omit them") + } + + private def validateHeader(name: String, value: String): Unit = { + validateHeaderName(name) + require(value != null && value.forall(c => + c == '\t' || (c >= ' ' && c <= '~') || + (c >= FirstLatin1HeaderCharacter && c <= LastLatin1HeaderCharacter)), + "Invalid HTTP header value") + } + + private def validateCallerHeaders(headers: scala.collection.Map[String, String]): Unit = + headers.keys.foreach(validateHeaderName) + + private def transportHeaders(prompt: AIPrompt, address: String, + resolution: ServiceAuthHeaders.Resolution): Seq[(String, String)] = { + val headers = if (prompt.usingImplicitFabricEndpoint(address)) { + OpenAIFabricHeaders.merge(Some(resolution.headers)) + } else resolution.headers + val bearerKey = prompt.getApiType == "decisions" || + AIRequestAdapter.isOpenAIHost(AIRequestAdapter.normalizedHost(new java.net.URI(address))) + val outgoing = headers.toSeq.filterNot { case (name, _) => HTTPRequestData.isFabricAuthMarker(name) } + .filterNot(_._1.equalsIgnoreCase("Content-Type")) + .map { case (name, value) => + if (bearerKey && resolution.usesSubscriptionKey && name.equalsIgnoreCase(prompt.subscriptionKeyHeaderName)) { + prompt.aadHeaderName -> ("Bearer " + value) + } else name -> value + } + val result = outgoing :+ ("Content-Type" -> "application/json") + if (resolution.usesFabricFallback) result :+ (HTTPRequestData.FabricAuthMarkerHeader -> "true") else result + } + + private def requestHeaders(prompt: AIPrompt, address: String, + string: ServiceParam[String] => Option[String], + map: ServiceParam[Map[String, String]] => Option[Map[String, String]], + includeFabricFallback: Boolean): Seq[(String, String)] = { + val custom = map(prompt.customHeaders) + val telemetry = map(prompt.telemHeaders) + Seq(custom, telemetry).flatten.foreach(validateCallerHeaders) + lazy val implicitAuth = if (includeFabricFallback && prompt.usingImplicitFabricEndpoint(address)) { + Option(prompt.fabricAuthHeader) + } else None + val resolution = ServiceAuthHeaders.resolve(string(prompt.subscriptionKey), prompt.subscriptionKeyHeaderName, + prompt.aadHeaderName, string(prompt.AADToken), string(prompt.CustomAuthHeader), custom, + implicitAuth, telemetry, Some("application/json")) + transportHeaders(prompt, address, resolution) + } + + private[ai] def validateScalarHeaders(prompt: AIPrompt, address: String): Unit = { + def scalar[T](param: ServiceParam[T]): Option[T] = + prompt.get(param).orElse(prompt.getDefault(param)).flatMap(_.left.toOption).flatMap(Option(_)) + def column[T](param: ServiceParam[T]): Boolean = + prompt.get(param).orElse(prompt.getDefault(param)).exists(_.isRight) + // A higher-priority column can suppress a scalar fallback; only its rows can validate that choice. + val conditionalAuth = Seq(prompt.subscriptionKey, prompt.AADToken, prompt.CustomAuthHeader) + .takeWhile(param => !scalar(param).exists(ServiceAuthHeaders.nonBlank)).exists(column(_)) + val conditionalTelemetry = column(prompt.telemHeaders) + requestHeaders(prompt, address, scalar[String], scalar[Map[String, String]], includeFabricFallback = false) + .filter { case (name, _) => + if (name.equalsIgnoreCase(prompt.subscriptionKeyHeaderName) || name.equalsIgnoreCase(prompt.aadHeaderName)) { + !conditionalAuth + } else !conditionalTelemetry + }.foreach { case (name, value) => validateHeader(name, value) } + } + + private val FirstRetryMs = 100 + private val SecondRetryMs = 500 + private val ThirdRetryMs = 1000 + private val RetryIntervals = Array(FirstRetryMs, SecondRetryMs, ThirdRetryMs) + private[ai] val DefaultHandler = UDFUtils.oldUdf(HandlingUtils.bounded(RetryIntervals: _*) _, StringType) + private[ai] def usesBoundedHandler(stage: AIPrompt): Boolean = + !stage.isSet(stage.handler) && stage.getDefault(stage.handler).contains(DefaultHandler) + + private val MessageSchema = Encoders.product[OpenAICompositeMessage].schema + private val MessageType = ArrayType(MessageSchema, containsNull = true) + + private[ai] def createHttpTransformer(stage: AIPrompt): HTTPTransformer = + new HTTPTransformer().setConcurrency(stage.getConcurrency).setTimeout(stage.getTimeout) + .setConcurrentTimeout(stage.get(stage.concurrentTimeout).orElse(stage.getDefault(stage.concurrentTimeout))) + .setHandler(stage.getHandler) + + private[ai] def withSession[T](spark: SparkSession)(action: => T): T = { + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + try action finally { + previous match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } + + private[ai] def withSafeTemplate[T](prepare: => T): T = try { + prepare + } catch { + // Parser and analyzer exceptions can contain the full prompt, including in their causes. + case _: IllegalArgumentException | _: AnalysisException => + throw new IllegalArgumentException("Invalid prompt template; check braces, expressions and input columns") + } + + private[ai] def bindTemplate(input: DataFrame, expression: Column): Column = { + // A command's analyzed plan differs from the CommandResult used by subsequent Dataset operations. + val base = input.select(input.columns.map(AIColumnPlan.column): _*) + withSafeTemplate(base.select(expression).queryExecution.analyzed) match { + case Project(Seq(resolved), child) if child == base.queryExecution.analyzed => new Column(resolved) + case _ => throw new IllegalArgumentException( + "Prompt templates require row-local expressions; materialize windows, aggregates and generators " + + "into input columns") + } + } +} + +private[ai] trait HasAIResponseSchema extends HasOpenAIResponseSchema { + override protected def validateResponseSchemaValues(schema: Map[String, Any]): Unit = { + AIJson.fromAny(schema) + } + + override protected def pyParamInputType(p: Param[_]): Option[String] = + AIPromptPython.paramInputType(p.name).orElse(super.pyParamInputType(p)) + + override protected def pyResponseSchemaToJava: String = + "self._java_obj.responseFormat().jsonDecode(self._ai_json({'left': schema})).left().get()" +} + +/** Distributed synchronous inference. Spark actions and recomputation may repeat paid requests. */ +class AIPrompt(override val uid: String) extends OpenAIPrompt(uid) + with AIPromptParams with HasHandler with ExplicitComplexParamsWritable with HasAIResponseSchema { + def this() = this(Identifiable.randomUID("AIPrompt")) + + override protected def createResponseFormatParam(name: String): ServiceParam[Map[String, Any]] = + new AIJson.ObjectParam(this, name, + "Response format for the completion. One of 'text', 'json_object', or 'json_schema'.", "response_format") + + setDefault(handler -> AIPrompt.DefaultHandler) + + override protected def supportedApiTypes: Array[String] = + AIRequestAdapter.SupportedApis.toArray + + override protected val subscriptionKeyHeaderName: String = "api-key" + + override protected def hasExplicitPostProcessing: Boolean = isSet(postProcessing) + + override protected val useBoundedAttachments: Boolean = true + + def setModelCol(value: String): this.type = setVectorParam(model, value) + def getModelCol: String = getVectorParam(model) + def setStoreCol(value: String): this.type = setVectorParam(store, value) + def getStoreCol: String = getVectorParam(store) + def setResponseFormatCol(value: String): this.type = setVectorParam(responseFormat, value) + def getResponseFormatCol: String = getVectorParam(responseFormat) + + override def setResponseFormat(value: Map[String, Any]): this.type = { + val formatted = AIRequestAdapter.responseFormat(value, "chat_completions") + setScalarParam(responseFormat, AIJson.orderedMap(formatted.fields.toSeq.map { case (key, entry) => + key -> AIJson.AnyFormat.read(entry) + })) + } + + override def copy(extra: ParamMap): AIPrompt = super.copy(extra).asInstanceOf[AIPrompt] + + override protected def validateForWrite(): Unit = { + // Explicit persistence replaces OpenAIPrompt.write, including its parser validation. + getParser + require(getDefault(handler).contains(AIPrompt.DefaultHandler), + "Saving AIPrompt requires the original handler default; use setHandler for a custom persisted handler. " + + s"Loading a custom handler requires ${Serializer.LegacyObjectDeserializationConfig}=true; " + + "enable this only for trusted artifacts.") + } + + private def setImplicitFabricEndpoint(): Unit = { + if (AIRequestAdapter.ImplicitFabricApis(getApiType) && + !isDefined(url) && !isDefined(customUrlRoot)) { + fabricEndpoint.foreach(endpoint => setDefault(url -> endpoint)) + } + } + + private def excludedGlobalDefaults(result: AIPrompt): Set[String] = { + val overriddenModelDefaults = if (isSet(model)) Set(deploymentName.name) + else if (isSet(deploymentName)) Set(model.name) else Set.empty[String] + val allowed = AIRequestAdapter.allowedOptions(result.getApiType) + val inactiveOptions = AIRequestAdapter.bodyOptions(result).filterNot(p => allowed(p.name)).map(_.name).toSet + val nativeDefaults = if (AIRequestAdapter.NativeApis(result.getApiType)) { + Set(model.name, deploymentName.name, url.name, subscriptionKey.name) + } else Set.empty[String] + val inactiveEndpointDefaults = if (result.isSet(customUrlRoot)) Set(url.name) else Set.empty[String] + Set(apiType.name, apiVersion.name) ++ overriddenModelDefaults ++ + inactiveOptions ++ nativeDefaults ++ inactiveEndpointDefaults + } + + private[ai] def effectiveCopy(): AIPrompt = { + val result = copy(ParamMap.empty) + if (!isSet(apiType)) { + GlobalParams.getParam(apiType).filter(value => supportedApiTypes.contains(value.toString)) + .foreach(value => result.setDefault(result.apiType -> value.toString)) + } + val excludedGlobals = excludedGlobalDefaults(result) + params.filterNot(p => isSet(p) || excludedGlobals(p.name)).foreach { p => + GlobalParams.getParam(p).foreach { value => + result.setDefault(result.getParam(p.name).asInstanceOf[Param[Any]] -> value) + } + } + result.setImplicitFabricEndpoint() + AIRequestAdapter.defaultApiVersion(result).foreach(value => result.setDefault(result.apiVersion -> value)) + if (result.getApiType == "embeddings" && !isSet(deploymentName) && !isSet(model)) { + GlobalParams.getGlobalParam(OpenAIEmbeddingDeploymentNameKey).foreach { v => + result.setDefault(result.deploymentName -> v) + } + } + if (result.isSet(result.model)) ParamInjections.clearDefault(result, result.deploymentName) + else if (result.isDefined(result.deploymentName)) ParamInjections.clearDefault(result, result.model) + result + } + + private[ai] def responseIdEnabled: Boolean = isSet(responseIdCol) || + (getApiType == "responses" && get(store).orElse(getDefault(store)).exists { + case Left(enabled) => enabled + case Right(_) => true + }) + + private def retainedMessagesName(schema: StructType): Option[String] = + if (AIRequestAdapter.GenerativeApis(getApiType) && !getDropPrompt && !isSet(messagesCol)) { + val outputs = Seq(getOutputCol, getErrorCol) ++ get(usageCol) ++ get(rawResponseCol) ++ + (if (responseIdEnabled) Seq(getResponseIdCol) else Seq.empty) + Some(AIColumnPlan.unused(getMessagesCol, schema.fieldNames.toSeq ++ outputs)) + } else None + + private def plan(schema: StructType, parsed: Boolean): AIColumnPlan = { + val responseType = if (parsed && AIRequestAdapter.GenerativeApis(getApiType)) { + AIValueCodec.nullable(getParser.outputSchema) + } else AIResponseCodec.outputType(getApiType) + val outputs = Seq(StructField(getOutputCol, responseType), StructField(getErrorCol, ErrorUtils.ErrorSchema)) ++ + get(usageCol).map(StructField(_, UsageUtils.UsageStructType)) ++ + get(rawResponseCol).map(StructField(_, StringType)) ++ + (if (responseIdEnabled) Seq(StructField(getResponseIdCol, StringType)) else Seq.empty) ++ + retainedMessagesName(schema).map(StructField(_, AIPrompt.MessageType)) + AIColumnPlan(schema, outputs) + } + + private[ai] def inferenceOutputSchema(schema: StructType): StructType = plan(schema, parsed = true).schema + + private[ai] def inferenceSchema(schema: StructType): StructType = { + AIValueCodec.validateMapKeys(schema) + val result = inferenceOutputSchema(schema) + new AIRequestAdapter(this, schema).validate() + result + } + + override def transformSchema(schema: StructType): StructType = { + val stage = effectiveCopy() + stage.inferenceSchema(schema) + } + + override def transform(dataset: Dataset[_]): DataFrame = AIPrompt.withSession(dataset.sparkSession) { + logTransform({ + val stage = effectiveCopy() + stage.inferenceSchema(dataset.schema) + stage.validateFabricModel(dataset.schema) + stage.infer(dataset) + }, dataset.columns.length) + } + + private def error(code: String, message: String, status: Option[Row] = None): Row = + Row(JsObject("code" -> JsString(code), "message" -> JsString(message)).compactPrint, status.orNull) + + private def attachmentColumns(schema: StructType): Seq[String] = + getColumnTypes.toSeq.collect { case (name, "path") => name } + .sortBy(name => (schema.fieldIndex(AIColumnPlan.resolve(schema, name)), name)) + + private def generatedMessages(row: Row, prompt: String, indices: Map[String, Int], paths: Seq[String], + conf: Option[Configuration]): Seq[Row] = { + val attachments = paths.map { name => + name -> row.getString(indices(name)) + }.toMap + Option(createMessagesForRow(prompt, attachments, paths, conf)).map(_.map { message => + new GenericRowWithSchema(Array[Any](message.role, message.content, message.name.orNull), + AIPrompt.MessageSchema): Row + }).getOrElse(throw new AIMissingInputException("attachment input is missing")) + } + + private def messages(row: Row, prompt: Option[String], schema: StructType, + indices: Map[String, Int], paths: Seq[String], + conf: Option[Configuration]): Option[Seq[Row]] = { + if (isSet(messagesCol)) { + val index = indices(getMessagesCol) + val messageSchema = schema(index).dataType.asInstanceOf[ArrayType].elementType.asInstanceOf[StructType] + Some(Option(row.get(index)).map(value => AIColumnPlan.arrayValues(value).iterator.map { message => + Option(message).map { + case struct: Row => AIColumnPlan.withSchema(struct, messageSchema) + case _ => throw new IllegalArgumentException("Expected a message struct") + }.orNull + }.toVector) + .getOrElse(throw new AIMissingInputException("messages is null"))) + } else if (AIRequestAdapter.GenerativeApis(getApiType)) { + val text = prompt.orElse(get(inputCol).flatMap(name => + Option(row.getString(indices(name))))) + .getOrElse(throw new AIMissingInputException("prompt is null")) + Some(generatedMessages(row, text, indices, paths, conf)) + } else None + } + + private def buildRequest(row: Row, address: String, body: String, adapter: AIRequestAdapter, + toRow: HTTPRequestData => Row): Row = { + val request = new HttpPost(address) + def string(p: ServiceParam[String]): Option[String] = adapter.value(row, p).map(_.asInstanceOf[String]) + def map(p: ServiceParam[Map[String, String]]): Option[Map[String, String]] = + adapter.value(row, p).map(_.asInstanceOf[scala.collection.Map[String, String]].toMap) + AIPrompt.requestHeaders(this, address, string, map, includeFabricFallback = true) + .foreach { case (name, value) => + AIPrompt.validateHeader(name, value) + request.addHeader(name, value) + } + request.setEntity(new StringEntity(AIJson.validatedString(body), ContentType.APPLICATION_JSON)) + toRow(new HTTPRequestData(request)) + } + + private[ai] def fabricEndpoint: Option[String] = + if (PlatformDetails.runningOnFabric()) Some(FabricClient.MLWorkloadEndpointOpenAI) else None + + private[ai] def fabricAuthHeader: String = FabricClient.getCognitiveMWCTokenAuthHeader + + private[ai] def usingImplicitFabricEndpoint(address: String): Boolean = + AIRequestAdapter.ImplicitFabricApis(getApiType) && + !isSet(url) && !isSet(customUrlRoot) && + fabricEndpoint.exists(root => FabricClient.isEndpointUnder(address, root.stripSuffix("/") + "/")) + + private[ai] def assertFabricModelStatus(modelName: String): Unit = + OpenAIFabricModelStatus.assertModelStatus(modelName) + + private[ai] def validateFabricModel(schema: StructType): Unit = + new AIRequestAdapter(this, schema).validateFabricModel() + + private def preparePrompt(input: DataFrame, promptName: String): DataFrame = { + if (!isSet(promptTemplate)) input + else { + val filenames = { + val filename = F.udf((path: String) => Option(path).map(_.trim).filter(_.nonEmpty).map { p => + try OpenAIAttachmentUtils.attachmentFilename(p) + catch { case _: IllegalArgumentException => "attachment" } + }.getOrElse("")) + attachmentColumns(input.schema).map { name => + name -> filename(AIColumnPlan.column(AIColumnPlan.resolve(input.schema, name))) + }.toMap + } + val expression = AIPrompt.withSafeTemplate(templateWithFilenames(filenames)) + // Bind against public columns without allowing the template to filter or expand rows. + val scoped = AIPrompt.bindTemplate(input, expression) + AIPrompt.withSafeTemplate(input.withColumn(promptName, scoped)) + } + } + + private def templateWithFilenames(filenames: Map[String, Column]): Column = { + val tokens = new PEP3101Parser(AIJson.validatedString(getPromptTemplate)).parse.map { token => + token.expr match { + case reference: UnresolvedAttribute if reference.nameParts.length == 1 => + filenames.collectFirst { + case (name, filename) if AIColumnPlan.matches(name, reference.nameParts.head) => filename + }.getOrElse(token) + case _ => token + } + } + F.concat(tokens: _*) + } + + private def prepareRequests(prepared: DataFrame, promptName: String, requestName: String, + preparationName: String, messagesName: String, + retainGeneratedMessages: Boolean): DataFrame = { + val inputSchema = prepared.schema + val promptIndex = if (isSet(promptTemplate)) Some(inputSchema.fieldIndex(promptName)) else None + val names = getColumnTypes.keys.toSeq ++ get(inputCol) ++ + (if (isSet(messagesCol)) Seq(getMessagesCol) else Seq.empty) + val indices = names.map(name => name -> inputSchema.fieldIndex(AIColumnPlan.resolve(inputSchema, name))).toMap + val paths = attachmentColumns(inputSchema) + val attachmentConf = if (paths.nonEmpty) { + Some(new SerializableConfiguration(prepared.sparkSession.sessionState.newHadoopConf())) + } else None + val adapter = new AIRequestAdapter(this, inputSchema) + val preparedSchema = inputSchema.add(requestName, HTTPSchema.Request) + .add(preparationName, ErrorUtils.ErrorSchema).add(messagesName, AIPrompt.MessageType) + prepared.mapPartitions { rows => + val toRequest = HTTPRequestData.makeToRowConverter + val conf = attachmentConf.map(_.value) + rows.map { row => + val prompt = promptIndex.flatMap(index => Option(row.getString(index))) + try { + val inputMessages = messages(row, prompt, inputSchema, indices, paths, conf) + val (address, body) = adapter.request(row, prompt, inputMessages) + Row.fromSeq(row.toSeq ++ Seq(buildRequest(row, address, body, adapter, toRequest), + Option.empty[Row].orNull, + if (retainGeneratedMessages) inputMessages.orNull else Option.empty[Seq[Row]].orNull)) + } catch { + case _: OpenAIAttachmentUtils.CancelledAttachmentReadException => + throw new CancellationException("Attachment read canceled") + case _: OpenAIAttachmentUtils.AttachmentReadException => + Row.fromSeq(row.toSeq ++ Seq(Option.empty[Row].orNull, + error("input_io_error", "Unable to read input attachment"), Option.empty[Seq[Row]].orNull)) + case e: AIMissingInputException => + Row.fromSeq(row.toSeq ++ Seq(Option.empty[Row].orNull, error("missing_input", e.getMessage), + Option.empty[Seq[Row]].orNull)) + case e: IllegalArgumentException => + Row.fromSeq(row.toSeq ++ Seq(Option.empty[Row].orNull, error("invalid_request", e.getMessage), + Option.empty[Seq[Row]].orNull)) + } + } + }(ExpressionEncoder(preparedSchema)) + } + + private def decodeResponse(response: HTTPResponseData, toStatus: StatusLineData => Row): Map[String, Any] = { + val (raw, decoded) = try { + val text = response.entity.map(e => StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT).onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(e.content)).toString).orNull + text -> AIResponseCodec.decode(getApiType, text) + } catch { + case _: CharacterCodingException => + Option.empty[String].orNull -> AIResponseCodec.Decoded(None.orNull, None.orNull, None.orNull, + Some("parse_error" -> "The provider response was not valid UTF-8")) + } + val exposed = get(rawResponseCol).map(_ -> raw).toMap[String, Any] + val observations = exposed ++ get(usageCol).map(_ -> decoded.usage) ++ + (if (responseIdEnabled) Map(getResponseIdCol -> decoded.responseId) else Map.empty) + if (response.statusLine.statusCode < 200 || response.statusLine.statusCode >= 300) { + observations + (getErrorCol -> error("http_error", s"HTTP ${response.statusLine.statusCode}", + Some(toStatus(response.statusLine)))) + } else { + observations ++ Map(getOutputCol -> decoded.value) ++ + decoded.diagnostic.map { case (code, message) => getErrorCol -> error(code, message) } + } + } + + private def decodeRows(responseRows: DataFrame, outputPlan: AIColumnPlan, responseName: String, + preparationName: String, messagesName: String, + retainedMessages: Option[String]): DataFrame = { + val responseIndex = responseRows.schema.fieldIndex(responseName) + val prepIndex = responseRows.schema.fieldIndex(preparationName) + val messagesIndex = responseRows.schema.fieldIndex(messagesName) + responseRows.mapPartitions { rows => + val fromResponse = HTTPResponseData.makeFromRowConverter + val toStatus = StatusLineData.makeToRowConverter + rows.map { row => + val messages = retainedMessages.map(_ -> row.get(messagesIndex)).toMap + val prepError = Option(row.getAs[Row](prepIndex)) + val values = if (prepError.isDefined) { + Map(getErrorCol -> prepError.get) + } else if (row.isNullAt(responseIndex)) { + Map(getErrorCol -> error("transport_error", "No HTTP response was received")) + } else { + decodeResponse(fromResponse(row.getStruct(responseIndex)), toStatus) + } + outputPlan.assemble(row, messages ++ values) + } + }(ExpressionEncoder(outputPlan.schema)) + } + + private def postprocess(rawResult: DataFrame, outputPlan: AIColumnPlan): DataFrame = { + if (AIRequestAdapter.GenerativeApis(getApiType) && getPostProcessing.nonEmpty) { + val scratch = outputPlan.scratch(s"${uid}_unparsed") + val parser = getParser + val parsedType = AIValueCodec.nullable(parser.outputSchema) + val parsed = rawResult.withColumn(scratch, AIColumnPlan.column(getOutputCol)) + .withColumn(getOutputCol, parser.parse(AIColumnPlan.column(scratch)).cast(parsedType)) + val invalid = parser match { + case json: OpenAIJsonParser => + AIColumnPlan.column(getOutputCol).isNull || json.hasParseError(AIColumnPlan.column(scratch)) + case _ => AIColumnPlan.column(getOutputCol).isNull + } + val failure = F.struct( + F.lit(error("parse_error", "Text does not match the configured postprocessing schema").getString(0)) + .as("response"), + F.lit(Option.empty[String].orNull).cast(StatusLineData.schema).as("status")) + parsed.withColumn(getErrorCol, F.when(AIColumnPlan.column(scratch).isNotNull && + invalid && AIColumnPlan.column(getErrorCol).isNull, failure) + .otherwise(AIColumnPlan.column(getErrorCol))).drop(scratch) + } else rawResult + } + + private[ai] def infer(dataset: Dataset[_]): DataFrame = AIPrompt.withSession(dataset.sparkSession) { + infer(dataset, AIPrompt.createHttpTransformer(this)) + } + + private[ai] def infer(dataset: Dataset[_], http: HTTPTransformer): DataFrame = { + AIPrompt.withSession(dataset.sparkSession) { + val original = dataset.toDF() + val outputPlan = plan(original.schema, parsed = false) + val retainedMessages = retainedMessagesName(original.schema) + val finalSchema = inferenceOutputSchema(original.schema) + new AIRequestAdapter(this, original.schema).validate() + val promptName = outputPlan.scratch(s"${uid}_prompt") + val requestName = outputPlan.scratch(s"${uid}_request", promptName) + val preparationName = outputPlan.scratch(s"${uid}_preparation", promptName, requestName) + val messagesName = outputPlan.scratch(s"${uid}_generated_messages", promptName, requestName, preparationName) + val responseName = outputPlan.scratch(s"${uid}_response", promptName, requestName, preparationName, messagesName) + val prepared = preparePrompt(original, promptName) + // Preserve SQL map-key identity while retained input traverses the inference Row encoders. + val protectedInput = AIValueCodec.encodeMapKeys(prepared) + val retainedSchema = StructType(original.schema.fields.map(field => protectedInput.schema(field.name))) + val protectedPlan = outputPlan.copy(input = retainedSchema) + val requestRows = prepareRequests(protectedInput, promptName, requestName, preparationName, messagesName, + retainedMessages.isDefined) + http.setInputCol(requestName).setOutputCol(responseName) + val responses = if (AIPrompt.usesBoundedHandler(this)) http.transformCancellable(requestRows) + else http.transform(requestRows) + val rawResult = decodeRows(responses, protectedPlan, responseName, + preparationName, messagesName, retainedMessages) + val processed = postprocess(rawResult, outputPlan) + AIValueCodec.reconcileSchema(AIValueCodec.decodeMapKeys(processed, finalSchema), finalSchema) + } + } + + override protected def pyParamSetter(p: Param[_]): String = + AIPromptPython.paramSetter(super.pyParamSetter(p), p.name) + + override protected val pyHasExplicitPostProcessing: String = "self.isSet(self.postProcessing)" + + override protected def pythonClass(): String = AIPromptPython.nativeReader(super.pythonClass()) + + override def pyAdditionalMethods: String = AIPromptPython.additionalMethods(super.pyAdditionalMethods, pyClassName) + + override def pyInitFunc(): String = AIPromptPython.initFunc(super.pyInitFunc()) + + override protected def rParamArg[T](p: Param[T]): String = p match { + case _: ServiceParam[_] => super.rParamArg(p) + case _ => s"${p.name}=NULL" + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptParams.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptParams.scala new file mode 100644 index 00000000000..5dfd5a5c8d6 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptParams.scala @@ -0,0 +1,93 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.contracts.HasInputCol +import com.microsoft.azure.synapse.ml.param.ServiceParam +import com.microsoft.azure.synapse.ml.services.HasServiceParams +import com.microsoft.azure.synapse.ml.services.ai.AIJson.AnyFormat +import org.apache.spark.ml.param.{Param, ParamValidators} +import spray.json.DefaultJsonProtocol._ +import spray.json.{JsObject, JsValue} + +trait AIPromptParams extends HasServiceParams with HasInputCol { + val extraBody: ServiceParam[Map[String, Any]] = new AIJson.ObjectParam( + this, "extraBody", + "Provider-owned JSON body extension; cannot override explicit typed options or request identity.", "extraBody") { + override def rValue(value: Either[Map[String, Any], String]): String = AIJson.rJsonValue(value) + } + + def setExtraBody(value: Map[String, Any]): this.type = setExtraBody(AIJson.obj(AIJson.fromAny(value))) + + def setExtraBody(value: java.util.Map[String, Object]): this.type = + setExtraBody(AIJson.obj(AIJson.fromAny(value))) + + def setExtraBody(value: JsObject): this.type = + setScalarParam(extraBody, AIJson.orderedMap(AIJson.obj(AIJson.fromAny(value)).fields.toSeq.map { + case (key, entry) => key -> AnyFormat.read(entry) + })) + + def setExtraBody(value: String): this.type = setExtraBody(AIJson.objectString(value)) + def getExtraBody: Map[String, Any] = getScalarParam(extraBody) + def setExtraBodyCol(value: String): this.type = setVectorParam(extraBody, value) + def getExtraBodyCol: String = getVectorParam(extraBody) + + val questions: ServiceParam[Any] = new ServiceParam[Any]( + this, "questions", "Native questions: System One named object or Decisions ordered array.") { + override def jsonDecode(json: String): Either[Any, String] = format.read(AIJson.parseServiceParam(json)) + override def rValue(value: Either[Any, String]): String = AIJson.rJsonValue(value) + } + + def setQuestions(value: String): this.type = setQuestions(AIJson.parse(value)) + def setQuestions(value: JsValue): this.type = setScalarParam(questions, AnyFormat.read(AIJson.fromAny(value))) + def setQuestions(value: Map[String, Any]): this.type = setQuestions(AIJson.fromAny(value)) + def setQuestions(value: Seq[Map[String, Any]]): this.type = setQuestions(AIJson.fromAny(value)) + def setQuestions(value: java.util.Map[String, Object]): this.type = setQuestions(AIJson.fromAny(value)) + def setQuestions(value: java.util.List[Object]): this.type = setQuestions(AIJson.fromAny(value)) + def getQuestions: Any = getScalarParam(questions) + def setQuestionsCol(value: String): this.type = setVectorParam(questions, value) + def getQuestionsCol: String = getVectorParam(questions) + + val stateCol: Param[String] = new Param[String](this, "stateCol", "System One state string column.") + def setStateCol(value: String): this.type = set(stateCol, value) + def getStateCol: String = $(stateCol) + + override val inputCol: Param[String] = new Param[String](this, "inputCol", "Input string column (text or JSON).") + + val inputMode: Param[String] = new Param[String]( + this, "inputMode", "Interpret native input as text or JSON.", + ParamValidators.inArray(Array("text", "json"))) + def setInputMode(value: String): this.type = set(inputMode, value) + def getInputMode: String = $(inputMode) + def setStateMode(value: String): this.type = setInputMode(value) + + val text: ServiceParam[String] = new ServiceParam[String](this, "text", "Embedding input text.") + def setText(value: String): this.type = setScalarParam(text, value) + def getText: String = getScalarParam(text) + def setTextCol(value: String): this.type = setVectorParam(text, value) + def getTextCol: String = getVectorParam(text) + + val dimensions: ServiceParam[Int] = new ServiceParam[Int]( + this, "dimensions", "Embedding dimensions; must be a positive integer.") + def setDimensions(value: Int): this.type = setScalarParam(dimensions, value) + def getDimensions: Int = getScalarParam(dimensions) + def setDimensionsCol(value: String): this.type = setVectorParam(dimensions, value) + def getDimensionsCol: String = getVectorParam(dimensions) + + val safetyIdentifier: ServiceParam[String] = new ServiceParam[String]( + this, "safetyIdentifier", "Decisions safety_identifier.") { + override val payloadName: String = "safety_identifier" + } + def setSafetyIdentifier(value: String): this.type = setScalarParam(safetyIdentifier, value) + def getSafetyIdentifier: String = getScalarParam(safetyIdentifier) + def setSafetyIdentifierCol(value: String): this.type = setVectorParam(safetyIdentifier, value) + def getSafetyIdentifierCol: String = getVectorParam(safetyIdentifier) + + val rawResponseCol: Param[String] = new Param[String]( + this, "rawResponseCol", "Complete response JSON, including unknown fields; never includes HTTP headers.") + def setRawResponseCol(value: String): this.type = set(rawResponseCol, value) + def getRawResponseCol: String = $(rawResponseCol) + + setDefault(inputMode -> "text") +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPython.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPython.scala new file mode 100644 index 00000000000..171502b3d40 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPython.scala @@ -0,0 +1,110 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.services.openai.OpenAIPromptPythonOverrides.replaceExactlyOnce +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions + +private[ai] object AIPromptPython { + private val ObjectParams = Set("responseFormat", "extraBody") + private val JsonParams = ObjectParams + "questions" + + def paramInputType(name: String): Option[String] = + if (ObjectParams(name)) Some("Union[str, Dict[str, Any]]") else None + + def nativeReader(defaultClass: String): String = { + val nativeClasses = Seq(classOf[AIPrompt], classOf[AISystemOne], classOf[OpenAIDecisions]) + val readers = nativeClasses.map { cls => + val name = cls.getName + s""""$name": "${name.replace("com.microsoft.azure.synapse.ml", "synapse.ml")}"""" + }.mkString(",\n ") + replaceExactlyOnce(defaultClass, " return from_java(java_stage, module_name)", + s""" native_readers = { + | $readers + | } + | module_name = native_readers.get(java_stage.getClass().getName(), module_name) + | return from_java(java_stage, module_name)""".stripMargin, + "AIPrompt Python reader template did not match") + } + + def initFunc(base: String): String = replaceExactlyOnce(base, + " if java_obj is None:\n kwargs = dict(kwargs)", + """ if java_obj is not None: + | self._resetUid(java_obj.uid()) + | self._transfer_params_from_java() + | if java_obj is None: + | kwargs = dict(kwargs)""".stripMargin, + "AIPrompt Python initializer template did not match") + + def paramSetter(base: String, name: String): String = { + if (!JsonParams(name)) base else { + val conversion = if (name == "responseFormat") { + """value if isinstance(value, str) else self._java_obj.responseFormat().jsonDecode('{"left":' + """ + + """self._ai_json(value, object_only=True) + '}').left().get()""" + } else s"self._ai_json(value, object_only=${if (ObjectParams(name)) "True" else "False"})" + replaceExactlyOnce(base, "value = self._service_param_value_to_java(value)", s"value = $conversion", + s"AIPrompt Python $name setter template did not match") + } + } + + def additionalMethods(base: String, className: String): String = { + val jsonParams = JsonParams.toSeq.sorted.map(name => s""""$name"""").mkString("(", ", ", ")") + base + s""" + |def _transfer_params_from_java(self, include_defaults=True): + | super($className, self)._transfer_params_from_java(include_defaults=include_defaults) + | + |def _service_param_scalar_to_python(self, name, value): + | if name not in $jsonParams: + | return super($className, self)._service_param_scalar_to_python(name, value) + | import json + | import math + | from decimal import Decimal + | def exact_number(token): + | number = float(token) + | exact = Decimal(token) + | return number if math.isfinite(number) and Decimal(str(number)) == exact else exact + | sc = SparkContext._active_spark_context + | parameter = self._java_obj.getParam(name) + | encoded = parameter.jsonEncode(sc._jvm.scala.util.Left.apply(value)) + | return json.loads(encoded, parse_float=exact_number)["left"] + |""".stripMargin + Methods + } + + val Methods: String = + """ + |def _ai_json(self, value, object_only=False): + | import json + | from decimal import Decimal + | if isinstance(value, str): + | # JVM parsing validates duplicate keys without losing the original JSON. + | return value + | if object_only and not isinstance(value, dict): + | raise TypeError("Expected a JSON object") + | active = set() + | def encode(item): + | if isinstance(item, Decimal): + | if not item.is_finite(): + | raise ValueError("Non-finite Decimal values are not JSON compliant") + | return str(item) + | if isinstance(item, (dict, list, tuple)): + | identity = id(item) + | if identity in active: + | raise ValueError("Circular reference detected") + | active.add(identity) + | try: + | if isinstance(item, dict): + | if any(not isinstance(key, str) for key in item): + | raise TypeError("JSON object keys must be strings") + | return "{" + ",".join(json.dumps(key) + ":" + encode(child) + | for key, child in item.items()) + "}" + | return "[" + ",".join(encode(child) for child in item) + "]" + | finally: + | active.remove(identity) + | return json.dumps(item, allow_nan=False) + | return encode(value) + | + |def setStateMode(self, value): + | return self.setInputMode(value) + |""".stripMargin +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIRequestAdapter.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIRequestAdapter.scala new file mode 100644 index 00000000000..85b261ca47a --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIRequestAdapter.scala @@ -0,0 +1,558 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.spark.Functions +import com.microsoft.azure.synapse.ml.param.{GlobalParams, ServiceParam} +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIChatCompletion, OpenAIEndpointUtils, + OpenAIResponses, ResponseFormatUtils} +import org.apache.spark.ml.param.Param +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.types._ +import spray.json._ + +import java.net.{URI, URISyntaxException, URLDecoder, URLEncoder} +import java.util.Locale + +private[ai] class AIMissingInputException(message: String) extends IllegalArgumentException(message) + +private[ai] object AIRequestAdapter { + val SupportedApis: Vector[String] = Vector("chat_completions", "responses", "systemone", "decisions", "embeddings") + val GenerativeApis: Set[String] = Set("chat_completions", "responses") + val NativeApis: Set[String] = Set("systemone", "decisions") + val ImplicitFabricApis: Set[String] = GenerativeApis + "embeddings" + val DefaultDatedApiVersion: String = "2025-04-01-preview" + + private val DeploymentRoute = "(?i).*/openai/deployments/([^/]+)/(?:chat/completions|embeddings)$".r + private val DeploymentBase = "(?i).*/openai/deployments/[^/]+$".r + private val AzureHostSuffixes = Seq(".azure.com", ".azure.us", ".azure.cn") + private val MaxPort = 65535 + + def endpointUri(stage: AIPrompt): URI = { + val base = stage.get(stage.customUrlRoot).getOrElse { + require(stage.isDefined(stage.url), "Configure an endpoint with setUrl or setCustomUrlRoot") + OpenAIEndpointUtils.stripTrailingSlashes(stage.getUrl) + } + val parameter = if (stage.isSet(stage.customUrlRoot)) stage.customUrlRoot else stage.url + val uri = try new URI(AIJson.validatedText(base, parameter.name)) catch { + case _: URISyntaxException => throw new IllegalArgumentException("Endpoint must be a valid absolute HTTP(S) URL") + } + require(Set("http", "https")(Option(uri.getScheme).getOrElse("").toLowerCase(Locale.ROOT)), + "Endpoint must be an absolute HTTP(S) URL") + require(normalizedHost(uri).nonEmpty, + "Endpoint must use a valid DNS host name or IP address and port; host names with underscores are not supported") + require(uri.getPort == -1 || (uri.getPort >= 1 && uri.getPort <= MaxPort), + "Endpoint port must be omitted or between 1 and 65535") + require(Option(uri.getUserInfo).isEmpty && Option(uri.getFragment).isEmpty, + "Endpoint must not contain credentials or a fragment") + if (!stage.isSet(stage.customUrlRoot) && isOpenAIHost(normalizedHost(uri)) && + routePath(uri).isEmpty && Option(uri.getRawQuery).isEmpty) { + new URI(base + "/v1") + } else uri + } + + def normalizedHost(uri: URI): String = + Option(uri.getHost).getOrElse("").stripSuffix(".").toLowerCase(Locale.ROOT) + + def isOpenAIHost(host: String): Boolean = + host == "api.openai.com" || host.endsWith(".api.openai.com") + + private def isAzureHost(host: String): Boolean = AzureHostSuffixes.exists(host.endsWith) + + private def routePath(uri: URI): String = + OpenAIEndpointUtils.stripTrailingSlashes(Option(uri.getRawPath).getOrElse("")).toLowerCase(Locale.ROOT) + + def isV1(uri: URI): Boolean = { + val path = routePath(uri) + !DeploymentRoute.pattern.matcher(path).matches() && !DeploymentBase.pattern.matcher(path).matches() && + (OpenAIEndpointUtils.isV1BaseUrl(uri.toString) || + Seq("/v1/chat/completions", "/v1/responses", "/v1/embeddings", "/v1/decisions").exists(path.endsWith)) + } + + def usesDatedApiVersion(stage: AIPrompt): Boolean = + !NativeApis(stage.getApiType) && !stage.isSet(stage.customUrlRoot) && !isV1(endpointUri(stage)) + + def defaultApiVersion(stage: AIPrompt): Option[Either[String, String]] = + if (!stage.isSet(stage.apiVersion) && stage.isDefined(stage.url) && usesDatedApiVersion(stage)) { + GlobalParams.getParam(stage.apiVersion).orElse(stage.getDefault(stage.apiVersion)) + .orElse(Some(Left(DefaultDatedApiVersion))) + } else None + + def bodyOptions(stage: AIPrompt): Seq[ServiceParam[_]] = Seq( + stage.maxTokens, stage.maxCompletionTokens, stage.temperature, stage.topP, stage.user, stage.n, + stage.stop, stage.echo, stage.cacheLevel, stage.presencePenalty, stage.frequencyPenalty, + stage.bestOf, stage.logProbs, stage.seed, stage.verbosity, stage.reasoningEffort, stage.responseFormat, + stage.store, stage.previousResponseId, stage.dimensions, stage.safetyIdentifier) + + def allowedOptions(api: String): Set[String] = api match { + case "chat_completions" => Set("maxTokens", "maxCompletionTokens", "temperature", "topP", "user", "n", + "stop", "presencePenalty", "frequencyPenalty", "seed", "verbosity", "reasoningEffort", "responseFormat") + case "responses" => Set("maxTokens", "maxCompletionTokens", "temperature", "topP", "user", + "verbosity", "reasoningEffort", "responseFormat", "store", "previousResponseId") + case "embeddings" => Set("user", "dimensions") + case "decisions" => Set("safetyIdentifier") + case "systemone" => Set.empty[String] + case _ => throw new IllegalArgumentException("Unsupported apiType") + } + + def responseFormat(value: Any, api: String): JsObject = { + val obj = value match { + case text: String => AIJson.objectString(text) + case other => AIJson.obj(AIJson.fromAny(other)) + } + obj.fields.get("type").foreach(v => + require(v.isInstanceOf[JsString], "responseFormat.type must be a string")) + obj.fields.get("json_schema").foreach(v => + require(v.isInstanceOf[JsObject], "responseFormat.json_schema must be an object")) + val schemas = Seq(obj) ++ obj.fields.get("json_schema").collect { case nested: JsObject => nested } + schemas.foreach(_.fields.get("name").foreach(v => + require(v.isInstanceOf[JsString], "responseFormat schema name must be a string"))) + val fields = AIJson.orderedMap(obj.fields.toSeq.map { case (key, entry) => + key -> AIJson.AnyFormat.read(entry) + }) + val normalized = AIJson.obj(AIJson.fromAny(ResponseFormatUtils.normalize(fields))) + val preserved = preserveFormatFields(obj, normalized) + val formatted = if (api == "chat_completions" && preserved.fields.get("type").contains(JsString("json_schema"))) { + JsObject("type" -> JsString("json_schema"), "json_schema" -> JsObject(preserved.fields - "type")) + } else preserved + AIJson.obj(AIJson.fromAny(formatted)) + } + + private def preserveFormatFields(original: JsObject, normalized: JsObject): JsObject = { + original.fields.get("type").collect { case JsString(value) => value.trim.toLowerCase(Locale.ROOT) } match { + case Some("json_schema") => + original.fields.get("json_schema") match { + case Some(nested: JsObject) => + require(original.fields.keySet.subsetOf(Set("type", "json_schema")), + "Nested responseFormat.json_schema cannot be combined with other top-level format fields") + require(!nested.fields.contains("type"), "Nested responseFormat.json_schema must not contain type") + JsObject(nested.fields.updated("type", JsString("json_schema"))) + case _ => JsObject(original.fields ++ normalized.fields) + } + case Some("text") | Some("json_object") => JsObject(original.fields ++ normalized.fields) + // Normalizing a bare envelope only adds type; avoid the legacy map's unordered intermediate. + case None if original.fields.keySet == (normalized.fields.keySet - "type") => + JsObject(original.fields.updated("type", normalized.fields("type"))) + case _ => normalized + } + } +} + +private[ai] final class AIRequestAdapter(stage: AIPrompt, schema: StructType) extends Serializable { + private val api = stage.getApiType + private val generative = AIRequestAdapter.GenerativeApis(api) + private val native = AIRequestAdapter.NativeApis(api) + private val bodyOptions = AIRequestAdapter.bodyOptions(stage) + private val allowed = AIRequestAdapter.allowedOptions(api) + + private def binding[T](param: ServiceParam[T]): Option[Either[T, String]] = + stage.get(param).orElse(stage.getDefault(param)) + + private val indices = stage.params.collect { case p: ServiceParam[_] => p }.flatMap { p => + binding(p).flatMap(_.right.toOption).map(name => + p.name -> schema.fieldIndex(AIColumnPlan.resolve(schema, name))) + }.toMap + private val inputIndex = stage.get(stage.stateCol).orElse(stage.get(stage.inputCol)) + .map(name => schema.fieldIndex(AIColumnPlan.resolve(schema, name))) + + def value(row: Row, param: ServiceParam[_]): Option[Any] = binding(param).flatMap { + case Left(v) => Option(v) + case Right(_) => Option(row.get(indices(param.name))) + } + + private def scalar(param: ServiceParam[_]): Option[Any] = + binding(param).flatMap(_.left.toOption).flatMap(Option(_)) + + private def path(param: ServiceParam[_]): List[String] = (api, param.name) match { + case ("responses", "maxTokens" | "maxCompletionTokens") => List("max_output_tokens") + case ("chat_completions", "maxTokens" | "maxCompletionTokens") => List("max_completion_tokens") + case ("responses", "responseFormat") => List("text", "format") + case ("responses", "verbosity") => List("text", "verbosity") + case ("responses", "reasoningEffort") => List("reasoning", "effort") + case _ => List(param.payloadName) + } + + private val owned = bodyOptions.filter(stage.isSet).map(path).toSet + + private def insert(base: JsObject, keys: List[String], value: JsValue): JsObject = keys match { + case key :: Nil => JsObject(base.fields.updated(key, value)) + case key :: tail => + val child = base.fields.get(key).map(AIJson.obj).getOrElse(JsObject()) + JsObject(base.fields.updated(key, insert(child, tail, value))) + case _ => base + } + + private def encodeOption(param: ServiceParam[_], value: Any): JsValue = { + val encoded = if (param.name == "responseFormat") AIRequestAdapter.responseFormat(value, api) + else AIJson.fromAny(value) + if (param.name == "dimensions") { + require(encoded match { + case JsNumber(n) => n.isWhole && n > 0 + case _ => false + }, "dimensions must be a positive integer") + } + encoded + } + + private val constantOptions = bodyOptions.filter(p => allowed(p.name)).foldLeft(JsObject()) { (body, param) => + scalar(param).map(v => insert(body, path(param), encodeOption(param, v))).getOrElse(body) + } + private val varyingOptions = bodyOptions.filter(p => allowed(p.name) && binding(p).exists(_.isRight)) + private val constantExtension = scalar(stage.extraBody).map(v => AIJson.obj(AIJson.fromAny(v))) + private val constantQuestions = binding(stage.questions).flatMap(_.left.toOption).map(AIJson.fromAny) + private val staticBody = varyingOptions.isEmpty && !binding(stage.extraBody).exists(_.isRight) + private val constantBody = AIExtraBody.merge(constantOptions, constantExtension.getOrElse(JsObject()), + owned, api, stage.isSet(stage.rawResponseCol)) + private val constantPrefix = + if (constantBody.fields.isEmpty) "{" else constantBody.compactPrint.dropRight(1) + "," + + private def mergedOptions(row: Row): JsObject = { + val perRow = varyingOptions.foldLeft(constantOptions) { (body, param) => + value(row, param).map(v => insert(body, path(param), encodeOption(param, v))).getOrElse(body) + } + val extension = constantExtension.getOrElse(value(row, stage.extraBody).map { + case json: String => AIJson.objectString(json) + case other => AIJson.obj(AIJson.fromAny(other)) + }.getOrElse(JsObject())) + AIExtraBody.merge(perRow, extension, owned, api, stage.isSet(stage.rawResponseCol)) + } + + private def validateOptions(): Unit = { + bodyOptions.filter(p => stage.isSet(p) && !allowed(p.name)).foreach { param => + throw new IllegalArgumentException( + s"${param.name} is not supported by apiType='$api'; use a supported typed option or explicit extraBody") + } + require(!(stage.isDefined(stage.maxTokens) && stage.isDefined(stage.maxCompletionTokens)), + "Cannot set both maxTokens and maxCompletionTokens") + require(stage.getConcurrency > 0, "concurrency must be positive") + require(stage.getTimeout >= 0.001 && !stage.getTimeout.isInfinity, + "timeout must be finite and at least 0.001 seconds") + stage.get(stage.concurrentTimeout).orElse(stage.getDefault(stage.concurrentTimeout)).foreach(v => + require(v >= 1e-9 && !v.isInfinity, "concurrentTimeout must be finite and at least 0.000000001 seconds")) + stage.get(stage.fileSizeLimitMB).orElse(stage.getDefault(stage.fileSizeLimitMB)).foreach(v => + require(v > 0, "fileSizeLimitMB must be positive")) + AIExtraBody.merge(constantOptions, constantExtension.getOrElse(JsObject()), + owned, api, stage.isSet(stage.rawResponseCol)) + } + + private def validateCapabilities(): Unit = { + if (!generative) { + Seq(stage.systemPrompt, stage.columnTypes, stage.fileSizeLimitMB, stage.postProcessing, + stage.postProcessingOptions, stage.messagesCol, stage.dropPrompt).filter(stage.isSet).foreach(p => + require(p == stage.messagesCol && api == "decisions", + s"${p.name} is not supported by apiType='$api'")) + } + if (stage.isSet(stage.messagesCol)) validateMessageConstruction() + validateNativeBindings() + require(api == "embeddings" || !stage.isDefined(stage.text), "text/textCol requires apiType='embeddings'") + require(api == "systemone" || !stage.isDefined(stage.stateCol), "stateCol requires apiType='systemone'") + require(!(stage.isSet(stage.apiVersion) && stage.isSet(stage.customUrlRoot)), + "apiVersion is not supported with customUrlRoot; include the version in the complete URL query instead") + if (stage.isSet(stage.apiVersion) && (native || v1)) { + throw new IllegalArgumentException(s"apiVersion is not supported for '$api' on this endpoint") + } + } + + private def validateMessageConstruction(): Unit = { + Seq(stage.systemPrompt, stage.fileSizeLimitMB, stage.dropPrompt, stage.inputMode) + .filter(stage.isSet).foreach(p => + throw new IllegalArgumentException(s"${p.name} cannot be combined with messagesCol")) + require(!stage.isSet(stage.columnTypes) || stage.getColumnTypes.isEmpty, + "Non-empty columnTypes cannot be combined with messagesCol") + } + + private def validateNativeBindings(): Unit = { + if (!native) { + require(!stage.isSet(stage.questions) && !stage.isSet(stage.stateCol) && !stage.isSet(stage.inputMode), + s"questions/stateCol/inputMode require a native API, not '$api'") + } else { + require(stage.isDefined(stage.questions), s"questions is required for apiType='$api'") + constantQuestions.foreach(NativeQuestions.validate(api, _)) + } + } + + private def validateInputs(): Unit = { + val inputs = Seq( + stage.isSet(stage.promptTemplate), stage.isSet(stage.messagesCol), stage.isSet(stage.stateCol), + stage.isSet(stage.inputCol), stage.isDefined(stage.text)).count(identity) + require(inputs == 1, "Configure exactly one of promptTemplate, messagesCol, stateCol, inputCol, or text/textCol") + Seq(stage.stateCol, stage.inputCol).filter(stage.isSet).foreach(checkStringColumn) + stage.params.collect { case p: ServiceParam[_] => p }.foreach { p => + binding(p).flatMap(_.right.toOption).foreach { name => + validateServiceColumn(p.name, schema(AIColumnPlan.resolve(schema, name)).dataType) + } + } + if (stage.isSet(stage.messagesCol)) { + schema(AIColumnPlan.resolve(schema, stage.getMessagesCol)).dataType match { + case ArrayType(message: StructType, _) => validateMessageSchema(message) + case _ => throw new IllegalArgumentException("messagesCol must contain an array of message structs") + } + } + AIColumnPlan.distinct(stage.getColumnTypes.keys.toSeq, "columnTypes keys") + stage.getColumnTypes.foreach { case (name, kind) => + require(kind == "text" || kind == "path", "columnTypes supports only 'text' and 'path'") + val field = schema(AIColumnPlan.resolve(schema, name)) + require(kind != "path" || field.dataType == StringType, "Path columns must have string type") + } + if (stage.isSet(stage.promptTemplate)) { + val input = SparkSession.active.createDataFrame(java.util.Collections.emptyList[Row](), schema) + val expression = AIPrompt.withSafeTemplate(Functions.template(AIJson.validatedString(stage.getPromptTemplate))) + val rendered = AIPrompt.bindTemplate(input, expression) + require(rendered.expr.dataType == StringType, "Prompt templates must produce strings") + } + } + + private def validateMessageSchema(dataType: DataType): Unit = dataType match { + case struct: StructType => + require(struct.fieldNames.distinct.length == struct.length, + "messagesCol must not contain duplicate protocol field names") + struct.fields.foreach(field => validateMessageSchema(field.dataType)) + case ArrayType(element, _) => validateMessageSchema(element) + case MapType(key, value, _) => + validateMessageSchema(key) + validateMessageSchema(value) + case _ => () + } + + private def jsonCompatible(dataType: DataType): Boolean = dataType match { + case StringType | BooleanType | ByteType | ShortType | IntegerType | LongType | + FloatType | DoubleType | NullType => true + case _: DecimalType => true + case ArrayType(element, _) => jsonCompatible(element) + case MapType(StringType, value, _) => jsonCompatible(value) + case _ => false + } + + private def validateServiceColumn(name: String, dataType: DataType): Unit = { + val integers = Set("maxTokens", "maxCompletionTokens", "n", "logProbs", + "cacheLevel", "bestOf", "seed", "dimensions") + val doubles = Set("temperature", "topP", "presencePenalty", "frequencyPenalty") + val headerMaps = Set("customHeaders", "telemHeaders") + val valid = if (integers(name)) dataType == IntegerType + else if (doubles(name)) dataType == DoubleType + else if (Set("store", "echo")(name)) dataType == BooleanType + else if (headerMaps(name)) dataType match { + case MapType(StringType, StringType, _) => true + case _ => false + } + else if (name == "responseFormat") { + dataType == StringType || (dataType.isInstanceOf[MapType] && jsonCompatible(dataType)) + } + else dataType == StringType + require(valid, s"${name}Col has an unsupported data type: ${dataType.simpleString}") + } + + def validate(): Unit = { + validateOptions() + validateCapabilities() + validateInputs() + require(!(stage.isSet(stage.deploymentName) && stage.isSet(stage.model)), + "Set model or deploymentName, not both; neither explicit binding may be silently ignored") + require(stage.isDefined(stage.deploymentName) || stage.isDefined(stage.model), + "Set deploymentName/model or the applicable operation-specific default") + Seq(stage.deploymentName, stage.model).foreach { p => + scalar(p).foreach(v => require(v.isInstanceOf[String] && v.asInstanceOf[String].trim.nonEmpty, + "Model/deployment must be a non-empty string")) + } + AIPrompt.validateScalarHeaders(stage, endpoint( + scalar(selectedModelParam).map(_.asInstanceOf[String]).getOrElse("validation-model"), + scalar(stage.apiVersion).map(_.toString), + validateModel = scalar(selectedModelParam).isDefined)) + } + + private def checkStringColumn(param: Param[String]): Unit = + require(schema(AIColumnPlan.resolve(schema, stage.getOrDefault(param))).dataType == StringType, + s"${param.name} must reference a string column") + + private def selectedModelParam: ServiceParam[String] = + if (stage.isSet(stage.deploymentName)) stage.deploymentName + else if (stage.isSet(stage.model)) stage.model + else if (stage.isDefined(stage.deploymentName)) stage.deploymentName else stage.model + + def validateFabricModel(): Unit = { + scalar(selectedModelParam).foreach { value => + val modelName = value.asInstanceOf[String] + if (stage.usingImplicitFabricEndpoint(endpoint(modelName, None))) { + stage.assertFabricModelStatus(modelName) + } + } + } + + private def selectedModel(row: Row): String = { + value(row, selectedModelParam) match { + case Some(text: String) if text.trim.nonEmpty => text + case _ => throw new AIMissingInputException("explicit model/deployment is missing for this row") + } + } + + private val uri = AIRequestAdapter.endpointUri(stage) + private val base = uri.toString + private val host = AIRequestAdapter.normalizedHost(uri) + private val basePath = AIRequestAdapter.routePath(uri) + private val v1 = AIRequestAdapter.isV1(uri) + + private def nativeRoute: String = api match { + case "systemone" => + require(basePath.isEmpty || basePath.endsWith("/providers/microsoft/v1/systemone"), + "systemone requires a Foundry resource endpoint or its complete System One route") + "providers/microsoft/v1/systemone" + case "decisions" => + require(!AIRequestAdapter.isAzureHost(host), + "decisions requires a qualified OpenAI endpoint, not a Foundry endpoint") + require(basePath.isEmpty || v1 || basePath.endsWith("/v1/decisions"), + "decisions requires a /v1 endpoint") + if (v1) "decisions" else "v1/decisions" + case _ => throw new IllegalArgumentException("apiType does not select a native route") + } + + private def generativeRoute(modelName: String): String = api match { + case "responses" => if (v1) "responses" else "openai/responses" + case "embeddings" => + if (v1) "embeddings" + else s"openai/deployments/${deploymentSegment(modelName)}/embeddings" + case "chat_completions" => + if (v1) "chat/completions" + else if (host.endsWith(".services.ai.azure.com") && + selectedModelParam == stage.model) "models/chat/completions" + else s"openai/deployments/${deploymentSegment(modelName)}/chat/completions" + case _ => throw new IllegalArgumentException("apiType does not select a Chat, Responses or embeddings route") + } + + private def deploymentSegment(modelName: String): String = { + validateDeploymentSegment(modelName) + URLEncoder.encode(AIJson.validatedText(modelName, selectedModelParam.name), "UTF-8").replace("+", "%20") + } + + private def validateDeploymentSegment(modelName: String): Unit = + require(modelName.trim.nonEmpty && !modelName.contains("%") && + !modelName.split("[/\\\\]", -1).exists(part => part == "." || part == ".."), + "Deployment name must be non-empty and must not contain URL dot segments or percent signs") + + private def validateCompleteDeployment(modelName: String, validateModel: Boolean): Unit = { + OpenAIEndpointUtils.stripTrailingSlashes(Option(uri.getRawPath).getOrElse("")) match { + case AIRequestAdapter.DeploymentRoute(encoded) => + val decoded = URLDecoder.decode(encoded.replace("+", "%2B"), "UTF-8") + validateDeploymentSegment(decoded) + if (validateModel) require(decoded == modelName, "Endpoint route does not match deployment") + case _ => + } + if (basePath.endsWith("/models/chat/completions")) { + require(selectedModelParam != stage.deploymentName, + "The models endpoint requires model, not deploymentName") + } + } + + private def validateDeploymentBase(): Unit = if (!stage.isSet(stage.customUrlRoot)) { + require(!AIRequestAdapter.DeploymentBase.pattern.matcher(basePath).matches(), + "Configure setUrl with a resource endpoint or a complete API route, not an incomplete deployment path") + } + + private def endpoint(modelName: String, version: Option[String], validateModel: Boolean = true): String = { + if (api == "decisions") require(!AIRequestAdapter.isAzureHost(host), + "decisions requires a qualified OpenAI endpoint, not a Foundry endpoint") + require(api != "systemone" || !AIRequestAdapter.isOpenAIHost(host), + "systemone requires a qualified Foundry endpoint, not an OpenAI endpoint") + validateDeploymentBase() + val completeRoutes = Map("/providers/microsoft/v1/systemone" -> "systemone", "/v1/decisions" -> "decisions", + "/chat/completions" -> "chat_completions", "/responses" -> "responses", "/embeddings" -> "embeddings") + completeRoutes.foreach { case (suffix, expected) => + require(!basePath.endsWith(suffix) || api == expected, "Endpoint route does not match apiType") + } + validateCompleteDeployment(modelName, validateModel) + if (stage.isSet(stage.customUrlRoot)) base else { + require(Option(uri.getQuery).isEmpty, "Use customUrlRoot for a complete URL with query parameters") + val result = if (completeRoutes.keys.exists(basePath.endsWith)) base + else base + "/" + (if (native) nativeRoute else generativeRoute(modelName)) + if (native || v1) result + else result + "?api-version=" + URLEncoder.encode(AIJson.validatedText(version.getOrElse( + AIRequestAdapter.DefaultDatedApiVersion), stage.apiVersion.name), "UTF-8") + } + } + + private def input(row: Row, preparedText: Option[String]): JsValue = { + val text = preparedText.orElse(inputIndex.flatMap(index => Option(row.getString(index)))) + .orElse(value(row, stage.text).map(_.asInstanceOf[String])) + text match { + case Some(v) if v.nonEmpty => + if (native && stage.getInputMode == "json") AIJson.parse(v) else JsString(v) + case _ => throw new AIMissingInputException("input is null or empty") + } + } + + private def validateMessageType(row: Row): Unit = { + val messageType = row.schema.fieldNames.indexOf("type") + if (api != "chat_completions" && messageType >= 0 && !row.isNullAt(messageType)) { + require(row.schema(messageType).dataType == StringType && row.getString(messageType) == "message", + "Message type must be 'message'") + } + } + + private def protocolMessages(rows: Seq[Row]): Seq[Map[String, Any]] = { + rows.foreach { row => + require(Option(row).isDefined, "Each message must be an object") + val content = row.schema.fieldNames.indexOf("content") + require(content >= 0 && content < row.length && !row.isNullAt(content), + "Each message must contain non-null content") + val supported = if (api == "chat_completions") Set("role", "content", "name") + else Set("role", "content", "type") + row.schema.fields.zipWithIndex.foreach { case (field, index) => + require(supported(field.name) || row.isNullAt(index), + s"Message field '${field.name}' is not supported by apiType='$api'") + } + validateMessageType(row) + } + if (api == "chat_completions") { + OpenAIChatCompletion.strictRequestMessages(rows).zip(rows).map { case (encoded, row) => + val name = row.schema.fieldNames.indexOf("name") + if (name < 0 || row.isNullAt(name)) encoded + else { + require(row.schema(name).dataType == StringType, "Message name must be a string") + encoded + ("name" -> row.getString(name)) + } + } + } else { + require(rows.nonEmpty, "messages must not be empty") + OpenAIResponses.requestMessages(rows) + } + } + + private def nativeEnvelope(row: Row, preparedText: Option[String], + messages: Option[Seq[Row]]): Map[String, JsValue] = { + val questions = constantQuestions.getOrElse(value(row, stage.questions).map { + case s: String => AIJson.parse(s) + case other => AIJson.fromAny(other) + }.getOrElse(throw new AIMissingInputException("questions is null"))) + if (constantQuestions.isEmpty) NativeQuestions.validate(api, questions) + val content = if (api == "decisions" && messages.isDefined) { + AIJson.fromAny(protocolMessages(messages.get)) + } else input(row, preparedText) + val validated = if (api == "decisions") NativeQuestions.prepareDecisionsInput(content) else content + Map("questions" -> questions, (if (api == "systemone") "state" else "input") -> validated) + } + + def request(row: Row, preparedText: Option[String], messages: Option[Seq[Row]]): (String, String) = { + val modelName = selectedModel(row) + val version = value(row, stage.apiVersion).map(_.toString) + if (AIRequestAdapter.usesDatedApiVersion(stage) && binding(stage.apiVersion).isDefined && version.isEmpty) { + throw new AIMissingInputException("apiVersion is null") + } + val body = if (staticBody) None else Some(mergedOptions(row)) + val envelope = api match { + case "chat_completions" => + val mapped = protocolMessages(messages.get) + Map("messages" -> AIJson.fromAny(mapped), "model" -> JsString(modelName)) + case "responses" => + val mapped = protocolMessages(messages.get) + Map("input" -> AIJson.fromAny(mapped), "model" -> JsString(modelName)) + case "embeddings" => + Map("input" -> input(row, preparedText), "model" -> JsString(modelName)) + case "systemone" | "decisions" => + nativeEnvelope(row, preparedText, messages) + ("model" -> JsString(modelName)) + case _ => throw new IllegalArgumentException("Unsupported apiType") + } + val encoded = body.map(b => JsObject(b.fields ++ envelope).compactPrint) + .getOrElse(constantPrefix + JsObject(envelope).compactPrint.drop(1)) + endpoint(modelName, version) -> encoded + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodec.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodec.scala new file mode 100644 index 00000000000..c857259e975 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodec.scala @@ -0,0 +1,365 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.services.openai.UsageUtils +import org.apache.spark.ml.linalg.SQLDataTypes.VectorType +import org.apache.spark.ml.linalg.Vectors +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{ArrayType, BooleanType, DataType, DoubleType, IntegerType, + MapType, StringType, StructType} +import spray.json.{JsArray, JsBoolean, JsNull, JsNumber, JsObject, JsString, JsValue} + +import scala.util.control.NonFatal + +/** Fixed public projections; the caller retains the original response separately for forward compatibility. */ +private[services] object AIResponseCodec { + case class Decoded(value: Any, usage: Row, responseId: String, diagnostic: Option[(String, String)]) + + private val Missing: Null = None.orNull + private val ChoiceType = new StructType() + .add("type", StringType, nullable = false) + .add("stringValue", StringType) + .add("booleanValue", BooleanType) + private val ChoiceProbabilityType = new StructType() + .add("value", ChoiceType, nullable = false) + .add("probability", DoubleType, nullable = false) + private val ScoreProbabilityType = new StructType() + .add("value", IntegerType, nullable = false) + .add("label", StringType, nullable = false) + .add("probability", DoubleType, nullable = false) + private val AnswerType = new StructType() + .add("type", StringType, nullable = false) + .add("name", StringType) + .add("probability", DoubleType) + .add("choice", ChoiceType) + .add("score", DoubleType) + .add("confidence", DoubleType) + .add("choiceProbabilities", ArrayType(ChoiceProbabilityType, containsNull = false)) + .add("scoreProbabilities", ArrayType(ScoreProbabilityType, containsNull = false)) + private val DecisionsType = new StructType() + .add("model", StringType, nullable = false) + .add("answers", ArrayType(AnswerType, containsNull = false), nullable = false) + .add("usage", UsageUtils.UsageStructType) + + // The System One guide specifies no complete answer or usage member schema. JSON strings are + // SynapseML's lossless representation, not an assertion that the provider returns string values. + private val SystemOneType = new StructType() + .add("model", StringType, nullable = false) + .add("answers", MapType(StringType, StringType, valueContainsNull = false), nullable = false) + .add("usage", StringType) + + def outputType(apiType: String): DataType = apiType match { + case "chat_completions" | "responses" => StringType + case "embeddings" => VectorType + case "systemone" => SystemOneType + case "decisions" => DecisionsType + case _ => throw new IllegalArgumentException("Unsupported AI response API type") + } + + /** Parsing/type failures have a constant, non-sensitive diagnostic; no provider text enters errors. */ + def decode(apiType: String, raw: String): Decoded = { + outputType(apiType) + try { + val parsed = if (apiType == "responses") AIJson.parseResponsesResponse(raw) else AIJson.parse(raw) + val fields = obj(parsed) + val metadata = new Metadata + val usage = metadata.read(readUsage(apiType, fields)).orNull + val id = metadata.text(fields, "id").orNull + metadata.text(fields, "model") + val providerFailed = metadata.read(providerFailure(fields)).contains(true) + val responseFailed = apiType == "responses" && metadata.text(fields, "status").contains("failed") + val failed = providerFailed || responseFailed + val decoded = try decodeValue(apiType, fields, usage, id, failed, metadata) catch { + case NonFatal(_) => result(Missing, usage, id, Some(if (failed) "provider_error" else "parse_error")) + } + // An existing provider diagnostic takes precedence over malformed optional metadata. + if (metadata.malformed && decoded.diagnostic.isEmpty) { + result(decoded.value, usage, id, Some("parse_error")) + } else decoded + } catch { + case NonFatal(_) => result(Missing, Missing, Missing, Some("parse_error")) + } + } + + private class Metadata { + var malformed: Boolean = false + + def read[A](value: => A): Option[A] = try Some(value) catch { + case NonFatal(_) => + malformed = true + None + } + + def text(fields: Map[String, JsValue], name: String): Option[String] = + read(optionalString(fields, name)).flatten + } + + private def decodeValue(apiType: String, + fields: Map[String, JsValue], + usage: Row, + id: String, + failed: Boolean, + metadata: Metadata): Decoded = { + apiType match { + case "chat_completions" => chat(fields, usage, id, failed, metadata) + case "responses" => responses(fields, usage, id, failed, metadata) + case "embeddings" => + if (failed) result(Missing, usage, id, Some("provider_error")) + else result(embedding(fields), usage, id, None) + case "systemone" => systemOne(fields, usage, id, failed) + case "decisions" => decisions(fields, usage, id, failed) + case _ => invalid() + } + } + + private def result(value: Any, usage: Row, id: String, code: Option[String]): Decoded = + Decoded(value, usage, id, code.map { code => + val message = code match { + case "provider_error" => "The provider reported a failure" + case "refusal" => "The provider refused one or more answers" + case "incomplete" => "The provider returned incomplete output" + case "no_text" => "The response contained no assistant text" + case "parse_error" => "The provider response did not match the supported JSON schema" + } + code -> message + }) + + private def invalid(): Nothing = throw new IllegalArgumentException("Invalid AI response field") + + private def obj(value: JsValue): Map[String, JsValue] = value match { + case JsObject(fields) => fields + case _ => invalid() + } + + private def array(value: JsValue): Vector[JsValue] = value match { + case JsArray(values) => values + case _ => invalid() + } + + private def string(value: JsValue): String = value match { + case JsString(s) => s + case _ => invalid() + } + + private def required(fields: Map[String, JsValue], name: String): JsValue = + fields.getOrElse(name, invalid()) + + private def optional(fields: Map[String, JsValue], name: String): Option[JsValue] = + fields.get(name).filterNot(_ == JsNull) + + private def optionalString(fields: Map[String, JsValue], name: String): Option[String] = + optional(fields, name).map(string) + + private def count(value: JsValue): Long = value match { + case JsNumber(n) if n.isWhole && n.isValidLong && n >= 0 => n.toLong + case _ => invalid() + } + + private def number(value: JsValue): Double = value match { + case JsNumber(n) if !n.toDouble.isInfinity && !n.toDouble.isNaN => n.toDouble + case _ => invalid() + } + + private def probability(value: JsValue): Double = value match { + case JsNumber(n) if n >= 0 && n <= 1 => number(value) + case _ => invalid() + } + + private def providerFailure(fields: Map[String, JsValue]): Boolean = optional(fields, "error").exists { value => + val error = obj(value) + Seq("code", "message", "type", "param").foreach(optionalString(error, _)) + true + } + + private def readUsage(apiType: String, fields: Map[String, JsValue]): Row = { + optional(fields, "usage").map { value => + val usage = obj(value) + val chatNames = apiType match { + case "chat_completions" | "embeddings" => true + case "responses" | "systemone" | "decisions" => false + case _ => invalid() + } + val input = if (chatNames) "prompt_tokens" else "input_tokens" + val output = if (chatNames) "completion_tokens" else "output_tokens" + val inputDetails = if (chatNames) "prompt_tokens_details" else "input_tokens_details" + val outputDetails = if (chatNames) "completion_tokens_details" else "output_tokens_details" + def tokens(key: String): java.lang.Long = usage.get(key).map(v => Long.box(count(v))).orNull + def details(key: String, known: Set[String]): Map[String, Long] = + optional(usage, key).map { value => + obj(value).collect { case (name, tokenCount) if known(name) && tokenCount != JsNull => + name -> count(tokenCount) + } + }.getOrElse(Map.empty) + Row(tokens(input), tokens(output), tokens("total_tokens"), + details(inputDetails, Set("cached_tokens", "cache_write_tokens", "audio_tokens")), + details(outputDetails, Set("reasoning_tokens", "audio_tokens", + "accepted_prediction_tokens", "rejected_prediction_tokens"))) + }.orNull + } + + private def chat(fields: Map[String, JsValue], + usage: Row, + id: String, + failed: Boolean, + metadata: Metadata): Decoded = { + val choices = optional(fields, "choices").map(array).getOrElse { + if (failed) Vector.empty else invalid() + } + // Deliberately project only the first choice, while checking known fields on every choice. + val parsed = choices.map { choice => + if (failed) metadata.read(chatChoice(choice, metadata)) else Some(chatChoice(choice, metadata)) + } + val first = parsed.headOption.flatten + val text = first.flatMap(_._1) + val code = if (failed) Some("provider_error") + else if (first.exists(_._2)) Some("refusal") + else if (first.exists(_._3)) Some("incomplete") + else if (text.isEmpty) Some("no_text") + else None + result(text.orNull, usage, id, code) + } + + private def chatChoice(choice: JsValue, metadata: Metadata): (Option[String], Boolean, Boolean) = { + val fields = obj(choice) + count(required(fields, "index")) + val message = obj(required(fields, "message")) + if (string(required(message, "role")) != "assistant") invalid() + val content = optionalString(message, "content") + val refusal = metadata.text(message, "refusal").exists(_.nonEmpty) + val finish = metadata.text(fields, "finish_reason") + (content, refusal || finish.contains("content_filter"), finish.contains("length")) + } + + private def incompleteResponse(status: Option[String], details: Option[JsValue]): Boolean = + details.isDefined || status.exists(s => Set("incomplete", "in_progress", "queued", "cancelled")(s)) + + private def responses(fields: Map[String, JsValue], + usage: Row, + id: String, + failed: Boolean, + metadata: Metadata): Decoded = { + val status = metadata.text(fields, "status") + val incomplete = optional(fields, "incomplete_details") + incomplete.foreach { value => metadata.read(optionalString(obj(value), "reason")) } + val output = optional(fields, "output").map(array).getOrElse { + if (failed) Vector.empty else invalid() + } + val messages = output.flatMap { item => + if (failed) metadata.read(responseMessage(item, metadata, failed)).toSeq.flatten + else responseMessage(item, metadata, failed).toSeq + } + val parts = messages.flatMap(_._1) + val texts = parts.flatMap(_._1) + val code = if (failed) Some("provider_error") + else if (parts.exists(_._2)) Some("refusal") + else if (incompleteResponse(status, incomplete) || messages.exists(_._2)) Some("incomplete") + else if (texts.isEmpty) Some("no_text") + else None + result(if (texts.isEmpty) Missing else texts.mkString, usage, id, code) + } + + private def responseMessage(item: JsValue, + metadata: Metadata, + failed: Boolean): Option[(Vector[(Option[String], Boolean)], Boolean)] = { + val fields = obj(item) + if (string(required(fields, "type")) == "message") { + val role = string(required(fields, "role")) + val parts = array(required(fields, "content")).flatMap { part => + if (failed) metadata.read(responsePart(part)) else Some(responsePart(part)) + } + val incomplete = metadata.text(fields, "status").exists(Set("incomplete", "in_progress")) + if (role == "assistant") Some((parts, incomplete)) else None + } else None + } + + private def responsePart(part: JsValue): (Option[String], Boolean) = { + val fields = obj(part) + string(required(fields, "type")) match { + case "output_text" => (Some(string(required(fields, "text"))), false) + case "refusal" => + string(required(fields, "refusal")) + (None, true) + case _ => (None, false) + } + } + + private def embedding(fields: Map[String, JsValue]): org.apache.spark.ml.linalg.Vector = { + optional(fields, "object").foreach { value => if (string(value) != "list") invalid() } + val data = array(required(fields, "data")) + if (data.size != 1) invalid() + val item = obj(data.head) + if (count(required(item, "index")) != 0) invalid() + optional(item, "object").foreach { value => if (string(value) != "embedding") invalid() } + val values = array(required(item, "embedding")) + if (values.isEmpty) invalid() + Vectors.dense(values.map(number).toArray) + } + + private def tagged(value: JsValue): Row = value match { + case JsString(s) => Row("string", s, Missing) + case JsBoolean(b) => Row("boolean", Missing, b) + case _ => invalid() + } + + private def systemOne(fields: Map[String, JsValue], usage: Row, id: String, failed: Boolean): Decoded = { + if (failed) result(Missing, usage, id, Some("provider_error")) + else { + val answers = obj(required(fields, "answers")).map { case (key, value) => key -> value.compactPrint } + if (answers.isEmpty) invalid() + val nativeUsage = optional(fields, "usage").map(_.compactPrint).orNull + result(Row(string(required(fields, "model")), answers, nativeUsage), usage, id, None) + } + } + + private def decisions(fields: Map[String, JsValue], usage: Row, id: String, failed: Boolean): Decoded = { + if (failed) result(Missing, usage, id, Some("provider_error")) + else { + val answers = array(required(fields, "answers")).map(answer) + if (answers.isEmpty) invalid() + val names = answers.filterNot(_.isNullAt(1)).map(_.getString(1)) + if (names.distinct.size != names.size) invalid() + val value = Row(string(required(fields, "model")), answers, usage) + result(value, usage, id, if (answers.exists(_.getString(0) == "refusal")) Some("refusal") else None) + } + } + + private def answer(value: JsValue): Row = { + val fields = obj(value) + val kind = string(required(fields, "type")) + val name = optionalString(fields, "name").orNull + kind match { + case "predicate" => + Row(kind, name, probability(required(fields, "probability")), Missing, Missing, Missing, Missing, Missing) + case "refusal" => Row(kind, name, Missing, Missing, Missing, Missing, Missing, Missing) + case "choice" => choiceAnswer(fields, name) + case "score" => scoreAnswer(fields, name) + case _ => invalid() + } + } + + private def choiceAnswer(fields: Map[String, JsValue], name: String): Row = { + val choice = tagged(required(fields, "choice")) + val probabilities = array(required(fields, "probabilities")).map { value => + val fields = obj(value) + Row(tagged(required(fields, "value")), probability(required(fields, "probability"))) + } + val values = probabilities.map(_.getStruct(0)) + if (values.size < 2 || values.distinct.size != values.size || !values.contains(choice)) invalid() + Row("choice", name, Missing, choice, Missing, probability(required(fields, "confidence")), probabilities, Missing) + } + + private def scoreAnswer(fields: Map[String, JsValue], name: String): Row = { + val probabilities = array(required(fields, "probabilities")).zipWithIndex.map { case (value, index) => + val fields = obj(value) + if (count(required(fields, "value")) != index) invalid() + Row(index, string(required(fields, "label")), probability(required(fields, "probability"))) + } + val score = required(fields, "score") match { + case value @ JsNumber(n) if n >= 0 && n <= probabilities.size - 1 => number(value) + case _ => invalid() + } + Row("score", name, Missing, Missing, score, probability(required(fields, "confidence")), Missing, probabilities) + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOne.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOne.scala new file mode 100644 index 00000000000..4aa5fe6487c --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOne.scala @@ -0,0 +1,19 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.ml.ExplicitComplexParamsReadable +import org.apache.spark.ml.util.Identifiable + +object AISystemOne extends ExplicitComplexParamsReadable[AISystemOne] { + override protected def expectedClass: Class[AISystemOne] = classOf[AISystemOne] +} + +/** Native Microsoft System One requests; no chat emulation or default model. */ +class AISystemOne(override val uid: String) extends AIPrompt(uid) { + def this() = this(Identifiable.randomUID("AISystemOne")) + + override protected def supportedApiTypes: Array[String] = Array("systemone") + override protected def defaultApiType: String = "systemone" +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodec.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodec.scala new file mode 100644 index 00000000000..0d74c3d3e58 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodec.scala @@ -0,0 +1,152 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.injections.UDFUtils +import org.apache.spark.ml.linalg.{DenseVector, SQLDataTypes, SparseVector, Vector, Vectors} +import org.apache.spark.sql.{Column, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{GetStructField, KnownNotNull, KnownNullable} +import org.apache.spark.sql.functions.{lit, struct, transform, transform_keys, transform_values, when} +import org.apache.spark.sql.types._ + +/** Preserve Spark values and schema while input rows traverse inference encoders. */ +private[ai] object AIValueCodec { + private val VectorStorage = StructType(Seq( + StructField("sparse", BooleanType), + StructField("size", IntegerType), + StructField("indices", ArrayType(IntegerType, containsNull = true)), + StructField("values", ArrayType(DoubleType, containsNull = true)))) + def nullable(dataType: DataType): DataType = dataType match { + case s: StructType => StructType(s.fields.map(f => f.copy(dataType = nullable(f.dataType), nullable = true))) + case ArrayType(t, _) => ArrayType(nullable(t), containsNull = true) + case MapType(k, v, _) => MapType(nullable(k), nullable(v), valueContainsNull = true) + case other => other + } + + def validateMapKeys(dataType: DataType, inKey: Boolean = false): Unit = dataType match { + case s: StructType => s.fields.foreach(field => validateMapKeys(field.dataType, inKey)) + case ArrayType(t, _) => validateMapKeys(t, inKey) + case MapType(k, v, _) => validateMapKeys(k, inKey = true); validateMapKeys(v, inKey) + case t: UserDefinedType[_] => + require(!inKey || t == SQLDataTypes.VectorType, + "AIPrompt supports only Spark ML vectors as UDT map keys; other UDT keys can lose SQL identity.") + case _ => + } + + def encodeMapKeys(frame: DataFrame): DataFrame = convertMapKeyColumns(frame, frame.schema, encode = true) + + def decodeMapKeys(frame: DataFrame, schema: StructType): DataFrame = + convertMapKeyColumns(frame, schema, encode = false) + + def reconcileSchema(frame: DataFrame, schema: StructType): DataFrame = { + val columns = schema.fields.map { field => + val source = frame.schema(field.name) + require(DataType.equalsIgnoreNullability(source.dataType, field.dataType), + "AIPrompt output types do not match the planned schema") + val column = AIColumnPlan.column(field.name) + // Struct casts preserve ordinal fields; Dataset.to resolves nested names and can introduce ambiguity. + val converted = if (source.dataType == field.dataType) column else column.cast(field.dataType) + val widened = if (field.nullable && !source.nullable) { + new Column(KnownNullable(converted.expr)) + } else converted + widened.as(field.name, field.metadata) + } + val result = frame.select(columns: _*) + require(result.schema == schema, "AIPrompt output schema does not match the planned schema") + result + } + + private def containsVectorMapKey(dataType: DataType): Boolean = dataType match { + case s: StructType => s.fields.exists(field => containsVectorMapKey(field.dataType)) + case ArrayType(t, _) => containsVectorMapKey(t) + case MapType(k, v, _) => containsVector(k) || containsVectorMapKey(v) + case _ => false + } + + private def convertMapKeyColumns(frame: DataFrame, schema: StructType, encode: Boolean): DataFrame = { + val converted = schema.fields.filter(field => containsVectorMapKey(field.dataType)).map(_.name).toSet + if (converted.isEmpty) frame else frame.select(schema.fields.map { field => + val column = AIColumnPlan.column(field.name) + if (converted(field.name)) { + convertColumn(column, field.dataType, field.nullable, encode).as(field.name, field.metadata) + } else column + }: _*) + } + + private def containsVector(dataType: DataType): Boolean = dataType match { + case t if t == SQLDataTypes.VectorType => true + case s: StructType => s.fields.exists(field => containsVector(field.dataType)) + case ArrayType(t, _) => containsVector(t) + case MapType(k, v, _) => containsVector(k) || containsVector(v) + case _ => false + } + + private def convertVectorColumn(column: Column, nullable: Boolean, encode: Boolean): Column = { + val conversion = if (encode) { + UDFUtils.oldUdf((value: Vector) => Option(value).map(encodeVector).orNull, VectorStorage) + } else { + UDFUtils.oldUdf((value: Row) => Option(value).map(decodeVector).orNull, SQLDataTypes.VectorType) + } + (if (nullable) conversion else conversion.asNonNullable())(column) + } + + private def convertMapColumn(column: Column, dataType: MapType, encode: Boolean): Column = { + // Catalyst preserves SQL key identity; JVM vector equality can collapse dense and sparse keys. + val values = if (containsVector(dataType.valueType)) { + transform_values(column, (_, value) => + convertColumn(value, dataType.valueType, dataType.valueContainsNull, encode)) + } else column + if (containsVector(dataType.keyType)) { + transform_keys(values, (key, _) => convertColumn(key, dataType.keyType, nullable = false, encode = encode)) + } else values + } + + private def convertStructColumn(column: Column, schema: StructType, nullable: Boolean, encode: Boolean): Column = { + val fields = schema.fields.zipWithIndex.map { case (field, index) => + val extracted = GetStructField(column.expr, index, Some(field.name)) + // The parent-null guard makes declared non-null children safe to retain as non-null. + val child = new Column(if (field.nullable) extracted else KnownNotNull(extracted)) + convertColumn(child, field.dataType, field.nullable, encode).as(field.name, field.metadata) + } + val converted = struct(fields: _*) + if (nullable) when(column.isNotNull, converted) else converted + } + + private def convertColumn(column: Column, dataType: DataType, nullable: Boolean, encode: Boolean): Column = { + if (!containsVector(dataType)) column else dataType match { + case t if t == SQLDataTypes.VectorType => convertVectorColumn(column, nullable, encode) + case s: StructType => convertStructColumn(column, s, nullable, encode) + case ArrayType(t, containsNull) => + transform(column, value => convertColumn(value, t, containsNull, encode)) + case map: MapType => convertMapColumn(column, map, encode) + case _ => column + } + } + + private def encodeVector(vector: Vector): Row = vector match { + case sparse: SparseVector => Row(true, sparse.size, sparse.indices.toSeq, sparse.values.toSeq) + case dense: DenseVector => Row(false, dense.size, Seq.empty[Int], dense.values.toSeq) + } + + private def decodeVector(row: Row): Vector = { + require(row.toSeq.forall(value => Option(value).isDefined), "Invalid encoded AIPrompt vector fields.") + val size = row.getInt(1) + val rawIndices = row.getSeq[Any](2) + val rawValues = row.getSeq[Any](3) + require((rawIndices ++ rawValues).forall(value => Option(value).isDefined), + "Invalid encoded AIPrompt vector elements.") + val indices = rawIndices.map(_.asInstanceOf[Int]).toArray + val values = rawValues.map(_.asInstanceOf[Double]).toArray + require(size >= 0, "Invalid encoded AIPrompt vector size.") + if (row.getBoolean(0)) { + require(indices.length == values.length && indices.forall(i => i >= 0 && i < size) && + indices.sliding(2).forall(pair => pair.length < 2 || pair.head < pair.last), + "Invalid encoded AIPrompt sparse vector.") + Vectors.sparse(size, indices, values) + } else { + require(indices.isEmpty && values.length == size, "Invalid encoded AIPrompt dense vector.") + Vectors.dense(values) + } + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestions.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestions.scala new file mode 100644 index 00000000000..838c72dcef7 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestions.scala @@ -0,0 +1,196 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.services.openai.OpenAIAttachmentUtils +import spray.json.{JsArray, JsBoolean, JsNull, JsObject, JsString, JsValue} + +/** Shared validation for the native leaves and the facade, without rewriting provider JSON. */ +private[services] object NativeQuestions { + private val MaxChoices = 255 + private val MaxImages = 128 + private val InlineImagePrefix = "data:image/" + private val UrlEscape = "%[0-9a-fA-F]{2}".r + private val UrlToken = "[a-zA-Z0-9!$&'*+._~-]+".r + private def check(condition: Boolean, message: String): Unit = + require(condition, message) + + private def obj(value: JsValue): Map[String, JsValue] = value match { + case JsObject(fields) => fields + case _ => throw new IllegalArgumentException("Native question or input member must be an object") + } + + private def array(value: JsValue): Vector[JsValue] = value match { + case JsArray(values) if values.nonEmpty => values + case _ => throw new IllegalArgumentException("Native question or input array must be nonempty") + } + + private def text(value: JsValue): String = value match { + case JsString(s) if s.trim.nonEmpty => s + case _ => throw new IllegalArgumentException("Native question or input text must be nonempty") + } + + private def required(fields: Map[String, JsValue], name: String): JsValue = + fields.getOrElse(name, throw new IllegalArgumentException(s"Native request requires $name")) + + private def optionalText(fields: Map[String, JsValue], name: String): Option[String] = + fields.get(name).filterNot(_ == JsNull).map { + case JsString(value) => value + case _ => throw new IllegalArgumentException("Optional native text fields must be strings") + } + + private def choiceValue(value: JsValue): Unit = value match { + case _: JsString | _: JsBoolean => + case _ => throw new IllegalArgumentException("Decision choice values must be strings or booleans") + } + + def validate(apiType: String, questions: JsValue): Unit = apiType match { + case "systemone" => systemOne(questions) + case "decisions" => decisions(questions) + case _ => throw new IllegalArgumentException("Native questions require systemone or decisions") + } + + private def systemOne(questions: JsValue): Unit = { + val named = obj(questions) + check(named.nonEmpty, "System One questions must be nonempty") + named.foreach { case (name, question) => + check(name.trim.nonEmpty, "System One question names must be nonempty") + val fields = obj(question) + text(required(fields, "instructions")) + text(required(fields, "type")) match { + case "noul" => + fields.get("criteria").foreach { value => + val criteria = obj(value) + check(criteria.keySet == Set("true", "false"), + "System One noul criteria must describe true and false") + criteria.values.foreach(text) + } + case "choice" => + val criteria = obj(required(fields, "criteria")) + check(criteria.size >= 2, "System One choice requires at least two criteria") + criteria.foreach { case (label, description) => + check(label.trim.nonEmpty, "System One choice labels must be nonempty") + text(description) + } + case "score" => + array(required(fields, "criteria")).foreach(text) + case _ => throw new IllegalArgumentException("Unsupported System One question type") + } + } + } + + private def decisions(questions: JsValue): Unit = { + val names = array(questions).flatMap { question => + val fields = obj(question) + text(required(fields, "instructions")) + text(required(fields, "type")) match { + case "predicate" => + case "choice" => + val choices = array(required(fields, "choices")) + check(choices.size >= 2 && choices.size <= MaxChoices, "Decisions requires 2 to 255 choices") + val values = choices.map { choice => + val option = obj(choice) + val value = required(option, "value") + choiceValue(value) + optionalText(option, "description") + value + } + check(values.distinct.size == values.size, "Decision choice values must be unique") + case "score" => + array(required(fields, "levels")).foreach { level => + val fields = obj(level) + optionalText(fields, "description") + text(required(fields, "label")) + } + case _ => throw new IllegalArgumentException("Unsupported Decisions question type") + } + optionalText(fields, "name") + } + check(names.distinct.size == names.size, "Decision question names must be unique") + } + + /** The guide and reference disagree on hosted images; accept only their inline-image common subset. */ + def validateDecisionsInput(input: JsValue): Unit = input match { + case _: JsString => text(input) + case _: JsArray => + val images = array(input).map { message => + val fields = obj(message) + check(required(fields, "role") == JsString("user"), "Decisions accepts only user messages") + fields.get("type").filterNot(_ == JsNull).foreach { value => + check(value == JsString("message"), "Decisions message type must be message") + } + required(fields, "content") match { + case value: JsString => text(value); 0 + case value: JsArray => array(value).map(validatePart).sum + case _ => throw new IllegalArgumentException("Unsupported Decisions message content") + } + }.sum + check(images <= MaxImages, "Decisions accepts at most 128 images") + case _ => throw new IllegalArgumentException("Decisions input must be text or user messages") + } + + def prepareDecisionsInput(input: JsValue): JsValue = { + validateDecisionsInput(input) + input match { + case JsArray(messages) => JsArray(messages.map { message => + val fields = omitNulls(obj(message), Seq("type")) + val content = fields("content") match { + case JsArray(parts) => JsArray(parts.map { part => + val fields = obj(part) + if (fields("type") == JsString("input_image")) JsObject(omitNulls(fields, Seq("file_id", "detail"))) + else part + }) + case text => text + } + JsObject(fields.updated("content", content)) + }) + case text => text + } + } + + private def omitNulls(fields: Map[String, JsValue], names: Seq[String]): Map[String, JsValue] = + names.foldLeft(fields) { case (result, name) => + if (result.get(name).contains(JsNull)) result - name else result + } + + private def inlineImageHeader(header: String): Boolean = { + // RFC 2397 uses URL escapes, rather than quoted strings, for MIME parameter values. + def token(value: String): Boolean = UrlToken.pattern.matcher(UrlEscape.replaceAllIn(value, "x")).matches() + val parts = header.split(";", -1) + parts.length >= 2 && parts.head.regionMatches(true, 0, InlineImagePrefix, 0, InlineImagePrefix.length) && + token(parts.head.substring(InlineImagePrefix.length)) && parts.last.equalsIgnoreCase("base64") && + parts.slice(1, parts.length - 1).forall { parameter => + val equals = parameter.indexOf('=') + equals > 0 && token(parameter.substring(0, equals)) && token(parameter.substring(equals + 1)) + } + } + + private def validatePart(part: JsValue): Int = { + val fields = obj(part) + text(required(fields, "type")) match { + case "input_text" => + text(required(fields, "text")) + 0 + case "input_image" => + check(fields.get("file_id").forall(_ == JsNull), "Decisions file_id images are unsupported") + fields.get("detail").filterNot(_ == JsNull).foreach { value => + check(Set("low", "high", "auto", "original").contains(text(value)), "Unsupported image detail") + } + val url = text(required(fields, "image_url")) + val comma = url.indexOf(',') + check(comma > 0 && inlineImageHeader(url.substring(0, comma)), + "Decisions images must be inline base64 image data URLs") + val encoded = url.substring(comma + 1) + try { + OpenAIAttachmentUtils.validateBase64Content(encoded) + } catch { + case _: IllegalArgumentException => + throw new IllegalArgumentException("Decisions image data must be valid base64") + } + check(encoded.nonEmpty, "Decisions image data must be nonempty") + 1 + case _ => throw new IllegalArgumentException("Unsupported Decisions input part") + } + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/HasOpenAIResponseSchema.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/HasOpenAIResponseSchema.scala index 2d236862c8f..7f2b10c54ee 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/HasOpenAIResponseSchema.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/HasOpenAIResponseSchema.scala @@ -10,6 +10,13 @@ trait HasOpenAIResponseSchema extends Wrappable { def setResponseFormat(value: Map[String, Any]): this.type + protected def validateResponseSchemaValues(schema: Map[String, Any]): Unit = { + AnyJsonFormat.anyFormat.write(schema) + } + + protected def pyResponseSchemaToJava: String = + "jvm.com.microsoft.azure.synapse.ml.param.ServiceParam.toMap(_convert(schema))" + /** Request strict structured output from an inner JSON Schema, without its response-format envelope. */ def setResponseSchema(schema: Map[String, Any]): this.type = setResponseSchema(schema, "response_schema", strict = true) @@ -31,7 +38,7 @@ trait HasOpenAIResponseSchema extends Wrappable { require(name != null && name.matches("[A-Za-z0-9_-]{1,64}"), "schema name must contain 1 to 64 letters, digits, underscores, or hyphens") // Reject unsupported nested values before changing the stage's parameters. - AnyJsonFormat.anyFormat.write(schema) + validateResponseSchemaValues(schema) setResponseFormat(Map( "type" -> "json_schema", "name" -> name, @@ -77,7 +84,7 @@ trait HasOpenAIResponseSchema extends Wrappable { | return result | return value | - | java_schema = jvm.com.microsoft.azure.synapse.ml.param.ServiceParam.toMap(_convert(schema)) + | java_schema = $pyResponseSchemaToJava | self._java_obj = self._java_obj.setResponseSchema(java_schema, name, strict) | self._paramMap.pop(self.responseFormat, None) | return self diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala index 1b1bb55e902..d2e8729ebd3 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala @@ -38,8 +38,8 @@ trait HasMessagesInput extends Params { case object OpenAIDeploymentNameKey extends GlobalKey[Either[String, String]] case object OpenAIEmbeddingDeploymentNameKey extends GlobalKey[Either[String, String]] -private[openai] object OpenAIEndpointUtils { - private def stripTrailingSlashes(value: String): String = value.replaceAll("/+$", "") +private[services] object OpenAIEndpointUtils { + def stripTrailingSlashes(value: String): String = value.replaceAll("/+$", "") private def withoutQueryOrFragment(value: String): String = { val stopAt = Seq(value.indexOf("?"), value.indexOf("#")).filter(_ >= 0) match { @@ -555,9 +555,17 @@ abstract class OpenAIServicesBase(override val uid: String) extends CognitiveSer } } -private[openai] object OpenAIFabricHeaders { +private[services] object OpenAIFabricModelStatus extends OpenAIFabricSetting + +private[services] object OpenAIFabricHeaders { lazy val Values: Map[String, String] = build(PlatformDetails.FabricRuntime) + def merge(headers: Option[Map[String, String]]): Map[String, String] = { + val remaining = headers.map(ServiceAuthHeaders.sanitizeHeaderMap).getOrElse(Map.empty) + .filterNot { case (name, _) => Values.keys.exists(_.equalsIgnoreCase(name)) } + remaining ++ Values + } + private[openai] def build(runtime: String): Map[String, String] = Map( "X-Taxonomy-TrafficType" -> "Background", "X-Llm-Service-Tier" -> "flex", @@ -591,10 +599,7 @@ trait HasOpenAIFabricHeaders extends HasCognitiveServiceInput { val headers = super.getCustomHeaders(row) if (usingImplicitFabricEndpoint) { // SynapseML owns the complete workload classification on its implicit Fabric endpoint. - val fabricHeaders = OpenAIFabricHeaders.Values - val remainingHeaders = headers.map(ServiceAuthHeaders.sanitizeHeaderMap).getOrElse(Map.empty) - .filterNot { case (name, _) => fabricHeaders.keys.exists(_.equalsIgnoreCase(name)) } - Some(if (remainingHeaders.isEmpty) fabricHeaders else remainingHeaders ++ fabricHeaders) + Some(OpenAIFabricHeaders.merge(headers)) } else { headers } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala index 95705b8032b..3caf39a3737 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala @@ -7,17 +7,19 @@ import org.apache.commons.io.IOUtils import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.{Path => HPath} import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.injections.TaskContextUtils -import java.io.{ByteArrayInputStream, InterruptedIOException} +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InterruptedIOException} import java.net.{URI, URLConnection} import java.nio.channels.ClosedByInterruptException import java.nio.charset.StandardCharsets import java.util.concurrent.CancellationException import java.util.{Base64, Collections, IdentityHashMap, Locale} +import scala.annotation.tailrec import scala.util.{Try, Using} import scala.util.control.NonFatal -private[openai] object OpenAIAttachmentUtils { +private[services] object OpenAIAttachmentUtils { sealed class AttachmentReadException(fileName: String) extends IllegalArgumentException(s"Unable to read attachment '$fileName'") @@ -25,6 +27,8 @@ private[openai] object OpenAIAttachmentUtils { /** Storage CancellationException, not task interruption; OpenAIPrompt keeps this sanitized error row-local. */ final class CancelledAttachmentReadException(fileName: String) extends AttachmentReadException(fileName) + private class AttachmentSizeException(message: String) extends IllegalArgumentException(message) + private def interruptedRead(): Nothing = { Thread.currentThread().interrupt() throw new InterruptedException("Attachment read interrupted") @@ -32,7 +36,10 @@ private[openai] object OpenAIAttachmentUtils { private def taskCancelled: Boolean = Option(TaskContext.get()).exists(_.isInterrupted()) - private def taskKilled(): Nothing = throw new TaskKilledException("Attachment read canceled") + private def taskKilled(): Nothing = { + throw new TaskKilledException( + TaskContextUtils.getKillReason(TaskContext.get()).getOrElse("Attachment read canceled")) + } private def checkCancellation(): Unit = { if (Thread.currentThread().isInterrupted) interruptedRead() @@ -97,6 +104,7 @@ private[openai] object OpenAIAttachmentUtils { case NonFatal(_) if Thread.currentThread().isInterrupted => interruptedRead() case NonFatal(_) if taskCancelled => taskKilled() case _: CancellationException => throw new CancelledAttachmentReadException(fileName) + case error: AttachmentSizeException => throw error case NonFatal(_) => throw new AttachmentReadException(fileName) } checkCancellation() @@ -111,6 +119,7 @@ private[openai] object OpenAIAttachmentUtils { originalUrl: String) private val MimeSniffPrefixLength = 64 + private val AttachmentBufferSize = 8192 private val Utf8Bom = Base64.getDecoder.decode("77u/") private val XmlDeclarationPrefix = " limitBytes) { val fileSizeMB = fileSizeBytes / (1024.0 * 1024.0) - throw new IllegalArgumentException( + throw new AttachmentSizeException( f"$fileDescription size $fileSizeMB%.2f MB exceeds limit $limit%.2f MB") } } @@ -184,7 +193,7 @@ private[openai] object OpenAIAttachmentUtils { char >= '0' && char <= '9' || char == '+' || char == '/' - private def validateBase64Content(encodedContent: String): Unit = { + def validateBase64Content(encodedContent: String): Unit = { val padding = encodedContent.reverseIterator.takeWhile(_ == '=').length val dataLength = encodedContent.length - padding val remainder = dataLength % 4 @@ -305,6 +314,63 @@ private[openai] object OpenAIAttachmentUtils { imageExtensions: Set[String], audioExtensions: Set[String], textExtensions: Set[String] + ): (String, Array[Byte], String, String) = + prepareFile(filePathStr, sizeLimitMB, imageExtensions, audioExtensions, textExtensions, bounded = false) + + // Using retains both read and cleanup failures for interruption classification. + private def readFile(path: HPath, conf: Configuration): Array[Byte] = + Using.resource(path.getFileSystem(conf).open(path))(IOUtils.toByteArray) + + private def readBoundedFile(path: HPath, limit: Double, conf: Configuration): Array[Byte] = { + val limitBytes = (limit * 1024 * 1024).toLong + if (limitBytes == Long.MaxValue) readFile(path, conf) + else Using.resource(path.getFileSystem(conf).open(path)) { input => + val output = new ByteArrayOutputStream() + val buffer = new Array[Byte](AttachmentBufferSize) + @tailrec + def fill(total: Long): Unit = { + val length = math.min(buffer.length.toLong, limitBytes - total + 1).toInt + val read = input.read(buffer, 0, length) + val count = if (read == 0) { + val next = input.read() + if (next < 0) -1 else { + buffer(0) = next.toByte + 1 + } + } else read + if (count >= 0) { + val nextTotal = total + count + if (nextTotal > limitBytes) { + throw new AttachmentSizeException(s"Attachment exceeds the $limitBytes-byte limit") + } + output.write(buffer, 0, count) + fill(nextTotal) + } + } + fill(0L) + output.toByteArray + } + } + + def prepareFile( + filePathStr: String, + sizeLimitMB: Option[Double], + imageExtensions: Set[String], + audioExtensions: Set[String], + textExtensions: Set[String], + bounded: Boolean + ): (String, Array[Byte], String, String) = + prepareFile(filePathStr, sizeLimitMB, imageExtensions, audioExtensions, textExtensions, + bounded, new Configuration()) + + def prepareFile( + filePathStr: String, + sizeLimitMB: Option[Double], + imageExtensions: Set[String], + audioExtensions: Set[String], + textExtensions: Set[String], + bounded: Boolean, + conf: Configuration ): (String, Array[Byte], String, String) = { val isDataUrl = filePathStr.regionMatches(true, 0, "data:", 0, "data:".length) val (fileName, fileBytes, declaredMimeType) = @@ -314,8 +380,8 @@ private[openai] object OpenAIAttachmentUtils { val fileName = attachmentFilename(filePathStr) val fileBytes = readAttachment(fileName) { val path = new HPath(filePathStr) - // Using retains both read and cleanup failures for interruption classification. - Using.resource(path.getFileSystem(new Configuration()).open(path))(IOUtils.toByteArray) + if (bounded && sizeLimitMB.isDefined) readBoundedFile(path, sizeLimitMB.get, conf) + else readFile(path, conf) } (fileName, fileBytes, None) } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala index 20be6fdfa3d..0f8c88c92b7 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala @@ -22,32 +22,24 @@ import scala.language.existentials trait HasOpenAITextParamsExtended extends HasOpenAITextParams with HasOpenAIResponseSchema { - val responseFormat: ServiceParam[Map[String, Any]] = new ServiceParam[Map[String, Any]]( - this, - "responseFormat", - "Response format for the completion. One of 'text', 'json_object', or 'json_schema'.", - isRequired = false) { - override val payloadName: String = "response_format" - } + // The argument keeps this factory out of Spark's zero-argument Param discovery. + /** Called during trait initialization; overrides must not read uninitialized subclass fields. */ + protected def createResponseFormatParam(name: String): ServiceParam[Map[String, Any]] = + new ServiceParam[Map[String, Any]]( + this, + name, + "Response format for the completion. One of 'text', 'json_object', or 'json_schema'.", + isRequired = false) { + override val payloadName: String = "response_format" + } + + // Initialize before sharedTextParams, including when a subclass supplies a specialized Param. + val responseFormat: ServiceParam[Map[String, Any]] = createResponseFormatParam("responseFormat") def getResponseFormat: Map[String, Any] = getScalarParam(responseFormat) def setResponseFormat(value: Map[String, Any]): this.type = { - val normalized = ResponseFormatUtils.normalize(value) - val payload = - if (normalized.get("type").exists(_.toString.equalsIgnoreCase("json_schema"))) { - Map( - "type" -> "json_schema", - "json_schema" -> { - val base = Map( - "name" -> normalized("name"), - "schema" -> normalized("schema") - ) - normalized.get("strict").map(v => base ++ Map("strict" -> v)).getOrElse(base) - } - ) - } else normalized - setScalarParam(responseFormat, payload) + setScalarParam(responseFormat, ResponseFormatUtils.chat(value)) } @@ -96,12 +88,12 @@ trait HasOpenAITextParamsExtended extends HasOpenAITextParams with HasOpenAIResp ) } -object OpenAIChatCompletion extends ComplexParamsReadable[OpenAIChatCompletion] - class OpenAIChatCompletion(override val uid: String) extends OpenAIServicesBase(uid) with HasOpenAITextParamsExtended with HasMessagesInput with HasCognitiveServiceInput with HasOpenAIFabricHeaders with HasInternalJsonOutputParser with SynapseMLLogging with HasRAIContentFilter with HasTextOutput { + import OpenAIChatCompletion._ + logClass(FeatureNames.AiServices.OpenAI) def this() = this(Identifiable.randomUID("OpenAIChatCompletion")) @@ -207,6 +199,38 @@ class OpenAIChatCompletion(override val uid: String) extends OpenAIServicesBase( super.transformSchema(schema) } + private[openai] def getStringEntity(messages: Seq[Row], optionalParams: Map[String, Any]): StringEntity = { + getStringEntityCollectionSeq(messages, optionalParams) + } + + private def getStringEntityCollectionSeq( + messages: CollectionSeq[Row], + optionalParams: Map[String, Any] + ): StringEntity = { + val mappedMessages = encodedMessageMaps(messages) + val fullPayload = optionalParams.updated("messages", mappedMessages) + val encoded = requestBody.encode(fullPayload, get(responseFormat).orElse(getDefault(responseFormat))) + new StringEntity(encoded, ContentType.APPLICATION_JSON) + } + + override private[openai] def getOutputMessageText(outputColName: String): org.apache.spark.sql.Column = { + F.element_at(F.col(outputColName).getField("choices"), 1) + .getField("message").getField("content") + } + + override private[openai] def isContentFiltered(outputRow: Row): Boolean = { + val result = ChatModelResponse.makeFromRowConverter(outputRow) + val firstChoice = result.choices.head + Option(firstChoice.message.content).isEmpty + } + + override private[openai] def getFilterReason(outputRow: Row): String = { + val result = ChatModelResponse.makeFromRowConverter(outputRow) + result.choices.head.finish_reason + } +} + +object OpenAIChatCompletion extends ComplexParamsReadable[OpenAIChatCompletion] { private def runtimeType(value: Any): String = { Option(value).map(_.getClass.getSimpleName).getOrElse("null") } @@ -539,34 +563,25 @@ class OpenAIChatCompletion(override val uid: String) extends OpenAIServicesBase( } } - private[openai] def getStringEntity(messages: Seq[Row], optionalParams: Map[String, Any]): StringEntity = { - getStringEntityCollectionSeq(messages, optionalParams) - } - - private def getStringEntityCollectionSeq( - messages: CollectionSeq[Row], - optionalParams: Map[String, Any] - ): StringEntity = { - val mappedMessages = encodedMessageMaps(messages) - val fullPayload = optionalParams.updated("messages", mappedMessages) - val encoded = requestBody.encode(fullPayload, get(responseFormat).orElse(getDefault(responseFormat))) - new StringEntity(encoded, ContentType.APPLICATION_JSON) - } - - override private[openai] def getOutputMessageText(outputColName: String): org.apache.spark.sql.Column = { - F.element_at(F.col(outputColName).getField("choices"), 1) - .getField("message").getField("content") - } - - override private[openai] def isContentFiltered(outputRow: Row): Boolean = { - val result = ChatModelResponse.makeFromRowConverter(outputRow) - val firstChoice = result.choices.head - Option(firstChoice.message.content).isEmpty - } - - override private[openai] def getFilterReason(outputRow: Row): String = { - val result = ChatModelResponse.makeFromRowConverter(outputRow) - result.choices.head.finish_reason + private[services] def strictRequestMessages(messages: Seq[Row]): Seq[Map[String, Any]] = { + validateMessages(messages) + messages.zipWithIndex.foreach { case (message, messageIndex) => + val (fieldIndex, dataType) = contentField(message, messageIndex) + dataType match { + case ArrayType(_: MapType, _) => + val parts = mapBackedContentParts(message.get(fieldIndex), messageIndex) + if (parts.isEmpty) { + throw new IllegalArgumentException(s"messages[$messageIndex].content must not be empty") + } + if (!parts.exists(_.get("type").contains("image_url"))) { + parts.zipWithIndex.foreach { case (part, partIndex) => + validateContentPart(part, messageIndex, partIndex) + } + } + case _ => + } + } + encodedMessageMaps(messages) } } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisions.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisions.scala new file mode 100644 index 00000000000..cfeaa11aa03 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisions.scala @@ -0,0 +1,20 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.openai + +import com.microsoft.azure.synapse.ml.services.ai.AIPrompt +import org.apache.spark.ml.ExplicitComplexParamsReadable +import org.apache.spark.ml.util.Identifiable + +object OpenAIDecisions extends ExplicitComplexParamsReadable[OpenAIDecisions] { + override protected def expectedClass: Class[OpenAIDecisions] = classOf[OpenAIDecisions] +} + +/** Native OpenAI Decisions requests; endpoint availability must be qualified separately. */ +class OpenAIDecisions(override val uid: String) extends AIPrompt(uid) { + def this() = this(Identifiable.randomUID("OpenAIDecisions")) + + override protected def supportedApiTypes: Array[String] = Array("decisions") + override protected def defaultApiType: String = "decisions" +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala index 433f32fbb19..02f07e2b20d 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala @@ -114,6 +114,8 @@ class OpenAIPrompt(override val uid: String) extends Transformer private var postProcessingExplicitlySet: Boolean = false + protected def hasExplicitPostProcessing: Boolean = postProcessingExplicitlySet + def setPostProcessing(value: String): this.type = { OpenAIPromptPostProcessing.inferMode(getPostProcessingOptions) .foreach(expectedMode => OpenAIPromptPostProcessing.validateModeValue(value, expectedMode)) @@ -130,7 +132,7 @@ class OpenAIPrompt(override val uid: String) extends Transformer def setPostProcessingOptions(value: Map[String, String]): this.type = { def setOrValidatePostProcessing(expected: String): Unit = { if (isSet(postProcessing)) { - if (getPostProcessing.isEmpty && !postProcessingExplicitlySet) { + if (getPostProcessing.isEmpty && !hasExplicitPostProcessing) { set(postProcessing, expected) } else { OpenAIPromptPostProcessing.validateModeValue(getPostProcessing, expected) @@ -166,8 +168,11 @@ class OpenAIPrompt(override val uid: String) extends Transformer override protected def pySetParamsFunc: String = OpenAIPromptPythonOverrides.setParamsFunc(super.pySetParamsFunc) + protected def pyHasExplicitPostProcessing: String = "self._post_processing_explicitly_set" + override def pyAdditionalMethods: String = - super.pyAdditionalMethods + OpenAIPromptPythonOverrides.AdditionalMethods + super.pyAdditionalMethods + + OpenAIPromptPythonOverrides.additionalMethods(pyClassName, pyHasExplicitPostProcessing) override def pyInitFunc(): String = OpenAIPromptPythonOverrides.initFunc(super.pyInitFunc()) @@ -186,9 +191,15 @@ class OpenAIPrompt(override val uid: String) extends Transformer def setSystemPrompt(value: String): this.type = set(systemPrompt, value) + /** Called during superclass construction; overrides must not depend on subclass instance state. */ + protected def supportedApiTypes: Array[String] = Array("chat_completions", "responses") + + /** Called during superclass construction; overrides must not depend on subclass instance state. */ + protected def defaultApiType: String = "chat_completions" + val apiType = new Param[String]( - this, "apiType", "The OpenAI API type to use: 'chat_completions' or 'responses'", - isValid = ParamValidators.inArray(Array("chat_completions", "responses"))) + this, "apiType", s"The OpenAI API type to use: ${supportedApiTypes.map(t => s"'$t'").mkString(" or ")}", + isValid = ParamValidators.inArray(supportedApiTypes)) GlobalParams.registerParam(apiType, OpenAIApiTypeKey) @@ -275,7 +286,7 @@ class OpenAIPrompt(override val uid: String) extends Transformer messagesCol -> (this.uid + "_messages"), dropPrompt -> true, systemPrompt -> defaultSystemPrompt, - apiType -> "chat_completions", + apiType -> defaultApiType, columnTypes -> Map.empty, timeout -> 360.0, store -> Left(false) @@ -579,10 +590,17 @@ class OpenAIPrompt(override val uid: String) extends Transformer ) } - private[openai] def createMessagesForRow( + private[services] def createMessagesForRow( userMessage: String, attachmentMap: Map[String, String], attachmentOrder: Seq[String] + ): Seq[OpenAICompositeMessage] = createMessagesForRow(userMessage, attachmentMap, attachmentOrder, None) + + private[services] def createMessagesForRow( + userMessage: String, + attachmentMap: Map[String, String], + attachmentOrder: Seq[String], + conf: Option[Configuration] ): Seq[OpenAICompositeMessage] = { // Filter to get only non-null, non-empty path values val orderedAttachments = attachmentOrder.flatMap { columnName => @@ -593,46 +611,22 @@ class OpenAIPrompt(override val uid: String) extends Transformer if (attachmentOrder.nonEmpty && orderedAttachments.isEmpty) { null //scalastyle:ignore null } else { - val contentParts = buildContentParts(userMessage, orderedAttachments) - val messages = getPromptsForMessage(Left(contentParts)) - messages + val contentParts = buildContentParts(userMessage, orderedAttachments, conf) + getPromptsForMessage(Left(contentParts)) } } - private def buildContentParts(promptText: String, attachmentPaths: Seq[String]): Seq[Map[String, String]] = { - var parts = Seq(stringMessageWrapper(promptText)) - if (!attachmentPaths.isEmpty) { - parts = parts ++ attachmentPaths.flatMap(wrapFileToMessagesList) - } - parts - } - - private def prepareFile(filePathStr: String): (String, Array[Byte], String, String) = { - val limit = if (isSet(fileSizeLimitMB)) Some(getFileSizeLimitMB) else None - OpenAIAttachmentUtils.prepareFile( - filePathStr, limit, imageExtensions, audioExtensions, textExtensions) - } + private def buildContentParts(promptText: String, attachmentPaths: Seq[String], + conf: Option[Configuration]): Seq[Map[String, String]] = + Seq(stringMessageWrapper(promptText)) ++ attachmentPaths.flatMap(path => wrapFileToMessagesList(path, conf)) - private def makeResponsesFileMessage( - fileName: String, - fileBytes: Array[Byte], - fileType: String, - mimeType: String - ): Map[String, String] = { - OpenAIAttachmentUtils.responsesMessage(fileName, fileBytes, fileType, mimeType, stringMessageWrapper) - } + private def attachmentSizeLimit: Option[Double] = + get(fileSizeLimitMB).orElse(if (useBoundedAttachments) getDefault(fileSizeLimitMB) else None) - private def makeChatCompletionsFileMessage( - fileName: String, - fileBytes: Array[Byte], - fileType: String, - mimeType: String - ): Map[String, String] = { - OpenAIAttachmentUtils.chatCompletionsMessage(fileBytes, fileType, mimeType, stringMessageWrapper) - } + protected def useBoundedAttachments: Boolean = false - private def wrapFileToMessagesList(filePathStr: String): Seq[Map[String, String]] = { - val limit = if (isSet(fileSizeLimitMB)) Some(getFileSizeLimitMB) else None + private def wrapFileToMessagesList(filePathStr: String, conf: Option[Configuration]): Seq[Map[String, String]] = { + val limit = attachmentSizeLimit val fileMessage = if (OpenAIAttachmentUtils.isDataUrl(filePathStr)) { this.getApiType match { case "responses" => @@ -643,12 +637,14 @@ class OpenAIPrompt(override val uid: String) extends Transformer filePathStr, limit, imageExtensions, audioExtensions, textExtensions, stringMessageWrapper) } } else { - val (fileName, fileBytes, fileType, mimeType) = prepareFile(filePathStr) + val (fileName, fileBytes, fileType, mimeType) = OpenAIAttachmentUtils.prepareFile( + filePathStr, limit, imageExtensions, audioExtensions, textExtensions, + useBoundedAttachments, conf.getOrElse(new Configuration())) this.getApiType match { case "responses" => - makeResponsesFileMessage(fileName, fileBytes, fileType, mimeType) + OpenAIAttachmentUtils.responsesMessage(fileName, fileBytes, fileType, mimeType, stringMessageWrapper) case "chat_completions" => - makeChatCompletionsFileMessage(fileName, fileBytes, fileType, mimeType) + OpenAIAttachmentUtils.chatCompletionsMessage(fileBytes, fileType, mimeType, stringMessageWrapper) } } Seq(fileMessage) @@ -712,7 +708,7 @@ class OpenAIPrompt(override val uid: String) extends Transformer case Some(inferredMode) => val configuredMode = get(postProcessing) .getOrElse(throw new IllegalArgumentException(s"postProcessing must be '$inferredMode'")) - if (configuredMode.isEmpty && postProcessingExplicitlySet) { + if (configuredMode.isEmpty && hasExplicitPostProcessing) { throw new IllegalArgumentException(s"postProcessing must be '$inferredMode'") } if (configuredMode.nonEmpty) { @@ -725,7 +721,7 @@ class OpenAIPrompt(override val uid: String) extends Transformer effectivePostProcessing } - private def getParser: OutputParser = { + private[services] def getParser: OutputParser = { val opts = getPostProcessingOptions val effectivePostProcessing = getEffectivePostProcessing effectivePostProcessing.toLowerCase match { diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParsers.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParsers.scala index 52bf0bf0add..83c35b084e3 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParsers.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParsers.scala @@ -3,7 +3,10 @@ package com.microsoft.azure.synapse.ml.services.openai +import com.fasterxml.jackson.core.{JsonProcessingException, JsonToken} +import com.microsoft.azure.synapse.ml.services.ai.AIJson import org.apache.spark.sql.{Column, functions => F, types => T} +import org.apache.spark.sql.internal.SQLConfInjections import org.apache.spark.sql.types.DataType trait OutputParser { @@ -35,6 +38,46 @@ class JsonParser(val schema: String, options: Map[String, String]) extends Outpu def parse(responseCol: Column): Column = F.from_json(cleanJsonString(responseCol), schema, options) + private def acceptsSingletonObject: Boolean = outputSchema match { + case T.ArrayType(_: T.StructType, _) => true + case _ => false + } + + private[services] def hasParseError(responseCol: Column): Column = { + val factory = SQLConfInjections.jsonFactory(options) + val acceptsSingleton = acceptsSingletonObject + val diagnosticJson = F.udf((text: String) => Option(text).flatMap { value => + val parser = factory.createParser(value) + try { + Option(parser.nextToken()).flatMap { root => + if (!AIJson.hasValidStringTokens(parser) || Option(parser.nextToken()).nonEmpty) None + else { + // Spark accepts an object as a singleton array only at the document root. + val nested = if (acceptsSingleton && root == JsonToken.START_OBJECT) "[" + value + "\n]" else value + Some("{\"value\":" + nested + "\n}") + } + } + } catch { + case _: JsonProcessingException => None + } finally { + parser.close() + } + }) + // A struct envelope exposes corrupt records for nested conversions in array and map roots as well. + val corrupt = "__synapseml_corrupt_record" + val checkedSchema = new T.StructType().add("value", outputSchema).add(corrupt, T.StringType) + val cleaned = cleanJsonString(responseCol) + val wrapped = diagnosticJson(cleaned) + // The root check enforces the original limit before adding either diagnostic envelope. + val extraDepth = if (acceptsSingleton) 2 else 1 + val depth = math.min(Int.MaxValue.toLong, + factory.streamReadConstraints().getMaxNestingDepth.toLong + extraDepth) + val checkedOptions = options.filterNot { case (key, _) => + key.equalsIgnoreCase("maxNestingDepth") || key.equalsIgnoreCase("columnNameOfCorruptRecord") + } ++ Map("maxNestingDepth" -> depth.toString, "columnNameOfCorruptRecord" -> corrupt) + wrapped.isNull || F.from_json(wrapped, checkedSchema, checkedOptions).getField(corrupt).isNotNull + } + def outputSchema: T.DataType = DataType.fromDDL(schema) } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptPythonOverrides.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptPythonOverrides.scala index d78a00e6d6c..3121731f7d8 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptPythonOverrides.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptPythonOverrides.scala @@ -3,8 +3,9 @@ package com.microsoft.azure.synapse.ml.services.openai -private[openai] object OpenAIPromptPythonOverrides { - private def replaceExactlyOnce(value: String, target: String, replacement: String, errorMessage: String): String = { +private[services] object OpenAIPromptPythonOverrides { + private[services] def replaceExactlyOnce(value: String, target: String, + replacement: String, errorMessage: String): String = { require(value.sliding(target.length).count(_ == target) == 1, errorMessage) value.replace(target, replacement) } @@ -46,7 +47,7 @@ private[openai] object OpenAIPromptPythonOverrides { | ) | if ( | post_processing == "" - | and not self._post_processing_explicitly_set + | and not self._post_processing_is_explicit() | ): | post_processing = None | inferred_post_processing = self._validate_post_processing_options( @@ -117,7 +118,7 @@ private[openai] object OpenAIPromptPythonOverrides { | if ( | post_processing == "" | and not post_processing_explicit - | and not self._post_processing_explicitly_set + | and not self._post_processing_is_explicit() | ): | post_processing = None | inferred_post_processing = self._validate_post_processing_options( @@ -139,7 +140,32 @@ private[openai] object OpenAIPromptPythonOverrides { ) } - val AdditionalMethods: String = + def additionalMethods(className: String, postProcessingIsExplicit: String): String = + AdditionalMethods + s""" + |def _post_processing_is_explicit(self): + | return $postProcessingIsExplicit + | + |def clear(self, param): + | if param == self.postProcessing: + | self._post_processing_explicitly_set = False + | return super($className, self).clear(param) + | + |def copy(self, extra=None): + | if extra is None: + | extra = {} + | result = super($className, self).copy(extra) + | result._post_processing_explicitly_set = ( + | self._post_processing_explicitly_set + | or self.postProcessing in extra + | ) + | if self.postProcessingOptions in extra: + | result.setPostProcessingOptions(result.getPostProcessingOptions()) + | elif self.postProcessing in extra: + | result._validate_post_processing(result.getPostProcessing()) + | return result + |""".stripMargin + + private val AdditionalMethods: String = """ |def _to_java_post_processing_options(self, value): | value = self._normalize_post_processing_options(value) @@ -257,24 +283,5 @@ private[openai] object OpenAIPromptPythonOverrides { | if inferred_post_processing: | self._set(postProcessing=inferred_post_processing) | return self - | - |def clear(self, param): - | if param == self.postProcessing: - | self._post_processing_explicitly_set = False - | return super(OpenAIPrompt, self).clear(param) - | - |def copy(self, extra=None): - | if extra is None: - | extra = {} - | result = super(OpenAIPrompt, self).copy(extra) - | result._post_processing_explicitly_set = ( - | self._post_processing_explicitly_set - | or self.postProcessing in extra - | ) - | if self.postProcessingOptions in extra: - | result.setPostProcessingOptions(result.getPostProcessingOptions()) - | elif self.postProcessing in extra: - | result._validate_post_processing(result.getPostProcessing()) - | return result |""".stripMargin } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIResponses.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIResponses.scala index 427e37032a3..658a21d0c67 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIResponses.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIResponses.scala @@ -120,12 +120,12 @@ trait HasOpenAITextParamsResponses extends HasOpenAITextParams with HasOpenAIRes ) } -object OpenAIResponses extends ComplexParamsReadable[OpenAIResponses] - class OpenAIResponses(override val uid: String) extends OpenAIServicesBase(uid) with HasOpenAITextParamsResponses with HasMessagesInput with HasCognitiveServiceInput with HasOpenAIFabricHeaders with HasInternalJsonOutputParser with SynapseMLLogging with HasCustomHeaders with HasRAIContentFilter with HasTextOutput { + import OpenAIResponses._ + logClass(FeatureNames.AiServices.OpenAI) def this() = this(Identifiable.randomUID("OpenAIResponses")) @@ -221,6 +221,112 @@ class OpenAIResponses(override val uid: String) extends OpenAIServicesBase(uid) override def responseDataType: DataType = ResponsesModelResponse.schema + private def validatePublicColumnNames(): Unit = + OpenAIColumnUtils.validateDistinctColumns( + "messagesCol" -> getMessagesCol, + "outputCol" -> getOutputCol, + "errorCol" -> getErrorCol + ) + + override def transform(dataset: Dataset[_]): DataFrame = { + transferGlobalParamsToParamMap() + validatePublicColumnNames() + logTransform[DataFrame]({ + val df = dataset.toDF() + val colsToAvoid = df.schema.fieldNames.toSet ++ Set(getErrorCol, getOutputCol) + val originalMessagesCol = OpenAIColumnUtils.findUnusedColumnName("originalMessages")(colsToAvoid) + val validationErrorCol = OpenAIColumnUtils.findUnusedColumnName( + "responsesMessageValidationError")(colsToAvoid + originalMessagesCol) + val resolvedMessagesCol = OpenAIColumnUtils.resolvedColumnName(df, getMessagesCol) + val messagesDataType = df.schema(resolvedMessagesCol).dataType + + val validationErrorUDF = UDFUtils.oldUdf( + (messages: CollectionSeq[Row]) => structuredMessageValidationError(messages).map(message => + Row(message, null) //scalastyle:ignore null + ).orNull, + ErrorUtils.ErrorSchema + ) + + val validatedMessages = df + .withColumn(originalMessagesCol, F.col(resolvedMessagesCol)) + .withColumn(validationErrorCol, validationErrorUDF(F.col(originalMessagesCol))) + .withColumn( + getMessagesCol, + F.when(F.col(validationErrorCol).isNotNull, F.lit(null).cast(messagesDataType)) //scalastyle:ignore null + .otherwise(F.col(originalMessagesCol)) + ) + + val validatedWithErrors = + OpenAIColumnUtils.existingColumnOfType(df, getErrorCol, ErrorUtils.ErrorSchema) match { + case Some(existingErrorCol) => + validatedMessages.withColumn( + getErrorCol, + F.coalesce(F.col(existingErrorCol), F.col(validationErrorCol)) + ) + case None => + validatedMessages.withColumn(getErrorCol, F.col(validationErrorCol)) + } + + getInternalTransformer(validatedWithErrors.schema).transform(validatedWithErrors) + .drop(validationErrorCol) + .withColumn(getMessagesCol, F.col(originalMessagesCol)) + .drop(originalMessagesCol) + }, dataset.columns.length) + } + + override def transformSchema(schema: StructType): StructType = { + validatePublicColumnNames() + super.transformSchema(schema) + } + + private[openai] def getStringEntity(messages: Seq[Row], optionalParams: Map[String, Any]): StringEntity = { + val fullPayload = optionalParams.updated("input", requestMessages(messages)) + val encoded = requestBody.encode(fullPayload, get(responseFormat).orElse(getDefault(responseFormat))) + new StringEntity(encoded, ContentType.APPLICATION_JSON) + } + + override private[openai] def getOutputMessageText(outputColName: String): org.apache.spark.sql.Column = { + val outputEntries = F.col(outputColName).getField("output") + val lastOutputEntry = F.element_at(outputEntries, -1) + val textValues = F.transform(lastOutputEntry.getField("content"), part => part.getField("text")) + val definedTextValues = F.filter(textValues, text => text.isNotNull) + F.element_at(definedTextValues, 1) + } + + private def outputEntries(outputRow: Row): Seq[Row] = { + Option(outputRow).flatMap(r => Option(r.getAs[Seq[Row]]("output"))).getOrElse(Seq.empty) + } + + private def lastOutputEntry(outputRow: Row): Option[Row] = { + outputEntries(outputRow).lastOption + } + + private def firstDefinedText(contentParts: Seq[Row]): Option[String] = { + contentParts.iterator + .flatMap(part => Option(part.getAs[String]("text"))) + .toSeq + .headOption + } + + private def lastOutputText(outputRow: Row): Option[String] = { + lastOutputEntry(outputRow).flatMap { outputEntry => + firstDefinedText(Option(outputEntry.getAs[Seq[Row]]("content")).getOrElse(Seq.empty)) + } + } + + override private[openai] def isContentFiltered(outputRow: Row): Boolean = { + lastOutputText(outputRow).isEmpty + } + + override private[openai] def getFilterReason(outputRow: Row): String = { + lastOutputEntry(outputRow).iterator + .flatMap(outputEntry => Option(outputEntry.getAs[String]("status"))) + .find(_.nonEmpty) + .getOrElse("content_filtered_or_empty") + } +} + +object OpenAIResponses extends ComplexParamsReadable[OpenAIResponses] { private def runtimeType(value: Any): String = { Option(value).map(_.getClass.getSimpleName).getOrElse("null") } @@ -344,64 +450,6 @@ class OpenAIResponses(override val uid: String) extends OpenAIServicesBase(uid) } } - private def validatePublicColumnNames(): Unit = - OpenAIColumnUtils.validateDistinctColumns( - "messagesCol" -> getMessagesCol, - "outputCol" -> getOutputCol, - "errorCol" -> getErrorCol - ) - - override def transform(dataset: Dataset[_]): DataFrame = { - transferGlobalParamsToParamMap() - validatePublicColumnNames() - logTransform[DataFrame]({ - val df = dataset.toDF() - val colsToAvoid = df.schema.fieldNames.toSet ++ Set(getErrorCol, getOutputCol) - val originalMessagesCol = OpenAIColumnUtils.findUnusedColumnName("originalMessages")(colsToAvoid) - val validationErrorCol = OpenAIColumnUtils.findUnusedColumnName( - "responsesMessageValidationError")(colsToAvoid + originalMessagesCol) - val resolvedMessagesCol = OpenAIColumnUtils.resolvedColumnName(df, getMessagesCol) - val messagesDataType = df.schema(resolvedMessagesCol).dataType - - val validationErrorUDF = UDFUtils.oldUdf( - (messages: CollectionSeq[Row]) => structuredMessageValidationError(messages).map(message => - Row(message, null) //scalastyle:ignore null - ).orNull, - ErrorUtils.ErrorSchema - ) - - val validatedMessages = df - .withColumn(originalMessagesCol, F.col(resolvedMessagesCol)) - .withColumn(validationErrorCol, validationErrorUDF(F.col(originalMessagesCol))) - .withColumn( - getMessagesCol, - F.when(F.col(validationErrorCol).isNotNull, F.lit(null).cast(messagesDataType)) //scalastyle:ignore null - .otherwise(F.col(originalMessagesCol)) - ) - - val validatedWithErrors = - OpenAIColumnUtils.existingColumnOfType(df, getErrorCol, ErrorUtils.ErrorSchema) match { - case Some(existingErrorCol) => - validatedMessages.withColumn( - getErrorCol, - F.coalesce(F.col(existingErrorCol), F.col(validationErrorCol)) - ) - case None => - validatedMessages.withColumn(getErrorCol, F.col(validationErrorCol)) - } - - getInternalTransformer(validatedWithErrors.schema).transform(validatedWithErrors) - .drop(validationErrorCol) - .withColumn(getMessagesCol, F.col(originalMessagesCol)) - .drop(originalMessagesCol) - }, dataset.columns.length) - } - - override def transformSchema(schema: StructType): StructType = { - validatePublicColumnNames() - super.transformSchema(schema) - } - private def requireOnlyFields(part: Map[String, Any], allowed: Set[String], location: String): Unit = { if ((part.keySet -- allowed).nonEmpty) { throw new IllegalArgumentException(s"$location contains unsupported fields") @@ -538,8 +586,8 @@ class OpenAIResponses(override val uid: String) extends OpenAIServicesBase(uid) } } - private[openai] def getStringEntity(messages: Seq[Row], optionalParams: Map[String, Any]): StringEntity = { - val mappedMessages = encodeResponsesMessages(messages) + private[services] def requestMessages(messages: Seq[Row]): Seq[Map[String, Any]] = { + encodeResponsesMessages(messages) .zipWithIndex .map { case (message, messageIndex) => validateEncodedMessage(message, messageIndex) @@ -554,49 +602,6 @@ class OpenAIResponses(override val uid: String) extends OpenAIServicesBase(uid) case _ => m } } - val fullPayload = optionalParams.updated("input", mappedMessages) - val encoded = requestBody.encode(fullPayload, get(responseFormat).orElse(getDefault(responseFormat))) - new StringEntity(encoded, ContentType.APPLICATION_JSON) - } - - override private[openai] def getOutputMessageText(outputColName: String): org.apache.spark.sql.Column = { - val outputEntries = F.col(outputColName).getField("output") - val lastOutputEntry = F.element_at(outputEntries, -1) - val textValues = F.transform(lastOutputEntry.getField("content"), part => part.getField("text")) - val definedTextValues = F.filter(textValues, text => text.isNotNull) - F.element_at(definedTextValues, 1) - } - - private def outputEntries(outputRow: Row): Seq[Row] = { - Option(outputRow).flatMap(r => Option(r.getAs[Seq[Row]]("output"))).getOrElse(Seq.empty) - } - - private def lastOutputEntry(outputRow: Row): Option[Row] = { - outputEntries(outputRow).lastOption - } - - private def firstDefinedText(contentParts: Seq[Row]): Option[String] = { - contentParts.iterator - .flatMap(part => Option(part.getAs[String]("text"))) - .toSeq - .headOption - } - - private def lastOutputText(outputRow: Row): Option[String] = { - lastOutputEntry(outputRow).flatMap { outputEntry => - firstDefinedText(Option(outputEntry.getAs[Seq[Row]]("content")).getOrElse(Seq.empty)) - } - } - - override private[openai] def isContentFiltered(outputRow: Row): Boolean = { - lastOutputText(outputRow).isEmpty - } - - override private[openai] def getFilterReason(outputRow: Row): String = { - lastOutputEntry(outputRow).iterator - .flatMap(outputEntry => Option(outputEntry.getAs[String]("status"))) - .find(_.nonEmpty) - .getOrElse("content_filtered_or_empty") } } diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/ResponseFormatUtils.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/ResponseFormatUtils.scala index 5b02017599c..731510c2cd1 100644 --- a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/ResponseFormatUtils.scala +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/ResponseFormatUtils.scala @@ -21,6 +21,15 @@ object ResponseFormatUtils { "null" ) + private[services] def chat(value: Map[String, Any]): Map[String, Any] = { + val normalized = normalize(value) + if (normalized.get("type").contains("json_schema")) { + val schema = Map("name" -> normalized("name"), "schema" -> normalized("schema")) ++ + normalized.get("strict").map("strict" -> _) + Map("type" -> "json_schema", "json_schema" -> schema) + } else normalized + } + def normalize(value: Map[String, Any]): Map[String, Any] = { requireNonNull(value) val tOpt = rfType(value) diff --git a/cognitive/src/test/java/synapsemltest/ai/AIIncompatibleStageProbe.java b/cognitive/src/test/java/synapsemltest/ai/AIIncompatibleStageProbe.java new file mode 100644 index 00000000000..54279cf47dc --- /dev/null +++ b/cognitive/src/test/java/synapsemltest/ai/AIIncompatibleStageProbe.java @@ -0,0 +1,14 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package synapsemltest.ai; + +public final class AIIncompatibleStageProbe { + static { + System.setProperty("synapseml.aiprompt.test.incompatible.initialized", "true"); + } + + public AIIncompatibleStageProbe(String uid) { + System.setProperty("synapseml.aiprompt.test.incompatible.constructed", uid); + } +} diff --git a/cognitive/src/test/python/synapsemltest/services/aiprompt_fabric_e2e.py b/cognitive/src/test/python/synapsemltest/services/aiprompt_fabric_e2e.py new file mode 100644 index 00000000000..8edda628bc7 --- /dev/null +++ b/cognitive/src/test/python/synapsemltest/services/aiprompt_fabric_e2e.py @@ -0,0 +1,157 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +"""Opt-in managed Spark qualification using synthetic native HTTP replies. + +Supply candidate core, cognitive and cognitive test JARs to Fabric Spark using +--expected-core-jar, --expected-cognitive-jar and --expected-tests-jar. +Each expected name must exactly match the URL-decoded loaded JAR basename. +Run without Python optimization so qualification assertions remain enabled. +This script does not provision storage, acquire model credentials, or claim +live-provider qualification. +""" + +import argparse +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import unquote, urlsplit + +from pyspark.sql import DataFrame, SparkSession + + +def main(): + if not __debug__: + raise RuntimeError( + "Fabric qualification requires Python assertions; disable -O and PYTHONOPTIMIZE" + ) + parser = argparse.ArgumentParser() + parser.add_argument("--expected-core-jar", required=True) + parser.add_argument("--expected-cognitive-jar", required=True) + parser.add_argument("--expected-tests-jar", required=True) + args = parser.parse_args() + spark = SparkSession.builder.getOrCreate() + jvm = spark._jvm + loader = jvm.java.lang.Thread.currentThread().getContextClassLoader() + + def source(name): + cls = jvm.java.lang.Class.forName(name, True, loader) + return str(cls.getProtectionDomain().getCodeSource().getLocation().toString()) + + provenance = { + "core": source("org.apache.spark.ml.ExplicitComplexParamsWritable"), + "cognitive": source("com.microsoft.azure.synapse.ml.services.ai.AIPrompt"), + "basePrompt": source( + "com.microsoft.azure.synapse.ml.services.openai.OpenAIPrompt" + ), + "probe": source( + "com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricProbe" + ), + } + for role, expected in ( + ("core", args.expected_core_jar), + ("cognitive", args.expected_cognitive_jar), + ("basePrompt", args.expected_cognitive_jar), + ("probe", args.expected_tests_jar), + ): + basename = unquote(urlsplit(provenance[role]).path.rsplit("/", 1)[-1]) + assert expected and basename == expected, provenance + implicit_apis = json.loads( + jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricProbe.verifyImplicitEndpoints( + spark._jsparkSession + ) + ) + assert implicit_apis == ["chat_completions", "responses", "embeddings"] + sends = [] + lock = threading.Lock() + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + request = json.loads( + self.rfile.read(int(self.headers["Content-Length"])).decode("utf-8") + ) + with lock: + sends.append(request["state"]) + assert self.path == "/providers/microsoft/v1/systemone" + body = json.dumps( + { + "model": "synthetic-systemone", + "answers": {"urgent": True, "team": "engineering", "severity": 2}, + "usage": {"input_tokens": 3, "output_tokens": 0, "total_tokens": 3}, + } + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + questions = { + "urgent": {"type": "noul", "instructions": "Is this urgent?"}, + "team": { + "type": "choice", + "instructions": "Choose a team.", + "criteria": {"billing": "Billing", "engineering": "Outages"}, + }, + "severity": { + "type": "score", + "instructions": "Rate severity.", + "criteria": ["Cosmetic", "Major", "Critical"], + }, + } + host = spark.conf.get("spark.driver.host") + server = ThreadingHTTPServer((host, 0), Handler) + endpoint = f"http://{host}:{server.server_port}" + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + stage = ( + jvm.com.microsoft.azure.synapse.ml.services.ai.AISystemOne() + .setUrl(endpoint) + .setModel("synthetic-model") + .setStateCol("text") + .setQuestions(json.dumps(questions)) + .setOutputCol("answer") + .setErrorCol("error") + .setUsageCol("usage") + .setRawResponseCol("raw") + .setConcurrency(2) + ) + data = spark.createDataFrame( + [(1, "outage"), (2, "outage"), (3, None), (4, "billing")], + "id INT, text STRING", + ).repartition(2) + expected = stage.transformSchema(data._jdf.schema()).json() + result = DataFrame(stage.transform(data._jdf), spark) + assert sends == [], "Transform must remain lazy" + assert json.loads(result.schema.json()) == json.loads(expected) + rows = result.collect() + assert len(rows) == 4 and len(sends) == 3 + assert sum(row.error is not None for row in rows) == 1 + assert sum(row.text == "outage" for row in rows) == 2 + assert all( + set(row.answer.answers) == set(questions) + for row in rows + if row.error is None + ) + evidence = { + "applicationId": spark.sparkContext.applicationId, + "sparkVersion": spark.version, + "classSources": provenance, + "rows": len(rows), + "physicalSends": len(sends), + "implicitApis": implicit_apis, + "synthetic": True, + } + print(json.dumps(evidence, sort_keys=True)) + finally: + server.shutdown() + server.server_close() + thread.join() + + +if __name__ == "__main__": + main() diff --git a/cognitive/src/test/python/synapsemltest/services/test_AIPrompt.py b/cognitive/src/test/python/synapsemltest/services/test_AIPrompt.py new file mode 100644 index 00000000000..8ea4072bd68 --- /dev/null +++ b/cognitive/src/test/python/synapsemltest/services/test_AIPrompt.py @@ -0,0 +1,1516 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +import ast +import importlib.util +import json +import tempfile +import threading +import unittest +from contextlib import contextmanager +from decimal import Decimal +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest import mock + +from pyspark.errors import IllegalArgumentException +from pyspark.ml import Pipeline +from pyspark.sql import SparkSession +from pyspark.sql.types import StructType +from synapse.ml.core.init_spark import init_spark +from synapse.ml.services.ai import AIPrompt, AISystemOne +from synapse.ml.services.openai import OpenAIDecisions, OpenAIEmbedding, OpenAIPrompt + +spark = SparkSession.getActiveSession() or init_spark() + + +@contextmanager +def endpoint(response, exact_numbers=False): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + requests.append( + ( + self.path, + json.loads( + self.rfile.read(int(self.headers["Content-Length"])), + parse_float=Decimal if exact_numbers else float, + ), + ) + ) + body = json.dumps(response).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + worker = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + yield f"http://127.0.0.1:{server.server_port}", requests + finally: + server.shutdown() + server.server_close() + worker.join() + + +class TestAIPrompt(unittest.TestCase): + def test_http_framing_headers_reject_after_python_copy_and_persistence(self): + data = spark.createDataFrame([("hello",)], "text STRING") + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest(stage=stage_type.__name__), endpoint({}) as ( + url, + requests, + ), tempfile.TemporaryDirectory() as directory: + stage = stage_type( + url=url if stage_type is AISystemOne else url + "/v1", + model="model", + inputCol="text", + outputCol="answer", + errorCol="error", + rawResponseCol="raw", + usageCol="usage", + responseIdCol="response_id", + ) + if stage_type is AISystemOne: + stage.setQuestions( + {"check": {"type": "noul", "instructions": "Check"}} + ) + elif stage_type is OpenAIDecisions: + stage.setQuestions( + [ + { + "type": "predicate", + "name": "check", + "instructions": "Check", + } + ] + ) + path = str(Path(directory) / "stage") + stage.write().save(path) + loaded = stage_type.load(path).copy({}) + loaded.setCustomHeaders({"cOnTeNt-LeNgTh": "0"}) + with self.assertRaisesRegex( + IllegalArgumentException, "HTTP framing headers" + ): + loaded.transform(data) + self.assertEqual(requests, []) + + def test_explicit_java_wrappers_preserve_uid_defaults_and_parameters(self): + names = ("outputCol", "errorCol", "messagesCol", "responseIdCol") + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + for index in range(2): + with self.subTest(stage=stage_type.__name__, index=index): + java_type = spark._jvm + for part in stage_type.getJavaPackage().split("."): + java_type = getattr(java_type, part) + original = ( + java_type(f"wrapped-{stage_type.__name__}-{index}") + .setConcurrency(index + 2) + .setInputCol("text") + .setModelCol("model") + ) + defaults = { + name: original.getDefault(original.getParam(name)).get() + for name in names + } + wrapped = stage_type(java_obj=original) + for name, expected in defaults.items(): + self.assertFalse(wrapped.isSet(name)) + self.assertEqual(wrapped.getOrDefault(name), expected) + self.assertEqual(wrapped.uid, original.uid()) + for candidate in (wrapped, wrapped.copy({})): + self.assertEqual(candidate.uid, original.uid()) + self.assertEqual(candidate.getConcurrency(), index + 2) + self.assertEqual(candidate.getInputCol(), "text") + self.assertEqual(candidate.getModelCol(), "model") + candidate._transfer_params_to_java() + for name, expected in defaults.items(): + self.assertFalse(candidate.isSet(name)) + self.assertEqual(candidate.getOrDefault(name), expected) + self.assertEqual( + original.getDefault(original.getParam(name)).get(), + expected, + ) + + def test_fresh_instances_preserve_their_own_java_defaults_and_explicit_columns( + self, + ): + names = ("outputCol", "errorCol", "messagesCol", "responseIdCol") + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest(stage=stage_type.__name__): + stages = [stage_type(), stage_type()] + for name in names: + self.assertNotEqual( + stages[0].getOrDefault(name), stages[1].getOrDefault(name) + ) + for stage in stages: + self.assertFalse(stage.isSet(name)) + self.assertEqual( + stage.getOrDefault(name), + stage._java_obj.getOrDefault( + stage._java_obj.getParam(name) + ), + ) + explicit = {name: "explicit_" + name for name in names} + stage = stage_type(**explicit) + for name, value in explicit.items(): + self.assertTrue(stage.isSet(name)) + self.assertEqual(stage.getOrDefault(name), value) + stage.clear(stage.getParam(name)) + self.assertFalse(stage.isSet(name)) + self.assertEqual( + stage.getOrDefault(name), + stage._java_obj.getDefault( + stage._java_obj.getParam(name) + ).get(), + ) + + def test_fresh_pipeline_stages_retain_both_default_outputs(self): + data = spark.createDataFrame([("hello",)], "text STRING") + replies = { + "chat_completions": { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ] + }, + "responses": { + "id": "response-id", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + }, + } + for api, reply in replies.items(): + with self.subTest(api=api), endpoint(reply) as (url, requests): + stages = [ + AIPrompt( + apiType=api, + url=url + "/v1", + model="model", + inputCol="text", + dropPrompt=False, + ) + for _ in range(2) + ] + if api == "responses": + for stage in stages: + stage.setStore(True) + result = Pipeline(stages=stages).fit(data).transform(data) + row = result.head() + self.assertEqual(len(requests), 2) + for name in ("outputCol", "errorCol", "messagesCol") + ( + ("responseIdCol",) if api == "responses" else () + ): + columns = [stage.getOrDefault(name) for stage in stages] + self.assertEqual(len(set(columns)), 2) + for column in columns: + self.assertIn(column, result.columns) + if name == "outputCol": + self.assertEqual(row[column], "answer") + elif name == "errorCol": + self.assertIsNone(row[column]) + elif name == "messagesCol": + self.assertEqual( + [ + message["content"][0]["text"] + for message in row[column] + if message["role"] == "user" + ], + ["hello"], + ) + else: + self.assertEqual(row[column], "response-id") + + def test_java_parameter_transfer_preserves_include_defaults_argument(self): + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + for include_defaults in (False, True, None): + for keyword in (False, True): + with self.subTest( + stage=stage_type.__name__, + include_defaults=include_defaults, + keyword=keyword, + ): + stage = stage_type() + java_default = stage.getConcurrency() + stage._setDefault(concurrency=java_default + 1) + if include_defaults is None: + stage._transfer_params_from_java() + elif keyword: + stage._transfer_params_from_java( + include_defaults=include_defaults + ) + else: + stage._transfer_params_from_java(include_defaults) + self.assertEqual( + stage.getConcurrency(), + java_default + 1 + if include_defaults is False + else java_default, + ) + self.assertFalse(stage.isSet(stage.concurrency)) + stage.setConcurrency(java_default + 2) + stage._transfer_params_to_java() + stage._transfer_params_from_java() + self.assertEqual(stage.getConcurrency(), java_default + 2) + self.assertTrue(stage.isSet(stage.concurrency)) + + def test_native_loader_preserves_json_and_runtime_class(self): + value = '{"decimal":0.12345678901234567890123456789,"large":1e400}' + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest( + stage=stage_type.__name__ + ), tempfile.TemporaryDirectory() as directory: + original = stage_type().setExtraBody(value) + saved = str(Path(directory) / "stage") + original.write().save(saved) + for reader in (stage_type, AIPrompt): + loaded = reader.load(saved) + self.assertIs(type(loaded), stage_type) + for candidate in (loaded, loaded.copy({})): + candidate._transfer_params_to_java() + restored = candidate._java_obj + self.assertEqual(candidate.getApiType(), original.getApiType()) + self.assertEqual( + restored.getClass().getName(), + original._java_obj.getClass().getName(), + ) + self.assertEqual( + restored.extraBody().jsonEncode( + restored.getOrDefault(restored.extraBody()) + ), + original._java_obj.extraBody().jsonEncode( + original._java_obj.getOrDefault( + original._java_obj.extraBody() + ) + ), + ) + self.assertTrue(restored.isSet(restored.extraBody())) + self.assertFalse(restored.isSet(restored.temperature())) + + def test_json_parameter_getters_preserve_exact_number_roundtrips(self): + numbers = ( + '{"precise":0.12345678901234567890123456789,"tiny":1e-400,' + '"huge":1e400,"integer":10000000000000000000000000000000000000000,' + '"ordinary":0.1,"half":0.5,"nested":[0.12345678901234567890123456789]}' + ) + values = { + "extraBody": '{"values":' + numbers + "}", + "responseFormat": '{"type":"json_schema","name":"answer","schema":' + + '{"type":"object","properties":{"value":{"type":"number"}},"values":' + + numbers + + "}}", + "questions": '[{"type":"predicate","instructions":"Check","values":' + + numbers + + "}]", + } + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + for name, text in values.items(): + with self.subTest( + stage=stage_type.__name__, parameter=name + ), tempfile.TemporaryDirectory() as directory: + suffix = name[0].upper() + name[1:] + original = stage_type() + parameter = original._java_obj.getParam(name) + if name == "responseFormat": + decoded = parameter.jsonDecode('{"left":' + text + "}") + original._java_obj.setResponseFormat(decoded.left().get()) + else: + getattr(original, "set" + suffix)(text) + encoded = parameter.jsonEncode( + original._java_obj.getOrDefault(parameter) + ) + path = str(Path(directory) / "stage") + original.write().save(path) + for candidate in ( + original, + original.copy({}), + stage_type.load(path), + ): + value = getattr(candidate, "get" + suffix)() + container = ( + value["json_schema"]["schema"] + if name == "responseFormat" + else value[0] + if name == "questions" + else value + ) + actual = container["values"] + for key in ("precise", "tiny", "huge"): + self.assertIsInstance(actual[key], Decimal) + self.assertEqual( + actual["precise"], + Decimal("0.12345678901234567890123456789"), + ) + self.assertEqual(actual["tiny"], Decimal("1e-400")) + self.assertEqual(actual["huge"], Decimal("1e400")) + self.assertEqual(actual["nested"], [actual["precise"]]) + self.assertIsInstance(actual["integer"], int) + self.assertIsInstance(actual["ordinary"], float) + self.assertEqual(actual["ordinary"], 0.1) + self.assertIsInstance(actual["half"], float) + for rebound in ( + getattr(candidate, "set" + suffix)(value), + stage_type(**{name: value}), + stage_type().setParams(**{name: value}), + ): + java_parameter = rebound._java_obj.getParam(name) + self.assertEqual( + java_parameter.jsonEncode( + rebound._java_obj.getOrDefault(java_parameter) + ), + encoded, + ) + + def test_json_parameters_reject_nonfinite_decimals_and_cycles_atomically(self): + cycle = [] + cycle.append(cycle) + for name in ("extraBody", "responseFormat", "questions"): + suffix = name[0].upper() + name[1:] + stage = getattr(AIPrompt(), "set" + suffix)({"type": "json_object"}) + for per_row in (False, True): + if per_row: + getattr(stage, "set" + suffix + "Col")("values") + parameter = stage._java_obj.getParam(name) + original = parameter.jsonEncode(stage._java_obj.getOrDefault(parameter)) + for invalid in ( + Decimal("NaN"), + Decimal("sNaN"), + Decimal("Infinity"), + Decimal("-Infinity"), + cycle, + ): + with self.subTest( + parameter=name, per_row=per_row, invalid=type(invalid).__name__ + ): + with self.assertRaises(ValueError): + getattr(stage, "set" + suffix)({"invalid": invalid}) + self.assertEqual( + parameter.jsonEncode( + stage._java_obj.getOrDefault(parameter) + ), + original, + ) + + def test_json_surrogates_reject_invalid_parameters_and_preserve_saved_unicode(self): + unicode = "caf\u00e9 \u96ea \U0001f600 \ufffd" + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + stage = stage_type() + for name in ("extraBody", "responseFormat", "questions"): + suffix = name[0].upper() + name[1:] + setter = getattr(stage, "set" + suffix) + setter({"type": "json_object"}) + parameter = stage._java_obj.getParam(name) + original = parameter.jsonEncode(stage._java_obj.getOrDefault(parameter)) + for malformed in ("\ud800", "\udc00", "\ud800x", "\udc00\ud800"): + for payload in ( + {"private-field": malformed}, + {malformed: "private-value", "?": 2}, + ): + values = ( + (payload,) + if name == "responseFormat" + else (payload, json.dumps(payload)) + ) + for value in values: + with self.subTest( + stage=stage_type.__name__, parameter=name + ): + with self.assertRaisesRegex( + IllegalArgumentException, "surrogate" + ): + setter(value) + self.assertEqual( + parameter.jsonEncode( + stage._java_obj.getOrDefault(parameter) + ), + original, + ) + stage.setExtraBody({unicode: [unicode]}) + with tempfile.TemporaryDirectory() as directory: + stage.write().save(directory + "/stage") + for candidate in ( + stage.copy({}), + stage_type.load(directory + "/stage"), + ): + self.assertEqual(candidate.getExtraBody(), {unicode: [unicode]}) + + def test_json_surrogates_remain_row_errors_without_corrupting_requests_or_answers( + self, + ): + reply = { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "\U0001f600"}} + ] + } + bodies = [ + json.dumps({"private-field": "\ud800"}), + json.dumps({"\udc00": "private-value", "?": 2}), + json.dumps({"\U0001f600": "\U0001f600", "?": 2}), + ] + data = spark.createDataFrame( + [(index, "hello", body) for index, body in enumerate(bodies)], + "id INT, text STRING, body STRING", + ) + with endpoint(reply) as (url, requests): + stage = AIPrompt( + url=url + "/v1", + model="model", + inputCol="text", + extraBodyCol="body", + outputCol="answer", + errorCol="error", + ) + rows = sorted(stage.transform(data).collect(), key=lambda row: row.id) + for row in rows[:2]: + self.assertIsNone(row.answer) + self.assertEqual( + json.loads(row.error.response)["code"], "invalid_request" + ) + self.assertIsNone(rows[2].error) + self.assertEqual(rows[2].answer, "\U0001f600") + self.assertEqual(len(requests), 1) + self.assertEqual(requests[0][1]["\U0001f600"], "\U0001f600") + self.assertEqual(requests[0][1]["?"], 2) + for malformed in ("\ud800", "\udc00"): + reply["choices"][0]["message"]["content"] = malformed + with endpoint(reply) as (url, requests): + stage = AIPrompt( + url=url + "/v1", + model="model", + inputCol="text", + outputCol="answer", + errorCol="error", + rawResponseCol="raw", + ) + row = stage.transform(data.limit(1)).head() + self.assertIsNone(row.answer) + self.assertEqual(json.loads(row.error.response)["code"], "parse_error") + self.assertEqual(row.raw, json.dumps(reply)) + self.assertEqual(len(requests), 1) + + def test_scalar_header_preflight_rejects_configuration_without_sending_requests( + self, + ): + with endpoint({}) as (url, requests): + data = spark.createDataFrame([("hello",)], "text STRING") + for configuration in ( + {"subscriptionKey": "synthetic-private\n"}, + {"AADToken": "synthetic-private\n"}, + {"customHeaders": {"X-Caller": "synthetic-private\n"}}, + ): + stage = AIPrompt( + url=url + "/v1", model="model", inputCol="text", **configuration + ) + with self.assertRaisesRegex( + IllegalArgumentException, "Invalid HTTP header" + ) as error: + stage.transform(data) + self.assertNotIn("synthetic-private", str(error.exception)) + self.assertEqual(requests, []) + + def test_mixed_header_preflight_validates_scalar_fallback_only_when_selected(self): + reply = { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ] + } + data = spark.createDataFrame( + [("valid", "synthetic-key"), ("null", None), ("blank", " ")], + "text STRING, key STRING", + ) + with endpoint(reply) as (url, requests): + stage = AIPrompt( + url=url + "/v1", + model="model", + inputCol="text", + subscriptionKeyCol="key", + AADToken="synthetic-private\n", + outputCol="answer", + errorCol="error", + ) + rows = {row.text: row for row in stage.transform(data).collect()} + self.assertIsNone(rows["valid"].error) + self.assertEqual(rows["valid"].answer, "answer") + for name in ("null", "blank"): + self.assertIsNone(rows[name].answer) + self.assertEqual( + json.loads(rows[name].error.response)["code"], "invalid_request" + ) + self.assertNotIn("synthetic-private", rows[name].error.response) + self.assertEqual(len(requests), 1) + + def test_structured_json_surrogates_produce_diagnostics_after_provider_decode(self): + data = spark.createDataFrame([("hello",)], "text STRING") + documents = ( + {"value": "\ud800"}, + {"\udc00": "private-value", "?": "other"}, + {"value": "\U0001f600", "\U0001f600": "unchanged"}, + ) + for api in ("chat_completions", "responses"): + for index, document in enumerate(documents): + text = json.dumps(document) + reply = ( + { + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + } + ] + } + if api == "chat_completions" + else { + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text}], + } + ] + } + ) + with self.subTest(api=api, document=index), endpoint(reply) as ( + url, + requests, + ): + stage = AIPrompt( + apiType=api, + url=url + "/v1", + model="model", + inputCol="text", + outputCol="answer", + errorCol="error", + rawResponseCol="raw", + postProcessingOptions={"jsonSchema": "MAP"}, + ) + row = stage.transform(data).head() + if index < 2: + self.assertIsNotNone(row.error) + self.assertEqual( + json.loads(row.error.response)["code"], "parse_error" + ) + self.assertNotIn("private-value", row.error.response) + else: + self.assertIsNone(row.error) + self.assertEqual(row.answer, document) + self.assertEqual(row.raw, json.dumps(reply)) + self.assertEqual(len(requests), 1) + + def test_leaf_loaders_reject_incompatible_saved_classes(self): + for original in (AIPrompt(), AISystemOne(), OpenAIDecisions()): + with self.subTest( + stage=type(original).__name__ + ), tempfile.TemporaryDirectory() as directory: + saved = str(Path(directory) / "stage") + original.write().save(saved) + loaded = AIPrompt.load(saved) + self.assertEqual( + loaded._java_obj.getClass().getName(), + original._java_obj.getClass().getName(), + ) + for reader in (AISystemOne, OpenAIDecisions): + if not isinstance(original, reader): + with self.assertRaisesRegex( + IllegalArgumentException, "incompatible" + ): + reader.load(saved) + + def test_public_stub_methods_exist_on_the_native_runtime_classes(self): + for stage_class in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest(stage=stage_class.__name__): + module = importlib.import_module(stage_class.__module__) + stub = ast.parse(Path(module.__file__).with_suffix(".pyi").read_text()) + definition = next( + node + for node in stub.body + if isinstance(node, ast.ClassDef) + and node.name == stage_class.__name__ + ) + for method in definition.body: + if isinstance( + method, ast.FunctionDef + ) and not method.name.startswith("_"): + self.assertTrue( + callable(getattr(stage_class, method.name, None)), + f"{stage_class.__name__}.{method.name} is stub-only", + ) + + def test_json_object_stub_inputs_accept_strings_and_dicts(self): + for stage_class in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest(stage=stage_class.__name__): + module = importlib.import_module(stage_class.__module__) + stub = ast.parse(Path(module.__file__).with_suffix(".pyi").read_text()) + definition = next( + node + for node in stub.body + if isinstance(node, ast.ClassDef) + and node.name == stage_class.__name__ + ) + methods = { + node.name: node + for node in definition.body + if isinstance(node, ast.FunctionDef) + } + input_type = "Union[str, Dict[str, Any]]" + for param in ("responseFormat", "extraBody"): + suffix = param[0].upper() + param[1:] + for method in ("__init__", "setParams"): + arguments = { + argument.arg: ast.unparse(argument.annotation) + for argument in methods[method].args.kwonlyargs + } + self.assertEqual(arguments[param], f"Optional[{input_type}]") + self.assertEqual(arguments[f"{param}Col"], "Optional[str]") + self.assertEqual( + ast.unparse(methods[f"set{suffix}"].args.args[1].annotation), + input_type, + ) + self.assertEqual( + ast.unparse(methods[f"get{suffix}"].returns), "Dict[str, Any]" + ) + self.assertEqual( + ast.unparse(methods[f"set{suffix}Col"].args.args[1].annotation), + "str", + ) + self.assertIn(f"{param} ({input_type})", stage_class.__doc__) + + def test_column_types_reject_unsupported_kinds_before_http(self): + data = spark.createDataFrame([("attachment.txt",)], "asset STRING") + for api in ("chat_completions", "responses"): + for binding in ("constructor", "setParams", "setter"): + with self.subTest(api=api, binding=binding), endpoint({}) as ( + url, + requests, + ): + arguments = { + "apiType": api, + "url": url + "/v1", + "model": "model", + "promptTemplate": "{asset}", + } + kinds = {"asset": "image"} + if binding == "constructor": + arguments["columnTypes"] = kinds + stage = AIPrompt(**arguments) + if binding == "setParams": + stage.setParams(columnTypes=kinds) + elif binding == "setter": + stage.setColumnTypes(kinds) + with self.assertRaisesRegex( + IllegalArgumentException, "columnTypes" + ): + stage.transform(data) + self.assertEqual(requests, []) + + def test_attachment_size_limits_reject_invalid_bindings_before_http(self): + data = spark.createDataFrame( + [("data:text/plain;base64,aGVsbG8=",)], "asset STRING" + ) + with endpoint({}) as ( + url, + requests, + ): + for api in ("chat_completions", "responses"): + for binding in ("constructor", "setParams", "setter", "default"): + for limit in (0.0, -1.0, float("nan")): + with self.subTest(api=api, binding=binding, limit=limit): + arguments = { + "apiType": api, + "url": url + "/v1", + "model": "model", + "promptTemplate": "{asset}", + "columnTypes": {"asset": "path"}, + } + if binding == "constructor": + arguments["fileSizeLimitMB"] = limit + stage = AIPrompt(**arguments) + if binding == "setParams": + stage.setParams(fileSizeLimitMB=limit) + elif binding == "setter": + stage.setFileSizeLimitMB(limit) + elif binding == "default": + stage._setDefault(fileSizeLimitMB=limit) + with self.assertRaisesRegex( + IllegalArgumentException, + "fileSizeLimitMB must be positive", + ): + stage.transform(data) + self.assertEqual(requests, []) + + def test_attachment_size_defaults_apply_without_overriding_explicit_values(self): + data = spark.createDataFrame( + [("data:text/plain;base64,aGVsbG8=",)], "asset STRING" + ) + replies = { + "chat_completions": { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ] + }, + "responses": { + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ] + }, + } + for api, reply in replies.items(): + with self.subTest(api=api), endpoint(reply) as (url, requests): + stage = AIPrompt( + apiType=api, + url=url + "/v1", + model="model", + promptTemplate="{asset}", + columnTypes={"asset": "path"}, + errorCol="error", + )._setDefault(fileSizeLimitMB=0.000001) + row = stage.transform(data).head() + self.assertIsNotNone(row.error) + self.assertEqual( + json.loads(row.error.response)["code"], "invalid_request" + ) + self.assertFalse(stage.isSet(stage.fileSizeLimitMB)) + self.assertEqual(requests, []) + stage.setFileSizeLimitMB(0.001) + self.assertIsNone(stage.transform(data).head().error) + self.assertEqual(len(requests), 1) + + def test_fabric_qualification_rejects_optimized_python_before_setup(self): + path = Path(__file__).with_name("aiprompt_fabric_e2e.py") + for optimize in (1, 2): + scenario = ModuleType("aiprompt_fabric_optimized") + exec( + compile( + path.read_text(encoding="utf-8"), + str(path), + "exec", + optimize=optimize, + ), + scenario.__dict__, + ) + with self.subTest(optimize=optimize), mock.patch.object( + scenario, "SparkSession" + ) as spark_type, mock.patch.object( + scenario.argparse.ArgumentParser, "parse_args" + ), mock.patch.object( + scenario, "ThreadingHTTPServer" + ) as server, mock.patch.object( + scenario, "threading" + ) as worker: + spark_type.builder.getOrCreate.side_effect = AssertionError( + "Spark setup must not be reached" + ) + with self.assertRaisesRegex(RuntimeError, "requires Python assertions"): + scenario.main() + spark_type.builder.getOrCreate.assert_not_called() + server.assert_not_called() + worker.Thread.assert_not_called() + + def test_fabric_jar_provenance_rejects_unexpected_artifacts_before_setup(self): + path = Path(__file__).with_name("aiprompt_fabric_e2e.py") + spec = importlib.util.spec_from_file_location( + "aiprompt_fabric_provenance", path + ) + scenario = importlib.util.module_from_spec(spec) + spec.loader.exec_module(scenario) + expected = { + "expected_core_jar": "candidate-core.jar", + "expected_cognitive_jar": "candidate-cognitive.jar", + "expected_tests_jar": "candidate-tests.jar", + } + sources = { + "org.apache.spark.ml.ExplicitComplexParamsWritable": "expected_core_jar", + "com.microsoft.azure.synapse.ml.services.ai.AIPrompt": "expected_cognitive_jar", + "com.microsoft.azure.synapse.ml.services.openai.OpenAIPrompt": "expected_cognitive_jar", + "com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricProbe": "expected_tests_jar", + } + for name, argument in sources.items(): + basename = expected[argument] + for actual, empty in ( + (f"file:/jars/backup-{basename}", False), + (f"file:/jars/{basename}.backup", False), + (f"file:/jars/{basename}/other.jar", False), + (f"file:/jars/other%2F{basename}", False), + (f"file:/jars/{basename}", True), + ): + args = SimpleNamespace(**expected) + if empty: + setattr(args, argument, "") + with self.subTest( + source=name, actual=actual, empty=empty + ), mock.patch.object( + scenario, "SparkSession" + ) as spark_type, mock.patch.object( + scenario.argparse.ArgumentParser, "parse_args", return_value=args + ), mock.patch.object( + scenario, "ThreadingHTTPServer" + ) as server, mock.patch.object( + scenario, "threading" + ) as worker: + session = spark_type.builder.getOrCreate.return_value + classes = {} + for class_name, expected_argument in sources.items(): + java_class = mock.Mock() + location = ( + java_class.getProtectionDomain.return_value.getCodeSource.return_value.getLocation.return_value + ) + location.toString.return_value = ( + actual + if class_name == name + else f"file:/jars/{expected[expected_argument]}" + ) + classes[class_name] = java_class + session._jvm.java.lang.Class.forName.side_effect = ( + lambda class_name, *unused: classes[class_name] + ) + with self.assertRaises(AssertionError): + scenario.main() + session._jvm.org.apache.hadoop.fs.Path.assert_not_called() + server.assert_not_called() + worker.Thread.assert_not_called() + + def test_fabric_route_probe_failure_does_not_start_http_server(self): + path = Path(__file__).with_name("aiprompt_fabric_e2e.py") + spec = importlib.util.spec_from_file_location("aiprompt_fabric_setup", path) + scenario = importlib.util.module_from_spec(spec) + spec.loader.exec_module(scenario) + args = SimpleNamespace( + expected_core_jar="candidate build.jar", + expected_cognitive_jar="candidate build.jar", + expected_tests_jar="candidate build.jar", + ) + with mock.patch.object( + scenario, "SparkSession" + ) as spark_type, mock.patch.object( + scenario.argparse.ArgumentParser, "parse_args", return_value=args + ), mock.patch.object( + scenario, "ThreadingHTTPServer" + ) as server, mock.patch.object( + scenario, "threading" + ) as worker: + session = spark_type.builder.getOrCreate.return_value + java_class = session._jvm.java.lang.Class.forName.return_value + location = ( + java_class.getProtectionDomain.return_value.getCodeSource.return_value.getLocation.return_value + ) + location.toString.return_value = "file:/jars/candidate%20build.jar" + probe = ( + session._jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricProbe + ) + probe.verifyImplicitEndpoints.side_effect = RuntimeError("Probe failed") + with self.assertRaisesRegex(RuntimeError, "Probe failed"): + scenario.main() + probe.verifyImplicitEndpoints.assert_called_once() + server.assert_not_called() + worker.Thread.assert_not_called() + + def test_jvm_saved_defaults_survive_python_loading_copy_and_transform(self): + cases = [ + ( + AIPrompt, + "chat_completions", + { + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "answer"}, + } + ] + }, + ), + ( + AIPrompt, + "responses", + { + "id": "response-id", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + }, + ), + ( + AISystemOne, + "systemone", + {"model": "model", "answers": {"urgent": {"result": True}}}, + ), + ( + OpenAIDecisions, + "decisions", + { + "model": "model", + "answers": [{"type": "predicate", "probability": 0.9}], + }, + ), + ] + data = spark.createDataFrame([("hello",)], "text STRING") + for stage_type, api, reply in cases: + with self.subTest(api=api), endpoint(reply) as (url, requests): + java_type = spark._jvm + for part in stage_type.getJavaPackage().split("."): + java_type = getattr(java_type, part) + original = ( + java_type("saved-" + api) + .setApiType(api) + .setUrl(url if api == "systemone" else url + "/v1") + .setModel("model") + .setInputCol("text") + ) + if api in ("chat_completions", "responses"): + original.setDropPrompt(False) + if api == "responses": + original.setStore(True) + elif api == "systemone": + original.setQuestions( + '{"urgent":{"type":"noul","instructions":"Is it urgent?"}}' + ) + elif api == "decisions": + original.setQuestions( + '[{"type":"predicate","instructions":"Is it urgent?"}]' + ) + defaults = { + name: original.getOrDefault(original.getParam(name)) + for name in ( + "outputCol", + "errorCol", + "messagesCol", + "responseIdCol", + ) + } + expected_schema = StructType.fromJson( + json.loads(original.transformSchema(data._jdf.schema()).json()) + ) + with tempfile.TemporaryDirectory() as directory: + original.write().save(directory + "/jvm") + loaded = stage_type.load(directory + "/jvm") + for candidate in (loaded, loaded.copy({})): + for name in ("columnTypes", "postProcessingOptions"): + value = candidate.getOrDefault(name) + self.assertIsInstance(value, dict) + self.assertEqual(value, {}) + self.assertFalse(value) + self.assertFalse(candidate.isSet(name), name) + for name, expected in defaults.items(): + self.assertFalse(candidate.isSet(name), name) + self.assertEqual(candidate.getOrDefault(name), expected) + result = candidate.transform(data) + self.assertEqual(result.schema, expected_schema) + self.assertIsNone(result.head()[defaults["errorCol"]]) + for name, expected in defaults.items(): + java_param = candidate._java_obj.getParam(name) + self.assertFalse(candidate._java_obj.isSet(java_param)) + self.assertEqual( + candidate._java_obj.getDefault(java_param).get(), + expected, + ) + candidate.setOutputCol("temporary_output") + candidate.clear(candidate.outputCol) + self.assertEqual( + candidate.getOutputCol(), defaults["outputCol"] + ) + self.assertEqual(len(requests), 2) + + def test_ordered_response_schema_preserves_exact_numbers_through_persistence_and_http( + self, + ): + names = ["e_one", "d_two", "c_three", "b_four", "a_five", "z_six"] + cases = [ + ( + "chat_completions", + { + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "answer"}, + } + ] + }, + ), + ( + "responses", + { + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + }, + ), + ] + for api, reply in cases: + for binding in ("schema", "setter", "constructor", "setParams"): + order = ( + names + if binding in ("schema", "constructor") + else list(reversed(names)) + ) + with self.subTest(api=api, binding=binding), endpoint( + reply, exact_numbers=True + ) as ( + url, + requests, + ): + schema = { + "type": "object", + "properties": { + name: { + "type": "number", + "maximum": 10**40 + index, + "minimum": Decimal("0.12345678901234567890123456789"), + } + for index, name in enumerate(order) + }, + "required": names, + "additionalProperties": False, + } + response_format = { + "type": "json_schema", + "name": "response_schema", + "schema": schema, + "strict": True, + } + arguments = dict(apiType=api, url=url + "/v1", model="model") + if binding == "constructor": + arguments["responseFormat"] = response_format + stage = ( + AIPrompt(**arguments).setInputCol("text").setErrorCol("error") + ) + stage.setExtraBody({name: name for name in order}) + if binding == "schema": + stage.setResponseSchema(schema) + elif binding == "setter": + stage.setResponseFormat(response_format) + elif binding == "setParams": + stage.setParams(responseFormat=response_format) + stage.setResponseFormat(stage.getResponseFormat()) + with tempfile.TemporaryDirectory() as directory: + stage.copy({}).write().save(directory + "/ordered") + loaded = AIPrompt.load(directory + "/ordered") + self.assertEqual(list(loaded.getExtraBody()), order) + row = loaded.transform( + spark.createDataFrame([("hello",)], "text STRING") + ).head() + self.assertIsNone(row.error) + self.assertEqual(len(requests), 1) + body = requests[0][1] + response_format = ( + body["response_format"]["json_schema"] + if api == "chat_completions" + else body["text"]["format"] + ) + self.assertEqual( + list(response_format["schema"]["properties"]), order + ) + self.assertEqual(response_format["schema"], schema) + + def test_response_format_rejects_non_json_values_without_changing_bindings(self): + for invalid in ( + {"nested": {1: "value"}}, + {"maximum": object()}, + ): + response_format = { + "type": "json_schema", + "name": "answer", + "schema": invalid, + } + with self.subTest(invalid=invalid, binding="constructor"): + with self.assertRaises(TypeError): + AIPrompt(responseFormat=response_format) + for binding in ("schema", "setter", "setParams"): + for per_row in (False, True): + with self.subTest( + invalid=invalid, binding=binding, per_row=per_row + ): + stage = AIPrompt(responseFormat={"type": "json_object"}) + if per_row: + stage.setResponseFormatCol("format") + original = stage._java_obj.responseFormat().jsonEncode( + stage._java_obj.getOrDefault( + stage._java_obj.responseFormat() + ) + ) + with self.assertRaises(TypeError): + if binding == "schema": + stage.setResponseSchema(invalid) + elif binding == "setter": + stage.setResponseFormat(response_format) + else: + stage.setParams(responseFormat=response_format) + self.assertEqual( + stage._java_obj.responseFormat().jsonEncode( + stage._java_obj.getOrDefault( + stage._java_obj.responseFormat() + ) + ), + original, + ) + + def test_response_format_metadata_survives_python_bindings_and_persistence(self): + details = { + "name": "answer", + "description": "Extract the answer", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "strict": True, + "provider_extension": {"value": 0.7}, + } + nested = {"type": "json_schema", "json_schema": details} + flat = {"type": "json_schema", **details} + reply = { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ], + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + } + data = spark.createDataFrame( + [("hello", json.dumps(flat))], "text STRING, format STRING" + ) + for api in ("chat_completions", "responses"): + for binding in ("constructor", "setter", "column"): + with self.subTest(api=api, binding=binding), endpoint(reply) as ( + url, + requests, + ), tempfile.TemporaryDirectory() as directory: + arguments = { + "apiType": api, + "url": url + "/v1", + "model": "model", + "inputCol": "text", + "outputCol": "answer", + "errorCol": "error", + } + if binding == "constructor": + arguments["responseFormat"] = nested + stage = AIPrompt(**arguments) + if binding == "setter": + stage.setResponseFormat(nested) + elif binding == "column": + stage.setResponseFormatCol("format") + saved = str(Path(directory) / "stage") + stage.copy({}).write().save(saved) + row = AIPrompt.load(saved).transform(data).head() + self.assertIsNone(row.error) + self.assertEqual(row.answer, "answer") + self.assertEqual(len(requests), 1) + body = requests[0][1] + actual = ( + body["response_format"] + if api == "chat_completions" + else body["text"]["format"] + ) + self.assertEqual( + actual, nested if api == "chat_completions" else flat + ) + + def test_response_format_string_shorthand_preserves_legacy_bindings(self): + reply = { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ], + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + } + data = spark.createDataFrame([("hello",)], "text STRING") + for api in ("chat_completions", "responses"): + for binding in ("constructor", "setter", "setParams"): + for selector in ("text", " JSON_OBJECT \n"): + with self.subTest( + api=api, binding=binding, selector=selector + ), endpoint(reply) as ( + url, + requests, + ), tempfile.TemporaryDirectory() as directory: + arguments = dict( + apiType=api, + url=url + "/v1", + model="model", + inputCol="text", + errorCol="error", + ) + if binding == "constructor": + arguments["responseFormat"] = selector + stage = AIPrompt(**arguments) + if binding == "setter": + stage.setResponseFormat(selector) + elif binding == "setParams": + stage.setParams(responseFormat=selector) + stage.setResponseFormat(" \n ") + stage.copy({}).write().save(directory + "/stage") + self.assertIsNone( + AIPrompt.load(directory + "/stage") + .transform(data) + .head() + .error + ) + self.assertEqual(len(requests), 1) + body = requests[0][1] + actual = ( + body["response_format"] + if api == "chat_completions" + else body["text"]["format"] + ) + self.assertEqual(actual, {"type": selector.strip().lower()}) + for value in ("json_schema", '{"type":"json_object"}', "unknown"): + with self.subTest(invalid=value): + with self.assertRaises(IllegalArgumentException): + AIPrompt(responseFormat=value) + for per_row in (False, True): + stage = AIPrompt(responseFormat={"type": "text"}) + if per_row: + stage.setResponseFormatCol("format") + for set_value in ( + stage.setResponseFormat, + lambda value: stage.setParams(responseFormat=value), + ): + with self.assertRaises(IllegalArgumentException): + set_value(value) + if per_row: + self.assertEqual(stage.getResponseFormatCol(), "format") + else: + self.assertEqual( + stage.getResponseFormat(), {"type": "text"} + ) + + def test_saved_explicit_empty_postprocessing_remains_disabled(self): + for stage_type in (AIPrompt, AISystemOne, OpenAIDecisions): + with self.subTest(stage=stage_type.__name__): + with tempfile.TemporaryDirectory() as directory: + stage = stage_type().setPostProcessing("") + stage.write().save(directory + "/explicit") + loaded = stage_type.load(directory + "/explicit") + for candidate in (stage, loaded, loaded.copy({})): + with self.assertRaises(IllegalArgumentException): + candidate.setPostProcessingOptions({"delimiter": ";"}) + implicit = stage_type() + implicit.write().save(directory + "/implicit") + restored = stage_type.load(directory + "/implicit") + restored.setPostProcessingOptions({"delimiter": ";"}) + self.assertEqual(restored.getPostProcessing(), "csv") + + def test_existing_prompt_and_embedding_saved_stages_remain_directly_usable(self): + cases = [ + ( + OpenAIPrompt, + { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "legacy answer", + }, + } + ] + }, + ), + (OpenAIEmbedding, {"data": [{"index": 0, "embedding": [0.25, 0.5]}]}), + ] + for stage_type, reply in cases: + with self.subTest(stage=stage_type.__name__), endpoint(reply) as ( + url, + requests, + ): + stage = ( + stage_type() + .setUrl(url + "/v1") + .setDeploymentName("existing-model") + .setOutputCol("answer") + .setErrorCol("error") + ) + if stage_type is OpenAIPrompt: + stage.setPromptTemplate("{text}") + else: + stage.setTextCol("text") + with tempfile.TemporaryDirectory() as directory: + stage.write().save(directory + "/stage") + key = "spark.synapseml.legacy.allowUnsafeJavaDeserialization" + previous = spark.conf.get(key, None) + try: + # Only load the trusted legacy stage created by this test. + spark.conf.set(key, "true") + loaded = stage_type.load(directory + "/stage") + finally: + if previous is None: + spark.conf.unset(key) + else: + spark.conf.set(key, previous) + row = loaded.transform( + spark.createDataFrame([("existing input",)], "text STRING") + ).head() + self.assertIsNone(row.error) + if stage_type is OpenAIPrompt: + self.assertEqual(row.answer, "legacy answer") + else: + self.assertEqual(row.answer.toArray().tolist(), [0.25, 0.5]) + self.assertEqual(len(requests), 1) + self.assertEqual(requests[0][1]["model"], "existing-model") + + def test_generated_constructor_setparams_copy_and_load_reach_http(self): + reply = { + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "answer"}} + ], + "future": {"retained": True}, + } + with endpoint(reply) as (url, requests): + original = AIPrompt( + url=url + "/v1", + deploymentName="test-model", + inputCol="text", + outputCol="answer", + errorCol="error", + rawResponseCol="raw", + extraBody={"provider_option": {"enabled": True}}, + ) + extension = { + "provider_option": [1, None, False, 9223372036854775808123456789] + } + original.setParams(extraBody=extension) + with tempfile.TemporaryDirectory() as directory: + original.copy({}).write().save(directory + "/stage") + loaded = AIPrompt.load(directory + "/stage") + self.assertEqual(loaded.getExtraBody(), extension) + rows = loaded.transform( + spark.createDataFrame([("hello",)], "text STRING") + ).collect() + self.assertEqual(rows[0].answer, "answer") + self.assertIsNone(rows[0].error) + self.assertTrue(json.loads(rows[0].raw)["future"]["retained"]) + self.assertEqual( + requests[0][1]["provider_option"], extension["provider_option"] + ) + + def test_native_dictionary_and_array_constructors_preserve_protocols(self): + names = ["e_one", "d_two", "c_three", "b_four", "a_five", "z_six"] + questions = { + "urgent": {"type": "noul", "instructions": "Is this urgent?"}, + **{ + name: { + "type": "choice", + "instructions": name, + "criteria": {label: label for label in names}, + } + for name in names + }, + } + cases = [ + ( + AISystemOne, + "systemone", + questions, + {"model": "returned-model", "answers": {"urgent": {"result": True}}}, + ), + ( + OpenAIDecisions, + "decisions", + [ + { + "type": "predicate", + "name": "urgent", + "instructions": "Is this urgent?", + } + ], + { + "model": "returned-model", + "answers": [ + {"type": "predicate", "name": "urgent", "probability": 0.9} + ], + }, + ), + ] + for stage_type, api, questions, reply in cases: + with self.subTest(api=api), endpoint(reply) as (url, requests): + stage = stage_type( + url=url if api == "systemone" else url + "/v1", + model="explicit-model", + inputCol="text", + questions=questions, + outputCol="answer", + errorCol="error", + ) + self.assertEqual(stage.getApiType(), api) + with tempfile.TemporaryDirectory() as directory: + stage.copy({}).write().save(directory + "/native") + loaded = stage_type.load(directory + "/native") + self.assertEqual(loaded.getApiType(), api) + row = loaded.transform( + spark.createDataFrame([("urgent outage",)], "text STRING") + ).head() + self.assertIsNone(row.error) + self.assertEqual(requests[0][1]["questions"], questions) + self.assertEqual( + json.dumps(requests[0][1]["questions"]), json.dumps(questions) + ) + if api == "systemone": + self.assertTrue(json.loads(row.answer.answers["urgent"])["result"]) + else: + self.assertEqual(row.answer.answers[0].probability, 0.9) + + def test_body_json_validation_does_not_silently_coerce_values(self): + cycle = {} + cycle["self"] = cycle + for value in ([1], {"x": float("nan")}, {1: "non-string key"}, cycle): + with self.subTest(value_type=type(value).__name__): + with self.assertRaises((TypeError, ValueError)): + AIPrompt(extraBody=value) + with self.assertRaises((TypeError, ValueError)): + AIPrompt().setParams(extraBody=value) + with self.assertRaises((TypeError, ValueError)): + AIPrompt().setExtraBody(value) + + def test_column_binding_then_scalar_clears_stale_column(self): + reply = {"data": [{"index": 0, "embedding": [0.25, 0.5]}]} + with endpoint(reply) as (url, requests): + stage = ( + AIPrompt() + .setApiType("embeddings") + .setUrl(url + "/v1") + .setDeploymentName("embedding-model") + .setInputCol("text") + .setExtraBodyCol("absent") + .setExtraBody({"encoding_format": "float"}) + .setUserCol("text") + .setOutputCol("vector") + ) + row = stage.transform( + spark.createDataFrame([("hello",)], "text STRING") + ).head() + self.assertEqual(row.vector.toArray().tolist(), [0.25, 0.5]) + self.assertEqual(requests[0][1]["user"], "hello") + self.assertEqual(requests[0][1]["encoding_format"], "float") + + +if __name__ == "__main__": + unittest.main() diff --git a/cognitive/src/test/python/synapsemltest/services/test_AIPromptLive.py b/cognitive/src/test/python/synapsemltest/services/test_AIPromptLive.py new file mode 100644 index 00000000000..199d1618bd8 --- /dev/null +++ b/cognitive/src/test/python/synapsemltest/services/test_AIPromptLive.py @@ -0,0 +1,176 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +"""Opt-in native inference with synthetic data and environment-only credentials.""" + +import json +import os +import unittest + +from pyspark.sql import SparkSession +from synapse.ml.core.init_spark import init_spark +from synapse.ml.services.ai import AIPrompt, AISystemOne +from synapse.ml.services.openai import OpenAIDecisions + + +SYSTEM_QUESTIONS = { + "repeat_contact": { + "type": "noul", + "instructions": "Has the customer contacted support before?", + }, + "team": { + "type": "choice", + "instructions": "Which team should handle this request?", + "criteria": { + "billing": "Charges, invoices and refunds", + "engineering": "Bugs, errors and outages", + }, + }, + "severity": { + "type": "score", + "instructions": "How severe is this issue?", + "criteria": ["Cosmetic", "Major", "Critical"], + }, +} +DECISION_QUESTIONS = [ + { + "type": "predicate", + "name": "repeat_contact", + "instructions": "Has the customer contacted support before?", + }, + { + "type": "choice", + "name": "team", + "instructions": "Which team should handle this request?", + "choices": [ + {"value": "billing", "description": "Charges, invoices and refunds"}, + {"value": "engineering", "description": "Bugs, errors and outages"}, + ], + }, + { + "type": "score", + "name": "severity", + "instructions": "How severe is this issue?", + "levels": [ + {"label": "Cosmetic", "description": "Appearance only"}, + {"label": "Major", "description": "A workaround exists"}, + {"label": "Critical", "description": "A task is fully blocked"}, + ], + }, +] + + +class TestAIPromptLive(unittest.TestCase): + def _run_native(self, api, direct, prefix, model_setting, questions): + def configured(name): + return os.environ.get(prefix + name, "").strip() + + endpoint = configured("_ENDPOINT") + model = configured(model_setting) + token = configured("_TOKEN") + key = configured("_KEY") + if not (endpoint and model and (token or key)): + self.skipTest( + "Live native endpoint, model and credentials are not configured" + ) + spark = SparkSession.getActiveSession() or init_spark() + dataset = spark.createDataFrame( + [ + ( + 1, + "Checkout is down again for every customer. " + "This is my third request.", + ), + (2, None), + ], + "id INT, ticket STRING", + ) + for stage in (direct(), AIPrompt().setApiType(api)): + with self.subTest(stage=type(stage).__name__): + stage.setUrl(endpoint).setModel(model).setQuestions(questions) + stage.setOutputCol("decision").setErrorCol("error") + stage.setUsageCol("usage").setRawResponseCol("raw").setConcurrency(2) + if api == "systemone": + stage.setStateCol("ticket") + else: + stage.setInputCol("ticket") + if token: + stage.setAADToken(token) + else: + stage.setSubscriptionKey(key) + result = stage.transform(dataset).cache() + try: + expected = stage._java_obj.transformSchema( + dataset._jdf.schema() + ).json() + self.assertEqual( + json.loads(result.schema.json()), json.loads(expected) + ) + rows = {row.id: row for row in result.collect()} + self.assertIsNone( + rows[1].error, "Live native inference returned an error" + ) + self.assertIsNotNone(rows[1].usage) + self.assertTrue(rows[1].decision.model) + raw = json.loads(rows[1].raw) + answers = rows[1].decision.answers + if api == "systemone": + self.assertEqual(set(answers), set(SYSTEM_QUESTIONS)) + self.assertEqual( + { + name: json.loads(value) + for name, value in answers.items() + }, + raw["answers"], + ) + else: + self.assertEqual( + [answer.name for answer in answers], + [question["name"] for question in DECISION_QUESTIONS], + ) + self.assertEqual( + [answer.type for answer in answers], + ["predicate", "choice", "score"], + ) + self.assertGreaterEqual(answers[0].probability, 0) + self.assertLessEqual(answers[0].probability, 1) + self.assertEqual(answers[1].choice.type, "string") + self.assertIn( + answers[1].choice.stringValue, ("billing", "engineering") + ) + self.assertGreaterEqual(answers[2].score, 0) + self.assertLessEqual(answers[2].score, 2) + self.assertIsNone(rows[2].decision) + self.assertIsNone(rows[2].usage) + self.assertEqual( + json.loads(rows[2].error.response)["code"], "missing_input" + ) + # Reuse the materialized result: a second action must not call the model. + pandas = result.toPandas().set_index("id") + self.assertEqual(list(pandas.columns), result.columns[1:]) + self.assertEqual(json.loads(pandas.loc[1, "raw"]), raw) + self.assertIsNone(pandas.loc[2, "decision"]) + finally: + result.unpersist() + + def test_live_systemone_pyspark_and_pandas_outputs(self): + self._run_native( + "systemone", + AISystemOne, + "SYNAPSEML_SYSTEMONE", + "_DEPLOYMENT", + SYSTEM_QUESTIONS, + ) + + def test_live_decisions_pyspark_and_pandas_outputs(self): + self._run_native( + "decisions", + OpenAIDecisions, + "SYNAPSEML_DECISIONS", + "_MODEL", + DECISION_QUESTIONS, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIJsonSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIJsonSuite.scala new file mode 100644 index 00000000000..63d64c9c989 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIJsonSuite.scala @@ -0,0 +1,786 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.codegen.CodegenConfig +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.param.AnyJsonFormat +import com.microsoft.azure.synapse.ml.services.openai.{JsonParser, OpenAIDecisions, OutputParser} +import org.apache.commons.io.FileUtils +import org.apache.spark.ml.param.{Param, ParamMap} +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.functions.lit +import org.scalatest.concurrent.TimeLimits +import org.scalatest.time.{Seconds, Span} +import spray.json.{JsArray, JsBoolean, JsNull, JsNumber, JsObject, JsString, enrichString} + +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.util.{ArrayList, LinkedHashMap} +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.collection.mutable +import scala.util.Using + +private[ai] class TestableAIDefaultJsonPrompt(override val uid: String) extends AIPrompt(uid) { + setDefault(extraBody -> Left(Map("value" -> BigDecimal("0.12345678901234567890123456789")))) +} + +private[ai] class TestableAIDefaultRawPrompt(override val uid: String) extends AIPrompt(uid) { + setDefault(rawResponseCol -> "raw") +} + +private[ai] class TestableAIDateParserPrompt(override val uid: String) extends AIPrompt(uid) { + override private[services] def getParser: OutputParser = + new JsonParser("value DATE", Map("dateFormat" -> "dd/MM/yyyy")) +} + +class AIJsonSuite extends TestBase with TimeLimits { + private val missing: Null = None.orNull + private val chatReply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + + test("ordered JSON retains insertion order through request builder map operations") { + val names = Vector("zebra", "alpha", "middle", "beta", "omega", "first") + val original = AIJson.orderedMap(names.zipWithIndex) + val updated = original.updated("middle", 20) + ("last" -> 6) + assert(updated.keys.toVector == names :+ "last") + val merged = updated ++ Seq("alpha" -> 10, "tail" -> 7) + assert(merged.keys.toVector == names ++ Vector("last", "tail")) + val removed = (merged - "zebra") -- Seq("beta", "omega") + assert(removed.keys.toVector == Vector("alpha", "middle", "first", "last", "tail")) + val filtered = removed.filter(_._2 > 5).filterNot(_._1 == "first") + assert(filtered.keys.toVector == Vector("alpha", "middle", "last", "tail")) + assert(filtered("alpha") == 10 && filtered("middle") == 20) + assert(original.keys.toVector == names) + } + + Seq("chat_completions", "responses").foreach { api => + Seq(false, true).foreach { columnBinding => + test(s"$api tool requests require explicitly projected raw responses with column binding $columnBinding") { + val reply = if (api == "chat_completions") { + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer", + |"tool_calls":[{"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{}"}}]}}]}""" + .stripMargin + } else { + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}, + |{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{}"}]}""".stripMargin + } + val tools = """{"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}""" + val data = spark.range(1).select(lit("hello").as("text"), lit(tools).as("body")) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new TestableAIDefaultRawPrompt("default-raw") + .setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (columnBinding) stage.setExtraBodyCol("body") else stage.setExtraBody(tools) + assert(stage.isDefined(stage.rawResponseCol) && !stage.isSet(stage.rawResponseCol)) + if (columnBinding) { + val rejected = stage.transform(data) + assert(!rejected.schema.fieldNames.contains("raw")) + val row = rejected.head() + assert(row.getAs[String]("answer") == null) + assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + } else { + intercept[IllegalArgumentException](stage.transform(data)) + } + assert(server.requests.isEmpty) + stage.setRawResponseCol("raw") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null) + assert(row.getAs[String]("answer") == "answer") + assert(row.getAs[String]("raw") == reply) + assert(server.requests.size() == 1) + } + } + } + } + + test("JSON postprocessing diagnostics use the same parser options as the parsed output") { + val content = JsObject("value" -> JsString("31/12/2025")).compactPrint + val reply = JsObject("choices" -> JsArray(JsObject("index" -> JsNumber(0), + "message" -> JsObject("role" -> JsString("assistant"), "content" -> JsString(content))))).compactPrint + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new TestableAIDateParserPrompt("date-parser") + .setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + .setPostProcessingOptions(Map("jsonSchema" -> "value DATE")) + val data = spark.range(1).select(lit("hello").as("text")) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[Row]("answer").getAs[java.sql.Date]("value") == java.sql.Date.valueOf("2025-12-31")) + assert(row.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + } + } + + test("JSON numeric limit errors are distinguished from malformed documents") { + val error = intercept[IllegalArgumentException](AIJson.parse("1" * 1001)) + assert(error.getMessage.contains("JSON exceeds supported parsing limits")) + assert(intercept[IllegalArgumentException](AIJson.parse("""{"broken":}""")) + .getMessage.contains("Invalid JSON")) + } + + test("container numeric setters enforce decodable numeric token boundaries") { + Seq(1000, 1001).foreach { size => + val digits = "1" * size + val numbers: Seq[java.lang.Number] = Seq(BigInt(digits), BigDecimal(digits), + new java.math.BigInteger(digits), new java.math.BigDecimal(digits)) + numbers.foreach { number => + val stage = AIPrompt.withSession(spark)(new AIPrompt()) + val scalaMap = Map[String, Any]("x" -> number) + val javaMap = new LinkedHashMap[String, Object]() + javaMap.put("x", number) + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setExtraBody(scalaMap), () => stage.setExtraBody(javaMap), + () => stage.setQuestions(scalaMap), () => stage.setQuestions(javaMap)) + setters.foreach { set => + if (size == 1001) intercept[IllegalArgumentException](set()) + else set() + } + if (size == 1000) { + assert(stage.extraBody.jsonDecode(stage.extraBody.jsonEncode(Left(stage.getExtraBody))) == + Left(stage.getExtraBody)) + assert(stage.questions.jsonDecode(stage.questions.jsonEncode(Left(stage.getQuestions))) == + Left(stage.getQuestions)) + } + } + } + } + + test("typed JSON numeric setters enforce the same parser and persistence boundaries") { + Seq("1" * 1000, "1" * 1001, "-" + "1" * 1000, "1." + "1" * 999).foreach { token => + val typed = JsObject("x" -> JsNumber(BigDecimal(token))) + val text = s"""{"x":$token}""" + val supported = scala.util.Try(AIJson.objectString(text)).isSuccess + val stage = AIPrompt.withSession(spark)(new AIPrompt()) + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setExtraBody(typed), () => stage.setExtraBody(text), + () => stage.setQuestions(typed), () => stage.setQuestions(text)) + setters.foreach { set => + if (!supported) intercept[IllegalArgumentException](set()) + else set() + } + if (supported) { + assert(stage.extraBody.jsonDecode(stage.extraBody.jsonEncode(Left(stage.getExtraBody))) == + Left(stage.getExtraBody)) + assert(stage.questions.jsonDecode(stage.questions.jsonEncode(Left(stage.getQuestions))) == + Left(stage.getQuestions)) + } + } + } + + test("large ordered JSON objects avoid quadratic per-row construction and merging") { + val entries = (0 until 50000).map(index => s"field_$index" -> s"value_$index") + val json = entries.map { case (key, value) => s""""$key":"$value"""" }.mkString("{", ",", "}") + val javaMap = new LinkedHashMap[String, String]() + entries.foreach { case (key, value) => javaMap.put(key, value) } + failAfter(Span(30, Seconds)) { + val parsed = AIJson.objectString(json) + assert(parsed.fields.keys.toVector == entries.map(_._1).toVector) + assert(AIJson.fromAny(javaMap).compactPrint == json) + val decoded = AIJson.AnyFormat.read(parsed) + assert(AIJson.fromAny(decoded).compactPrint == json) + val merged = AIExtraBody.merge(JsObject(), parsed, Set.empty, "chat_completions", lossless = true) + assert(merged.compactPrint == json) + val updated = merged.fields.updated("field_25000", JsString("replaced")) - "field_0" + assert(updated.keys.toVector == entries.map(_._1).tail.toVector) + assert(updated("field_25000") == JsString("replaced")) + } + } + + test("JSON-valued defaults preserve exact numbers and explicitness through save load") { + val data = spark.range(1).select(lit("hello").as("text")) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new TestableAIDefaultJsonPrompt("default-json-precision") + .setUrl(server.url + "/v1").setDeploymentName("explicit-model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + val expected = AIJson.fromAny(stage.getExtraBody) + val directory = Files.createTempDirectory("aiprompt-default-json-").resolve("stage").toString + try { + stage.copy(ParamMap.empty).write.session(spark).save(directory) + val loaded = AIPrompt.load(directory) + assert(!loaded.isSet(loaded.extraBody)) + assert(loaded.getDefault(loaded.extraBody) == stage.getDefault(stage.extraBody)) + assert(AIJson.fromAny(loaded.getExtraBody) == expected) + assert(stage.transformSchema(data.schema) == loaded.transformSchema(data.schema)) + assert(loaded.transform(data).head().getAs[String]("answer") == "answer") + assert(AIJson.objectString(server.requests.peek()._2).fields("value") == + expected.asJsObject.fields("value")) + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(new java.io.File(directory).getParentFile) + } + } + } + + test("extraBody preserves top-level order through setters persistence and HTTP") { + val data = spark.range(1).select(lit("hello").as("text")) + val names = Seq("e_one", "d_two", "c_three", "b_four", "a_five", "z_six") + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setDeploymentName("explicit-model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + Seq(names, names.reverse).foreach { order => + val fields = order.map(name => name -> JsString(name)) + val extension = JsObject(ListMap(fields: _*)) + val javaMap = new java.util.LinkedHashMap[String, Object]() + order.foreach(name => javaMap.put(name, name)) + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setExtraBody(extension), + () => stage.setExtraBody(extension.compactPrint), + () => stage.setExtraBody(ListMap(order.map(name => name -> (name: Any)): _*)), + () => stage.setExtraBody(javaMap)) + setters.foreach { set => + assert(set().getExtraBody.keys.toSeq == order) + } + val directory = Files.createTempDirectory("aiprompt-extension-order-").resolve("stage").toString + try { + stage.copy(ParamMap.empty).write.session(spark).save(directory) + val loaded = AIPrompt.load(directory) + assert(loaded.getExtraBody.keys.toSeq == order) + assert(stage.transformSchema(data.schema) == loaded.transformSchema(data.schema)) + assert(loaded.transform(data).head().getAs[String]("answer") == "answer") + val wireKeys = AIJson.objectString(server.requests.asScala.last._2).fields.keys.toSeq + assert(wireKeys.filter(names.contains) == order) + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(new java.io.File(directory).getParentFile) + } + } + } + } + + Seq("chat_completions", "systemone", "decisions").foreach { api => + test(s"generated R wrapper defaults preserve caller intent through $api inference") { + val stage = api match { + case "systemone" => new AISystemOne() + case "decisions" => new OpenAIDecisions() + case _ => new AIPrompt() + } + val directory = Files.createTempDirectory(tmpDir, "aiprompt-r-wrapper-").toFile + try { + val conf = CodegenConfig("test-module", None, directory.getAbsolutePath, directory.getAbsolutePath, + "1.0.0", "1.0.0", "1.0.0", "synapseml-test") + stage.makeRFile(conf) + val files = conf.rSrcDir.listFiles().filter(_.getName.endsWith(".R")) + assert(files.length == 1) + val source = FileUtils.readFileToString(files.head, StandardCharsets.UTF_8) + val arguments = source.split("\n").map(_.trim).dropWhile(!_.contains("<- function(")) + .drop(1).takeWhile(!_.startsWith("uid=")).filter(_.contains("=")).map { line => + val pair = line.stripSuffix(",").split("=", 2) + pair(0) -> pair(1) + }.toMap + // Reproduce the generated conditional setters for omitted R arguments. + stage.params.foreach { raw => + val param = raw.asInstanceOf[Param[Any]] + if (arguments(param.name) != "NULL") { + stage.set(param, stage.getDefault(param).getOrElse(fail(s"Unexpected R default for ${param.name}"))) + } + } + val response = api match { + case "systemone" => NativeRequestFixtures.SystemResponse + case "decisions" => NativeRequestFixtures.DecisionResponse + case _ => chatReply + } + Using.resource(new AIRecordingServer((_, _) => 200 -> response)) { server => + stage.setUrl(server.url + (if (api == "systemone") "" else "/v1")) + .setModel("model").setOutputCol("answer").setErrorCol("error") + api match { + case "systemone" => stage.setStateCol("text").setQuestions(NativeRequestFixtures.SystemQuestions) + case "decisions" => stage.setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions) + case _ => stage.setPromptTemplate("{text}") + } + val input = spark.range(1).select(lit("hello").as("text")) + val expectedSchema = stage.transformSchema(input.schema) + val result = stage.transform(input) + assert(result.schema == expectedSchema) + assert(result.head().getAs[Row]("error") == null) + assert(server.requests.size() == 1) + assert(!stage.isSet(stage.messagesCol) && !stage.isSet(stage.systemPrompt)) + } + } finally { + FileUtils.deleteDirectory(directory) + } + } + } + + test("native questions preserve object and array order through facade leaf persistence and column bindings") { + val data = spark.range(1).select(lit("hello").as("text")) + val names = Seq("e_one", "d_two", "c_three", "b_four", "a_five", "z_six") + val criteria = JsObject(ListMap(names.map(name => name -> JsString(name)): _*)) + val systemQuestions = JsObject(ListMap(names.map(name => name -> JsObject( + "type" -> JsString("choice"), "instructions" -> JsString(name), "criteria" -> criteria)): _*)) + val decisionQuestions = JsArray(names.map(name => JsObject("type" -> JsString("predicate"), + "name" -> JsString(name), "instructions" -> JsString(name), "extension" -> criteria)).toVector) + Seq( + (new AIPrompt().setApiType("systemone"), AIPrompt.load _, systemQuestions, NativeRequestFixtures.SystemResponse), + (new AISystemOne(), AISystemOne.load _, systemQuestions, NativeRequestFixtures.SystemResponse), + (new AIPrompt().setApiType("decisions"), AIPrompt.load _, decisionQuestions, + NativeRequestFixtures.DecisionResponse), + (new OpenAIDecisions(), OpenAIDecisions.load _, decisionQuestions, NativeRequestFixtures.DecisionResponse) + ).foreach { case (stage, load, questions, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + stage.setUrl(server.url).setModel("model").setInputCol("text").setQuestions(questions.compactPrint) + .setErrorCol("error") + val saved = Files.createTempDirectory(tmpDir, "aiprompt-native-question-order").resolve("model").toString + stage.write.save(saved) + val fromColumn = stage.copy(ParamMap.empty).setQuestionsCol("questions") + val savedColumn = Files.createTempDirectory(tmpDir, "aiprompt-native-question-column").resolve("model").toString + fromColumn.write.save(savedColumn) + Seq(stage -> data, stage.copy(ParamMap.empty) -> data, load(saved) -> data, + load(savedColumn) -> data.withColumn("questions", lit(questions.compactPrint))).foreach { + case (candidate, input) => + server.requests.clear() + assert(Option(candidate.transform(input).head().getAs[Row]("error")).isEmpty) + assert(server.requests.size() == 1) + val sent = AIJson.objectString(server.requests.peek()._2).fields("questions") + assert(sent.compactPrint == questions.compactPrint) + } + } + } + } + + Seq( + "chat_completions" -> """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""", + "responses" -> """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + ).foreach { case (api, reply) => + test(s"$api malformed per-row response schema names produce errors without failing valid rows") { + import spark.implicits._ + val invalid = Seq( + """{"properties":{},"name":null}""", + """{"type":"object","name":42}""", + """{"type":"json_schema","json_schema":{"name":null,"schema":{"type":"object"}}}""", + """{"type":"json_schema","json_schema":{"name":42,"schema":{"type":"object"}}}""") + val data = (invalid :+ """{"type":"text"}""").zipWithIndex.map { case (format, id) => + (id, "hello", format) + }.toDF("id", "text", "format") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setResponseFormatCol("format").setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect().sortBy(_.getAs[Int]("id")) + assert(rows.take(invalid.size).forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(rows.last.getAs[Row]("error") == null && rows.last.getAs[String]("answer") == "answer") + assert(server.requests.size() == 1) + } + } + + test(s"$api preserves ordered response schemas through setters persistence columns and body extensions") { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val names = Seq("e_one", "d_two", "c_three", "b_four", "a_five", "z_six") + val data = spark.range(1).select(lit("hello").as("text")) + def prompt: AIPrompt = new AIPrompt().setApiType(api).setUrl(server.url + "/v1") + .setModel("model").setInputCol("text").setErrorCol("error") + Seq(names, names.reverse).foreach { order => + val schema = Map[String, Any]("type" -> "object", + "properties" -> ListMap(order.map(_ -> Map("type" -> "string")): _*), + "required" -> names, "additionalProperties" -> false) + val stage = prompt.setResponseSchema(schema) + def verify(candidate: AIPrompt, data: DataFrame): Unit = { + server.requests.clear() + assert(Option(candidate.transform(data).head().getAs[Row]("error")).isEmpty) + assert(server.requests.size() == 1) + val body = server.requests.peek()._2 + val start = body.indexOf("\"properties\":{") + assert(start >= 0) + val positions = order.map(name => body.indexOf(s""""$name":""", start)) + assert(positions.forall(_ >= 0) && positions == positions.sorted) + } + val saved = Files.createTempDirectory(tmpDir, "aiprompt-ordered-schema").resolve("model").toString + stage.write.save(saved) + Seq("original" -> stage, "copy" -> stage.copy(ParamMap.empty), + "loaded" -> AIPrompt.load(saved)).foreach { case (label, candidate) => + withClue(s"$label: ")(verify(candidate, data)) + } + val formatJson = AnyJsonFormat.anyFormat.write(stage.getResponseFormat).compactPrint + verify(prompt.setResponseFormatCol("format"), data.withColumn("format", lit(formatJson))) + val details = Map[String, Any]("name" -> "response_schema", "strict" -> true, "schema" -> schema) + val extension = if (api == "chat_completions") { + Map("response_format" -> Map("type" -> "json_schema", "json_schema" -> details)) + } else Map("text" -> Map("format" -> (details + ("type" -> "json_schema")))) + val extended = prompt.setExtraBody(extension) + verify(extended, data) + val savedExtension = Files.createTempDirectory(tmpDir, "aiprompt-ordered-extension").resolve("model").toString + extended.copy(ParamMap.empty).write.save(savedExtension) + withClue("loaded extension: ")(verify(AIPrompt.load(savedExtension), data)) + verify(prompt.setExtraBodyCol("body"), + data.withColumn("body", lit(AnyJsonFormat.anyFormat.write(extension).compactPrint))) + } + } + } + } + + private def merge(base: String, + extension: String, + owned: Set[List[String]] = Set.empty, + api: String = "responses", + lossless: Boolean = true): JsObject = + AIExtraBody.merge(AIJson.objectString(base), AIJson.objectString(extension), owned, api, lossless) + + test("strict JSON parsing preserves booleans null arrays and exact numeric values") { + val text = """{"nested":{"items":[true,false,null,"true",1,1.125,9223372036854775808]}, + |"text":"quoted \"text\" and escaped \n newline"}""".stripMargin + assert(AIJson.parse(text) == text.parseJson) + assert(AIJson.objectString(text) == text.parseJson) + assert(AIJson.parse("null") == JsNull) + } + + test("duplicate object keys are rejected at every nesting level without disclosing payloads") { + Seq( + """{"private-key":1,"private-key":2}""", + """{"nested":{"private-key":1,"private-key":2}}""", + """{"items":[{"private-key":1,"private-key":2}]}""" + ).foreach { text => + val error = intercept[IllegalArgumentException](AIJson.parse(text)) + assert(!error.getMessage.contains("private-key")) + assert(!error.getMessage.contains(text)) + } + } + + test("malformed truncated multiple-root and nonfinite JSON is rejected safely") { + Seq( + "", " ", """{"private":"unterminated""", """{"x":""", "[1,", + """{"x":1,}""", """{"x":NaN}""", """{"x":Infinity}""", """{"x":-Infinity}""" + ).foreach { text => + val error = intercept[IllegalArgumentException](AIJson.parse(text)) + assert(!error.getMessage.contains("unterminated")) + } + intercept[IllegalArgumentException](AIJson.parse("{} {}")) + intercept[IllegalArgumentException](AIJson.parse("[] true")) + intercept[IllegalArgumentException](AIJson.parse(missing)) + } + + test("body bindings reject non-object roots while generic JSON values remain supported") { + Seq("[]", "null", "true", "1", "\"string\"").foreach { text => + intercept[IllegalArgumentException](AIJson.objectString(text)) + } + intercept[IllegalArgumentException](AIJson.obj(AIJson.fromAny(Seq("not", "an", "object")))) + assert(AIJson.fromAny(Seq("array", true, missing)) == JsArray(JsString("array"), JsBoolean(true), JsNull)) + } + + test("generic object bindings reject null before changing configuration") { + val stage = AIPrompt.withSession(spark)(new AIPrompt()) + Seq(stage.extraBody, stage.responseFormat).foreach { param => + val original = Map[String, Any]("type" -> "json_object") + stage.setScalarParam(param, original) + val invalid: Seq[Either[Map[String, Any], String]] = Seq(missing, Left(missing), Right(missing)) + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setScalarParam(param, missing), () => stage.setVectorParam(param, missing)) ++ + invalid.flatMap(value => Seq( + () => stage.set(param, value), () => stage.copy(ParamMap(param -> value)))) + setters.foreach { set => + intercept[IllegalArgumentException](set()) + assert(stage.getScalarParam(param) == original) + } + stage.setVectorParam(param, "") + assert(stage.getVectorParam(param) == "") + } + } + + test("generic object bindings preserve nested null and optional null cells through save load and HTTP") { + val data = spark.range(1).select(lit("hello").as("text"), lit(missing).cast("string").as("body")) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq(false, true).foreach { columnBinding => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (columnBinding) { + stage.setVectorParam(stage.extraBody, "body").setVectorParam(stage.responseFormat, "body") + } else { + stage.setScalarParam(stage.extraBody, Map[String, Any]("metadata" -> Map("optional" -> missing))) + stage.setScalarParam(stage.responseFormat, Map[String, Any]("type" -> "json_object")) + } + val saved = Files.createTempDirectory(tmpDir, "aiprompt-null-bindings").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val loaded = AIPrompt.load(saved) + assert(loaded.get(loaded.extraBody) == stage.get(stage.extraBody)) + assert(loaded.get(loaded.responseFormat) == stage.get(stage.responseFormat)) + val result = loaded.transform(data) + assert(result.schema == loaded.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer") + val body = AIJson.objectString(server.requests.asScala.last._2) + if (columnBinding) { + assert(!body.fields.contains("metadata") && !body.fields.contains("response_format")) + } else { + assert(body.fields("metadata") == JsObject("optional" -> JsNull)) + assert(body.fields("response_format") == JsObject("type" -> JsString("json_object"))) + } + } + assert(server.requests.size() == 2) + } + } + + test("Java and Scala containers convert without coercing typed values or rejecting shared children") { + val shared = mutable.ArrayBuffer[Any]("value", true, missing) + val scalaValues = Map[String, Any]( + "left" -> shared, "right" -> shared, "integer" -> BigInt("9223372036854775808"), + "decimal" -> BigDecimal("0.1234567890123456789"), "array" -> Array(1, 2)) + val javaValues = new LinkedHashMap[String, Any]() + val javaList = new ArrayList[Any]() + javaList.add("value") + javaList.add(true) + javaList.add(missing) + javaValues.put("left", javaList) + javaValues.put("right", javaList) + javaValues.put("integer", new java.math.BigInteger("9223372036854775808")) + javaValues.put("decimal", new java.math.BigDecimal("0.1234567890123456789")) + javaValues.put("array", Array(1, 2)) + val expected = """{"left":["value",true,null],"right":["value",true,null], + |"integer":9223372036854775808,"decimal":0.1234567890123456789,"array":[1,2]}""".stripMargin.parseJson + assert(AIJson.fromAny(scalaValues) == expected) + assert(AIJson.fromAny(javaValues) == expected) + } + + test("nonfinite Scala and Java numeric members are rejected before body serialization") { + val numbers: Seq[Any] = Seq(Double.NaN, Double.PositiveInfinity, Double.NegativeInfinity, + Float.NaN, Float.PositiveInfinity, Float.NegativeInfinity, + java.lang.Double.valueOf(Double.NaN), java.lang.Float.valueOf(Float.PositiveInfinity)) + numbers.foreach { number => + intercept[IllegalArgumentException](AIJson.fromAny(Map("value" -> number))) + val javaValues = new ArrayList[Any]() + javaValues.add(number) + intercept[IllegalArgumentException](AIJson.fromAny(javaValues)) + } + } + + test("finite Float values retain their decimal representation through setters persistence and HTTP") { + val values = Seq(0.7f, 0.1f, -0.7f, Float.MinPositiveValue, Float.MaxValue) + val scalaValues = values.zipWithIndex.map { case (value, index) => s"v$index" -> value }.toMap + val javaValues = new LinkedHashMap[String, Object]() + scalaValues.foreach { case (key, value) => javaValues.put(key, java.lang.Float.valueOf(value)) } + val expected = JsObject(scalaValues.map { case (key, value) => + key -> JsNumber(BigDecimal(java.lang.Float.toString(value))) + }) + val input = spark.range(1).select(lit("hello").as("text")) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setExtraBody(scalaValues), () => stage.setExtraBody(javaValues)) + setters.foreach { set => + set() + val saved = Files.createTempDirectory(tmpDir, "aiprompt-floats").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val loaded = AIPrompt.load(saved) + assert(AIJson.fromAny(loaded.getExtraBody) == expected) + assert(loaded.transform(input).head().getAs[Row]("error") == null) + val sent = AIJson.objectString(server.requests.asScala.last._2) + assert(JsObject(sent.fields.filter { case (key, _) => scalaValues.contains(key) }) == expected) + } + assert(server.requests.size() == 2) + val questionSetters: Seq[() => AIPrompt] = Seq( + () => stage.setQuestions(scalaValues), () => stage.setQuestions(javaValues)) + questionSetters.foreach { set => + val candidate = set() + assert(AIJson.fromAny(candidate.getQuestions) == expected) + assert(AIJson.fromAny(candidate.questions.jsonDecode( + candidate.questions.jsonEncode(Left(candidate.getQuestions))).left.get) == expected) + } + } + } + + test("non-string object keys and unsupported member types are rejected") { + intercept[IllegalArgumentException](AIJson.fromAny(Map(1 -> "value"))) + val javaValues = new LinkedHashMap[Any, Any]() + javaValues.put(true, "value") + intercept[IllegalArgumentException](AIJson.fromAny(javaValues)) + val error = intercept[IllegalArgumentException](AIJson.fromAny(Map("value" -> new Object()))) + assert(error.getMessage == "Unsupported JSON value type") + } + + test("Scala map sequence and array cycles are rejected without recursion overflow") { + val map = mutable.Map.empty[String, Any] + map.update("self", map) + val sequence = mutable.ArrayBuffer.empty[Any] + sequence += sequence + val array = new Array[AnyRef](1) + array(0) = array + Seq(map, sequence, array).foreach { cyclic => + val error = intercept[IllegalArgumentException](AIJson.fromAny(cyclic)) + assert(error.getMessage.contains("Cyclic")) + } + } + + test("Java map list and mixed Java Scala cycles are rejected") { + val map = new LinkedHashMap[String, Any]() + map.put("self", map) + val list = new ArrayList[Any]() + list.add(list) + val mixed = new ArrayList[Any]() + mixed.add(Map("back" -> mixed)) + Seq(map, list, mixed).foreach { cyclic => + val error = intercept[IllegalArgumentException](AIJson.fromAny(cyclic)) + assert(error.getMessage.contains("Cyclic")) + } + } + + test("excessive JSON and container nesting is rejected with a controlled error") { + val depth = 110 + val text = ("[" * depth) + "0" + ("]" * depth) + intercept[IllegalArgumentException](AIJson.parse(text)) + val nested = (1 to depth).foldLeft[Any]("leaf") { case (value, _) => Seq(value) } + intercept[IllegalArgumentException](AIJson.fromAny(nested)) + } + + test("maximum JSON payload depth survives public object and question parameter persistence") { + val text = ("{\"level\":" * 100) + "\"leaf\"" + ("}" * 100) + val expected = AIJson.parse(text) + val root = Files.createTempDirectory("aiprompt-json-depth") + try { + val stage = new AIPrompt().setExtraBody(text).setQuestions(text) + stage.write.session(spark).save(root.resolve("model").toString) + val loaded = AIPrompt.read.session(spark).load(root.resolve("model").toString) + assert(AIJson.fromAny(loaded.getExtraBody) == expected) + assert(AIJson.fromAny(loaded.getQuestions) == expected) + val array = ("[" * 100) + "\"leaf\"" + ("]" * 100) + stage.setQuestions(array) + assert(AIJson.fromAny(stage.questions.jsonDecode(stage.questions.jsonEncode(Left(stage.getQuestions))) + .left.get) == AIJson.parse(array)) + } finally org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + + test("JSON payload depth rejects the first unsupported scalar leaf for every setter representation") { + val text = ("{\"level\":" * 101) + "\"leaf\"" + ("}" * 101) + val nested = (1 to 101).foldLeft[Any]("leaf") { case (value, _) => Map("level" -> value) } + val json = (1 to 101).foldLeft[spray.json.JsValue](JsString("leaf")) { case (value, _) => + JsObject("level" -> value) + } + intercept[IllegalArgumentException](AIJson.parse(text)) + intercept[IllegalArgumentException](AIJson.fromAny(nested)) + intercept[IllegalArgumentException](new AIPrompt().setExtraBody(json.asJsObject)) + intercept[IllegalArgumentException](new AIPrompt().setQuestions(json)) + val stage = new AIPrompt() + intercept[IllegalArgumentException](stage.extraBody.jsonDecode("{\"left\":" + text + "}")) + intercept[IllegalArgumentException](stage.questions.jsonDecode("{\"left\":" + text + "}")) + } + + test("recursive disjoint members preserve explicit typed subtrees and unknown provider options") { + val base = """{"model":"deployment","text":{"format":{"type":"json_schema","name":"answer"}}, + |"metadata":{"retained":"yes"}}""".stripMargin + val extension = """{"text":{"verbosity":"low"},"metadata":{"new":"value"}, + |"future_provider_field":{"array":[1,true,{"value":null}]}}""".stripMargin + val expected = """{"model":"deployment","text":{"format":{"type":"json_schema","name":"answer"}, + |"verbosity":"low"},"metadata":{"retained":"yes","new":"value"}, + |"future_provider_field":{"array":[1,true,{"value":null}]}}""".stripMargin.parseJson + assert(merge(base, extension, Set(List("text", "format"))) == expected) + } + + test("implicit defaults can be overridden and null arrays and scalars replace atomically") { + val base = """{"temperature":1,"metadata":{"old":1},"stop":["old","retained"], + |"text":{"verbosity":"medium"}}""".stripMargin + val extension = """{"temperature":0.3,"metadata":null,"stop":["new"],"text":"replacement"}""" + assert(merge(base, extension) == extension.parseJson) + assert(merge("""{"metadata":{"old":1}}""", """{"metadata":[]}""") == """{"metadata":[]}""".parseJson) + assert(merge("""{"metadata":null}""", """{"metadata":{"new":1}}""") == + """{"metadata":{"new":1}}""".parseJson) + } + + test("equal explicit values and modifications below an owned subtree are conflicts") { + intercept[IllegalArgumentException] { + merge("""{"temperature":0.3}""", """{"temperature":0.3}""", Set(List("temperature"))) + } + val error = intercept[IllegalArgumentException] { + merge("""{"text":{"format":{"type":"json_schema"}}}""", + """{"text":{"format":{"name":"private-value"}}}""", Set(List("text", "format"))) + } + assert(error.getMessage.contains("text.format")) + assert(!error.getMessage.contains("private-value")) + intercept[IllegalArgumentException] { + merge("""{"text":{"verbosity":"medium"}}""", """{"text":{}}""", Set(List("text"))) + } + } + + test("atomic ancestor replacements cannot overwrite explicitly owned descendants") { + Seq("null", "[]", "\"replacement\"", "3", "true").foreach { replacement => + intercept[IllegalArgumentException] { + merge("""{"text":{"format":{"type":"json_schema"}}}""", s"""{"text":$replacement}""", + Set(List("text", "format", "type"))) + } + } + assert(merge("""{"text":{"format":{"type":"json_schema"}}}""", """{"text":{}}""", + Set(List("text", "format", "type"))) == """{"text":{"format":{"type":"json_schema"}}}""".parseJson) + } + + test("ownership checks also apply when an extension introduces a missing parent object") { + intercept[IllegalArgumentException] { + merge("{}", """{"text":{"format":{"type":"json_schema"}}}""", Set(List("text", "format"))) + } + assert(merge("{}", """{"text":{"verbosity":"low"}}""", Set(List("text", "format"))) == + """{"text":{"verbosity":"low"}}""".parseJson) + } + + test("explicit token ownership conflicts with every semantic alias including equal values") { + val aliases = Seq("max_tokens", "max_completion_tokens", "max_output_tokens") + aliases.foreach { typed => + aliases.foreach { extra => + intercept[IllegalArgumentException] { + merge(s"""{"$typed":12}""", s"""{"$extra":12}""", Set(List(typed))) + } + } + } + } + + test("raw token aliases replace implicit alternate defaults rather than leaving contradictory fields") { + val aliases = Seq("max_tokens", "max_completion_tokens", "max_output_tokens") + aliases.foreach { default => + aliases.foreach { extra => + assert(merge(s"""{"$default":12,"temperature":0.2}""", s"""{"$extra":6}""") == + s"""{"$extra":6,"temperature":0.2}""".parseJson) + } + } + intercept[IllegalArgumentException](merge("{}", """{"max_tokens":12,"max_completion_tokens":12}""")) + intercept[IllegalArgumentException](merge("""{"max_tokens":12,"max_output_tokens":12}""", "{}")) + } + + test("identity routing authentication transport and storage fields are reserved in body extensions") { + val reserved = Seq("model", "messages", "input", "state", "questions", "url", "endpoint", "base_url", + "api_version", "api_key", "subscription_key", "subscriptionKey", "AADToken", "CustomAuthHeader", + "customHeaders", "headers", "extra_headers", "extra_query", "query", "timeout", "max_retries", + "retry_policy", "outputPath", "outputBatchSize", "output_path", "output_batch_size", "customUrlRoot") + reserved.foreach { key => + val error = intercept[IllegalArgumentException](merge("{}", s"""{"$key":"private-value"}""")) + assert(error.getMessage.contains(key)) + assert(!error.getMessage.contains("private-value")) + } + assert(merge("{}", """{"provider_options":{"model":"nested-value"}}""") == + """{"provider_options":{"model":"nested-value"}}""".parseJson) + } + + test("streaming and background lifecycle protections validate the fully merged body") { + Seq("stream", "background").foreach { key => + Seq("true", "null", "\"true\"", "1", "{}").foreach { value => + intercept[IllegalArgumentException](merge("{}", s"""{"$key":$value}""")) + } + intercept[IllegalArgumentException](merge(s"""{"$key":true}""", "{}")) + assert(merge("{}", s"""{"$key":false}""") == s"""{"$key":false}""".parseJson) + } + } + + test("tool request fields require raw response exposure whether supplied by the base or extension") { + Seq("tools", "tool_choice", "functions", "function_call").foreach { key => + val request = s"""{"$key":[]}""" + intercept[IllegalArgumentException](merge("{}", request, lossless = false)) + intercept[IllegalArgumentException](merge(request, "{}", lossless = false)) + assert(merge("{}", request, lossless = true) == request.parseJson) + } + assert(merge("{}", """{"metadata":{"tools":"ordinary metadata"}}""", lossless = false) == + """{"metadata":{"tools":"ordinary metadata"}}""".parseJson) + } + + test("embedding encoding must remain float after merge and is not silently stripped") { + Seq("\"base64\"", "\"future\"", "null", "true", "1").foreach { encoding => + intercept[IllegalArgumentException](merge("{}", s"""{"encoding_format":$encoding}""", api = "embeddings")) + } + intercept[IllegalArgumentException](merge("""{"encoding_format":"base64"}""", "{}", api = "embeddings")) + assert(merge("{}", """{"encoding_format":"float"}""", api = "embeddings") == + """{"encoding_format":"float"}""".parseJson) + assert(merge("""{"encoding_format":"float"}""", "{}", api = "embeddings") == + """{"encoding_format":"float"}""".parseJson) + assert(merge("{}", "{}", api = "embeddings") == JsObject()) + assert(merge("{}", """{"encoding_format":"base64"}""", api = "responses") == + """{"encoding_format":"base64"}""".parseJson) + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptAttachmentSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptAttachmentSuite.scala new file mode 100644 index 00000000000..0ef512d118e --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptAttachmentSuite.scala @@ -0,0 +1,415 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.io.http.HTTPTransformer +import com.microsoft.azure.synapse.ml.services.openai.OpenAIPrompt +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{FSDataInputStream, FSInputStream, FileSystem, Path, RawLocalFileSystem} +import org.apache.spark.SparkException +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.ml.util.Identifiable +import org.apache.spark.sql.{DataFrame, Dataset, Row} +import org.apache.spark.sql.types.{StringType, StructType} +import spray.json._ + +import java.io.File +import java.net.URI +import java.util.concurrent.{CancellationException, ConcurrentLinkedQueue} +import java.util.concurrent.atomic.AtomicInteger +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptCancelledAttachmentFileSystem extends RawLocalFileSystem { + override def getUri: URI = URI.create("aipromptcancel:///") + + override def open(path: Path, bufferSize: Int): FSDataInputStream = + if (path.getName == "without-message.png") throw new CancellationException() + else throw new CancellationException("private cancellation details") +} + +class TestableAIPromptAttachmentDefaults(override val uid: String) extends AIPrompt(uid) { + def this() = this(Identifiable.randomUID("TestableAIPromptAttachmentDefaults")) + def defaultLimit(value: Double): this.type = setDefault(fileSizeLimitMB -> value) +} + +class TestableLegacyAttachmentDefaults(override val uid: String) extends OpenAIPrompt(uid) { + def this() = this(Identifiable.randomUID("TestableLegacyAttachmentDefaults")) + setDefault(fileSizeLimitMB -> 0.000001) +} + +class AIPromptBoundedAttachmentFileSystem extends RawLocalFileSystem { + val bytesRead = new AtomicInteger() + val closed = new AtomicInteger() + override def getUri: URI = URI.create("aipromptbounded:///") + + override def open(path: Path, bufferSize: Int): FSDataInputStream = new FSDataInputStream(new FSInputStream { + private val size = path.getName.stripSuffix(".png").toInt + private var position = 0 + override def getPos: Long = position + override def seek(value: Long): Unit = throw new UnsupportedOperationException() + override def seekToNewSource(value: Long): Boolean = false + override def read(): Int = if (position == size) -1 else { + position += 1 + bytesRead.incrementAndGet() + 1 + } + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + if (length == 0) 0 else if (position == size) -1 else { + val count = math.min(length, size - position) + java.util.Arrays.fill(buffer, offset, offset + count, 1.toByte) + position += count + bytesRead.addAndGet(count) + count + } + } + override def close(): Unit = { + closed.incrementAndGet() + super.close() + } + }) +} + +object AIPromptSessionAttachmentFileSystem { + val Instances = new ConcurrentLinkedQueue[AIPromptSessionAttachmentFileSystem]() +} + +class AIPromptSessionAttachmentFileSystem extends AIPromptBoundedAttachmentFileSystem { + override def getUri: URI = URI.create("aipromptsession:///") + override def initialize(uri: URI, conf: Configuration): Unit = { + require(conf.get("aiprompt.test.attachment.marker") == "session-only", "Missing session filesystem configuration") + super.initialize(uri, conf) + AIPromptSessionAttachmentFileSystem.Instances.add(this) + } +} + +class AIPromptAttachmentSuite extends AIPromptTestBase { + Seq("chat_completions", "responses").foreach { api => + test(s"$api attachment size limits reject invalid Param bindings before HTTP") { + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + val data = spark.createDataFrame(Seq(Row("data:text/plain;base64,aGVsbG8=")).asJava, + new StructType().add("asset", StringType)) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val base = new TestableAIPromptAttachmentDefaults() + .setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")).setErrorCol("error") + for { + limit <- Seq(0.0, -1.0, Double.NaN) + binding <- Seq("direct", "copy", "default") + } { + val stage = binding match { + case "copy" => base.copy(ParamMap(base.fileSizeLimitMB -> limit)) + case "default" => base.defaultLimit(limit).copy(ParamMap.empty) + case "direct" => + val direct = base.copy(ParamMap.empty) + direct.set(direct.fileSizeLimitMB, limit) + } + assert(intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + .getMessage.contains("fileSizeLimitMB must be positive")) + assert(intercept[IllegalArgumentException](stage.transform(data)) + .getMessage.contains("fileSizeLimitMB must be positive")) + assert(server.requests.isEmpty) + } + assert(base.setFileSizeLimitMB(0.001).transform(data).head().getAs[Row]("error") == null) + val oversized = base.setFileSizeLimitMB(0.000001).transform(data).head().getAs[Row]("error") + assert(AIJson.objectString(oversized.getString(0)).fields("code") == JsString("invalid_request")) + assert(server.requests.size() == 1) + } + } + + test(s"$api applies default attachment limits without making them explicit") { + import spark.implicits._ + val data = Seq("data:text/plain;base64,aGVsbG8=").toDF("asset") + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new TestableAIPromptAttachmentDefaults().defaultLimit(0.000001) + .setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")).setErrorCol("error") + val error = stage.transform(data).head().getAs[Row]("error") + assert(error != null && AIJson.objectString(error.getString(0)).fields("code") == JsString("invalid_request")) + assert(!stage.isSet(stage.fileSizeLimitMB) && server.requests.isEmpty) + val explicit = stage.copy(ParamMap(stage.fileSizeLimitMB -> 0.001)) + assert(explicit.transform(data).head().getAs[Row]("error") == null) + assert(stage.defaultLimit(Double.PositiveInfinity).transform(data).head().getAs[Row]("error") == null) + val legacy = new TestableLegacyAttachmentDefaults().setApiType(api).setUrl(server.url + "/v1") + .setDeploymentName("model").setSubscriptionKey("unused") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")).setErrorCol("error") + assert(legacy.transform(data).head().getAs[Row]("error") == null && server.requests.size() == 3) + } + } + + test(s"$api bounds filesystem attachment reads to the configured bytes plus one sentinel") { + import spark.implicits._ + val location = URI.create("aipromptbounded:///") + val conf = new Configuration() + conf.setClass("fs.aipromptbounded.impl", classOf[AIPromptBoundedAttachmentFileSystem], classOf[FileSystem]) + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + Using.resource(FileSystem.get(location, conf)) { fs => + val recorded = fs.asInstanceOf[AIPromptBoundedAttachmentFileSystem] + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(false, true).foreach { default => + val stage = new TestableAIPromptAttachmentDefaults() + .setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")).setErrorCol("error") + val limit = 16.0 / (1024 * 1024) + if (default) stage.defaultLimit(limit) else stage.setFileSizeLimitMB(limit) + Seq(32, 16).foreach { size => + recorded.bytesRead.set(0) + recorded.closed.set(0) + val error = stage.transform(Seq(s"${location}$size.png").toDF("asset")).head().getAs[Row]("error") + assert(recorded.bytesRead.get() == math.min(size, 17) && recorded.closed.get() == 1) + if (size == 32) { + assert(error != null) + val response = error.getString(0) + assert(AIJson.objectString(response).fields("code") == JsString("invalid_request")) + assert(!response.contains(location.toString) && !response.contains("32.png")) + assert(response.contains("16-byte limit") && !response.contains("0.00MB")) + } else assert(error == null) + } + } + assert(server.requests.size() == 2) + } + } + } + + Seq(None, Some(16.0 / (1024 * 1024)), Some(Double.PositiveInfinity)).foreach { limit => + test(s"$api reads cold filesystems from session configuration with limit=$limit") { + val session = spark.newSession() + import session.implicits._ + session.conf.set("fs.aipromptsession.impl", classOf[AIPromptSessionAttachmentFileSystem].getName) + session.conf.set("aiprompt.test.attachment.marker", "session-only") + val location = s"aipromptsession://${java.util.UUID.randomUUID()}/" + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + try { + assert(AIPromptSessionAttachmentFileSystem.Instances.isEmpty) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")).setErrorCol("error") + limit.foreach(stage.setFileSizeLimitMB) + Seq(16, 32).foreach { size => + AIPromptSessionAttachmentFileSystem.Instances.asScala.foreach { fs => + fs.bytesRead.set(0) + fs.closed.set(0) + } + val result = stage.transform(Seq(s"${location}$size.png").toDF("asset")) + assert(AIPromptSessionAttachmentFileSystem.Instances.asScala.forall(_.bytesRead.get() == 0)) + session.conf.set("aiprompt.test.attachment.marker", "changed-after-transform") + val error = try result.head().getAs[Row]("error") finally { + session.conf.set("aiprompt.test.attachment.marker", "session-only") + } + assert(AIPromptSessionAttachmentFileSystem.Instances.size() == 1) + val fs = AIPromptSessionAttachmentFileSystem.Instances.peek() + val overflow = limit.exists(value => !value.isInfinity) && size == 32 + assert(fs.bytesRead.get() == (if (overflow) 17 else size) && fs.closed.get() == 1) + if (overflow) { + assert(error != null && AIJson.objectString(error.getString(0)).fields("code") == + JsString("invalid_request")) + } else assert(error == null) + } + assert(server.requests.size() == (if (limit.exists(value => !value.isInfinity)) 1 else 2)) + } + } finally { + AIPromptSessionAttachmentFileSystem.Instances.asScala.foreach(_.close()) + AIPromptSessionAttachmentFileSystem.Instances.clear() + } + } + } + + test(s"$api buffers generated attachment messages only when retained in output") { + val image = "data:image/png;base64,AQID" + val data = spark.createDataFrame(Seq(Row(image)).asJava, new StructType().add("asset", StringType)) + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(false, true).foreach { drop => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error").setDropPrompt(drop) + val http = new HTTPTransformer() { + override private[ml] def transformCancellable(dataset: Dataset[_]): DataFrame = { + val generated = dataset.select(AIColumnPlan.column(s"${stage.uid}_generated_messages")).head() + assert(generated.isNullAt(0) == drop) + if (!drop) assert(generated.getSeq[Row](0).exists(_.toString.contains(image))) + super.transformCancellable(dataset) + } + } + val result = stage.infer(data, http) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + if (!drop) assert(row.getAs[Seq[Row]](stage.getMessagesCol).exists(_.toString.contains(image))) + else assert(!result.columns.contains(stage.getMessagesCol)) + } + val bodies = server.requests.asScala.map(_._2).toVector + assert(bodies.size == 2 && bodies.distinct.size == 1 && bodies.head.contains(image)) + } + } + + test(s"$api rejects resolver-equivalent columnTypes keys before HTTP") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + val schema = new StructType().add("asset", StringType) + val data = spark.createDataFrame(Seq(Row("data:image/png;base64,AQID")).asJava, schema) + try { + spark.conf.set("spark.sql.caseSensitive", "false") + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val base = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset}").setOutputCol("answer").setErrorCol("error") + Seq("path", "text").foreach { kind => + val stage = base.copy(ParamMap.empty).setColumnTypes(Map("asset" -> kind, "ASSET" -> "path")) + Seq[() => Any](() => stage.transformSchema(schema), () => stage.transform(data)).foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("Ambiguous columnTypes keys")) + assert(server.requests.isEmpty) + } + } + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + + test(s"$api preserves distinct case-sensitive columnTypes keys and attachment order") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + val first = "data:image/png;base64,AQID" + val second = "data:image/png;base64,BAUG" + val schema = new StructType().add("asset", StringType).add("ASSET", StringType) + val data = spark.createDataFrame(Seq(Row(first, second)).asJava, schema) + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + try { + spark.conf.set("spark.sql.caseSensitive", "true") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset} and {ASSET}") + .setColumnTypes(Map("ASSET" -> "path", "asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(schema)) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + val key = if (api == "chat_completions") "messages" else "input" + val imageType = if (api == "chat_completions") "image_url" else "input_image" + val content = AIJson.objectString(server.requests.peek()._2).fields(key).asInstanceOf[JsArray] + .elements.last.asJsObject.fields("content").asInstanceOf[JsArray].elements + val images = content.map(_.asJsObject.fields).filter(_("type") == JsString(imageType)) + .map(fields => if (api == "chat_completions") fields("image_url").asJsObject.fields("url") + else fields("image_url")) + assert(images == Seq(JsString(first), JsString(second))) + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + } + + Seq("chat_completions", "responses").foreach { api => + Seq("private-attachment.png", "without-message.png").foreach { filename => + test(s"legacy attachment CancellationException in $api for $filename remains sanitized") { + import spark.implicits._ + val location = URI.create(s"aipromptcancel:///$filename") + val conf = new Configuration() + conf.setClass("fs.aipromptcancel.impl", classOf[AIPromptCancelledAttachmentFileSystem], classOf[FileSystem]) + Using.resource(FileSystem.get(location, conf)) { _ => + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val stage = new OpenAIPrompt().setApiType(api).setUrl(server.url + "/v1") + .setDeploymentName("model").setSubscriptionKey("unused") + .setPromptTemplate("Review {asset}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val row = stage.transform(Seq(location.toString).toDF("asset")).head() + val error = row.getAs[Row]("error") + assert(error != null) + assert(error.getAs[String]("response") == s"Unable to read attachment '$filename'") + assert(row.getAs[String]("answer") == null && server.requests.isEmpty) + } + } + } + } + + test(s"$api propagates attachment CancellationException instead of committing a row error") { + import spark.implicits._ + val location = URI.create("aipromptcancel:///private-attachment.png") + val conf = new Configuration() + conf.setClass("fs.aipromptcancel.impl", classOf[AIPromptCancelledAttachmentFileSystem], classOf[FileSystem]) + Using.resource(FileSystem.get(location, conf)) { _ => + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val failure = intercept[SparkException] { + stage.transform(Seq(location.toString).toDF("asset")).collect() + } + val causes = Iterator.iterate[Throwable](failure)(_.getCause).takeWhile(_ != null).toVector + assert(causes.exists(_.isInstanceOf[CancellationException])) + assert(!causes.exists(error => Option(error.getMessage).exists(_.contains("private cancellation details")))) + val cancelled = causes.collectFirst { case error: CancellationException => error }.get + assert(cancelled.getMessage == "Attachment read canceled" && cancelled.getCause == null) + assert(server.requests.isEmpty) + } + } + } + + test(s"$api distinguishes attachment IO from invalid content without leaking paths or sending requests") { + import spark.implicits._ + val missing = new File(path(), "private-attachment.png").toURI.toString + val data = Seq("missing" -> missing, "invalid" -> "data:image/png;base64,@@@", + "wrong_fs" -> "file://invalid-host/private-attachment.png").toDF("kind", "asset") + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val rows = stage.transform(data).collect().map(row => row.getAs[String]("kind") -> row).toMap + Seq("missing", "wrong_fs").foreach { kind => + val readError = rows(kind).getAs[Row]("error").getString(0) + assert(AIJson.objectString(readError).fields("code") == JsString("input_io_error")) + assert(!readError.contains(missing) && !readError.contains("private-attachment")) + } + assert(AIJson.objectString(rows("invalid").getAs[Row]("error").getString(0)) + .fields("code") == JsString("invalid_request")) + assert(rows.values.forall(_.getAs[String]("answer") == null) && server.requests.isEmpty) + } + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"missing attachment placeholders preserve usable $api inputs") { + val inline = "data:image/png;base64,AQID" + val absent = Option.empty[String].orNull + val rows = Seq(absent, "", " \t ").flatMap { missing => + Seq(Row("hello", inline, missing), Row("hello", missing, missing)) + } :+ Row(absent, inline, absent) + val schema = new StructType().add("text", StringType).add("asset", StringType) + .add("optional", StringType) + val data = spark.createDataFrame(rows.asJava, schema) + val reply = if (api == "chat_completions") + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + else """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Review {asset}; optional={optional}; text={text}") + .setColumnTypes(Map("asset" -> "path", "optional" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data).collect() + assert(result.length == 7 && result.count(_.getAs[String]("answer") == "answer") == 3) + val errors = result.flatMap(row => Option(row.getAs[Row]("error"))) + assert(errors.length == 4 && errors.forall(error => + AIJson.objectString(error.getString(0)).fields("code") == JsString("missing_input"))) + assert(server.requests.size() == 3) + server.requests.asScala.foreach { case (_, body) => + val key = if (api == "chat_completions") "messages" else "input" + val content = AIJson.objectString(body).fields(key).asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsArray].elements + assert(content.exists(_.asJsObject.fields.get("text").contains( + JsString("Review attachment.png; optional=; text=hello")))) + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptBenchmarkSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptBenchmarkSuite.scala new file mode 100644 index 00000000000..1d1290503cf --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptBenchmarkSuite.scala @@ -0,0 +1,50 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.openai.OpenAIPrompt +import org.apache.spark.sql.functions.{col, expr} + +import scala.util.Using + +class AIPromptBenchmarkSuite extends TestBase { + test("compare public distributed chat paths at fixed payload, latency, partitions and concurrency") { + assume(sys.env.get("SYNAPSEML_AI_BENCHMARK").contains("true"), "Opt-in bounded local benchmark") + val count = 4096 + val reply = """{"id":"bench","choices":[{"index":0,"message":{"role":"assistant","content":"ok"}, + |"finish_reason":"stop"}]}""".stripMargin + val data = spark.range(count).withColumn("text", expr("repeat('x', 128 + pmod(id, 896))")).cache() + assert(data.count() == count) + try { + Using.resource(new AIRecordingServer((_, _) => { + Thread.sleep(5) + 200 -> reply + })) { server => + def configured(stage: OpenAIPrompt, concurrency: Int): OpenAIPrompt = + stage.setUrl(server.url + "/v1").setDeploymentName("local-benchmark") + .setPromptTemplate("{text}").setOutputCol("answer").setErrorCol("error") + .setConcurrency(concurrency) + Seq(new OpenAIPrompt(), new AIPrompt()).foreach { stage => + configured(stage, 1).transform(data.limit(64).coalesce(1)).select("answer").collect() + } + Seq(4 -> 1, 16 -> 4).foreach { case (partitions, concurrency) => + Seq(new OpenAIPrompt(), new AIPrompt()).foreach { stage => + val before = server.requests.size() + val started = System.nanoTime() + val rows = configured(stage, concurrency).transform(data.repartition(partitions)) + .select(col("answer"), col("error")).collect() + val seconds = (System.nanoTime() - started).toDouble / 1e9 + val sends = server.requests.size() - before + assert(rows.length == count && sends == count) + assert(rows.forall(row => row.getString(0) == "ok" && row.isNullAt(1))) + info(s"AI_BENCHMARK stage=${stage.getClass.getSimpleName} rows=$count " + + s"partitions=$partitions concurrency=$concurrency providerLatencyMs=5 sends=$sends " + + s"seconds=$seconds rowsPerSecond=${count / seconds}") + } + } + } + } finally data.unpersist() + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptCancellationSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptCancellationSuite.scala new file mode 100644 index 00000000000..194b5551efe --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptCancellationSuite.scala @@ -0,0 +1,122 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.injections.TaskContextUtils +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.{col, udf} +import org.apache.spark.sql.types.StructType + +import java.util.concurrent.{Callable, ConcurrentHashMap, CountDownLatch, ExecutionException, Executors, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import scala.collection.JavaConverters._ +import scala.util.Using + +private[ai] object AICancellationGates { + class Gate { + val entered = new CountDownLatch(1) + val release = new CountDownLatch(1) + val completed = new CountDownLatch(1) + val context = new AtomicReference[TaskContext]() + val failure = new AtomicReference[Throwable]() + val completionInterrupted = new AtomicBoolean() + } + + val Active = new ConcurrentHashMap[String, Gate]() + + def block(id: String, answer: String): String = { + val gate = Active.get(id) + val context = TaskContext.get() + gate.context.set(context) + context.addTaskFailureListener((_: TaskContext, error: Throwable) => gate.failure.set(error)) + context.addTaskCompletionListener[Unit] { _ => + gate.completionInterrupted.set(Thread.currentThread().isInterrupted) + gate.completed.countDown() + } + gate.entered.countDown() + require(gate.release.await(30, TimeUnit.SECONDS), "Downstream cancellation test gate timed out") + answer + } +} + +class AIPromptCancellationSuite extends TestBase { + Seq(false, true).foreach { activeRequest => + test(s"Spark cancellation drains workers with blocked downstream and active request=$activeRequest") { + val received = new CountDownLatch(1) + val releaseResponse = new CountDownLatch(1) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, body) => { + if (body.contains("retry")) { + received.countDown() + if (activeRequest) { + assert(releaseResponse.await(30, TimeUnit.SECONDS)) + 200 -> reply + } else 429 -> "{}" + } else 200 -> reply + }, retryAfter = "5")) { server => + val data = spark.createDataFrame(Seq(Row("ready"), Row("retry")).asJava, + new StructType().add("text", "string")).coalesce(1) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setConcurrency(2).setTimeout(30).setOutputCol("answer") + val context = spark.sparkContext + val group = java.util.UUID.randomUUID().toString + val gate = new AICancellationGates.Gate() + AICancellationGates.Active.put(group, gate) + val block = udf((answer: String) => AICancellationGates.block(group, answer)) + val result = stage.transform(data).select(block(col("answer"))) + val executor = Executors.newSingleThreadExecutor() + val job = executor.submit(new Callable[Unit] { + override def call(): Unit = { + context.setJobGroup(group, "AIPrompt blocked downstream cancellation", interruptOnCancel = false) + try result.collect() finally context.clearJobGroup() + } + }) + def workers: Iterable[Thread] = Thread.getAllStackTraces.keySet().asScala + .filter(thread => thread.isAlive && thread.getName.startsWith("synapseml-bounded-http-")) + try { + assert(gate.entered.await(15, TimeUnit.SECONDS), "The first response did not block downstream") + assert(received.await(15, TimeUnit.SECONDS), "The concurrent request never reached the server") + if (!activeRequest) { + val sleeping = System.nanoTime() + TimeUnit.SECONDS.toNanos(3) + def inBackoff: Boolean = workers.exists(_.getStackTrace.exists(frame => + frame.getClassName == "java.lang.Thread" && frame.getMethodName == "sleep")) + Iterator.continually(inBackoff) + .takeWhile(ready => !ready && System.nanoTime() < sleeping).foreach(_ => Thread.sleep(10)) + assert(inBackoff, "The retry worker never entered backoff") + } + context.cancelJobGroup(group) + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10) + Iterator.continually(gate.context.get().isInterrupted()) + .takeWhile(interrupted => !interrupted && System.nanoTime() < deadline).foreach(_ => Thread.sleep(10)) + assert(gate.context.get().isInterrupted(), "Spark did not deliver task cancellation") + val drained = System.nanoTime() + TimeUnit.SECONDS.toNanos(3) + Iterator.continually(workers.nonEmpty) + .takeWhile(running => running && System.nanoTime() < drained).foreach(_ => Thread.sleep(10)) + assert(workers.isEmpty, "Canceled HTTP workers outlived cancellation while downstream was blocked") + assert(!gate.context.get().isCompleted(), "Task completion incorrectly supplied the cancellation signal") + Thread.sleep(6000) + assert(server.requests.size() == 2, "A worker retried after task cancellation") + gate.release.countDown() + intercept[ExecutionException](job.get(15, TimeUnit.SECONDS)) + assert(gate.completed.await(15, TimeUnit.SECONDS)) + assert(!gate.completionInterrupted.get(), "Task-only cancellation interrupted executor cleanup") + val killed = Iterator.iterate(gate.failure.get())(_.getCause).takeWhile(_ != null) + .collectFirst { case error: TaskKilledException => error } + val reason = TaskContextUtils.getKillReason(gate.context.get()) + assert(reason.nonEmpty && killed.exists(error => reason.contains(error.reason))) + } finally { + releaseResponse.countDown() + gate.release.countDown() + context.cancelJobGroup(group) + job.cancel(true) + executor.shutdownNow() + assert(executor.awaitTermination(15, TimeUnit.SECONDS)) + AICancellationGates.Active.remove(group) + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptEndpointSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptEndpointSuite.scala new file mode 100644 index 00000000000..43b6991581e --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptEndpointSuite.scala @@ -0,0 +1,657 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.HTTPSchema +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.OpenAIApiVersion +import com.microsoft.azure.synapse.ml.services.aifoundry.AIFoundryModel +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIDecisions, OpenAIDeploymentNameKey, + OpenAIEmbeddingDeploymentNameKey, OpenAITemperatureKey, OpenAITopPKey} +import org.apache.spark.sql.{Encoders, Row} +import org.apache.spark.sql.types.{StringType, StructType} +import spray.json._ + +import java.net.URLDecoder +import java.nio.charset.StandardCharsets +import java.util.Locale +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptEndpointSuite extends TestBase { + private val inputSchema = new StructType().add("text", StringType) + private val cases = Seq( + ("chat_completions", "chat/completions", + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}"""), + ("responses", "responses", + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}"""), + ("embeddings", "embeddings", """{"data":[{"index":0,"embedding":[0.25,0.5]}]}"""), + ("decisions", "decisions", NativeRequestFixtures.DecisionResponse)) + + cases.filter(entry => Set("chat_completions", "embeddings")(entry._1)).foreach { case (api, _, reply) => + test(s"deployment URL Unicode rejects malformed scalar and default $api bindings before HTTP") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val keys = Seq(OpenAIDeploymentNameKey, OpenAIEmbeddingDeploymentNameKey, AIFoundryModel) + val previous = keys.map(key => key -> GlobalParams.getGlobalParam(key)) + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val bindings = Seq("deployment", "model", "default-deployment", "default-model", + "global-deployment", "global-model") ++ (if (api == "embeddings") Seq("global-embedding") else Seq.empty) + for (malformed <- Seq("private-\uD800", "private-\uDC00"); binding <- bindings) { + keys.foreach(GlobalParams.resetGlobalParam) + val stage = new TestableAIMutableJsonDefaultPrompt("deployment-unicode") + .setApiType(api).setUrl(server.url).setInputCol("text") + binding match { + case "deployment" => stage.setDeploymentName(malformed) + case "model" => stage.setModel(malformed) + case "default-deployment" => stage.replaceJsonDefault(stage.deploymentName, Left(malformed)) + case "default-model" => stage.replaceJsonDefault(stage.model, Left(malformed)) + case "global-deployment" => GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, Left(malformed)) + case "global-model" => GlobalParams.setGlobalParam(AIFoundryModel, Left(malformed)) + case "global-embedding" => GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, Left(malformed)) + } + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { invoke => + val error = intercept[IllegalArgumentException](invoke()) + assert(error.getMessage.contains("surrogate") && !error.getMessage.contains("private-")) + assert(error.getMessage.contains(if (binding.endsWith("model")) "model" else "deploymentName")) + assert(error.getCause == null) + } + assert(server.requests.isEmpty) + } + } + } finally { + previous.foreach { + case (key, Some(value)) => GlobalParams.setGlobalParam(key, value) + case (key, None) => GlobalParams.resetGlobalParam(key) + } + } + } + } + + test("deployment URL Unicode keeps column failures per row and preserves valid segment bytes") { + val unicode = "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD+?/" + val models = Seq("\uD800", "\uDC00", unicode, "\\" + "ud800") + val encoded = Set("caf%C3%A9%20%E9%9B%AA%20%F0%9F%98%80%20%EF%BF%BD%2B%3F%2F", "%5Cud800") + val schema = new StructType().add("text", StringType).add("target", StringType) + val input = spark.range(models.size).map(index => Row("hello", models(index.toInt)))(Encoders.row(schema)) + for ((api, route, reply) <- cases.filter(entry => Set("chat_completions", "embeddings")(entry._1)); + deployment <- Seq(false, true)) { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url).setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + if (deployment) stage.setDeploymentNameCol("target") else stage.setModelCol("target") + stage.transformSchema(schema) + val rows = stage.transform(input).collect() + val errors = rows.filter(_.getAs[Row]("error") != null) + assert(rows.length == 4 && errors.length == 2 && server.requests.size() == 2) + assert(errors.forall(row => row.getAs[Any]("answer") == null && + row.getAs[Row]("error").getString(0).contains("invalid_request") && + row.getAs[Row]("error").getString(0).contains("surrogate") && + row.getAs[Row]("error").getString(0).contains(if (deployment) "deploymentName" else "model"))) + val addresses = server.requests.asScala.map(_._1.split("\\?api-version=", 2).head).toSet + assert(addresses == encoded.map(value => s"/openai/deployments/$value/$route")) + val bodyModels = server.requests.asScala.map(entry => AIJson.objectString(entry._2).fields("model")).toSet + assert(bodyModels == models.drop(2).map(JsString(_)).toSet) + } + } + } + + test("deployment URL Unicode does not preflight model names used only in request bodies") { + val data = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val protocols = cases :+ ("systemone", "providers/microsoft/v1/systemone", NativeRequestFixtures.SystemResponse) + for ((api, _, reply) <- protocols; custom <- Seq(false, true); malformed <- Seq("\uD800", "\uDC00")) { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setModel(malformed).setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + if (custom) stage.setCustomUrlRoot(server.url + "/complete") + else stage.setUrl(if (api == "systemone") server.url else server.url + "/v1") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + else if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + stage.transformSchema(inputSchema) + val row = stage.transform(data).head() + assert(row.getAs[Any]("answer") == null && server.requests.isEmpty) + assert(row.getAs[Row]("error").getString(0).contains("invalid_request") && + row.getAs[Row]("error").getString(0).contains("surrogate")) + } + } + } + + test("deployment URL Unicode ignores model defaults overridden by the selected binding") { + val data = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + for ((api, _, reply) <- cases.filter(entry => Set("chat_completions", "embeddings")(entry._1)); + deployment <- Seq(false, true)) { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new TestableAIMutableJsonDefaultPrompt("ignored-deployment") + .setApiType(api).setUrl(server.url).setInputCol("text").setErrorCol("error") + if (deployment) { + stage.setDeploymentName("selected") + stage.replaceJsonDefault(stage.model, Left("\uD800")) + } else { + stage.setModel("selected") + stage.replaceJsonDefault(stage.deploymentName, Left("\uDC00")) + } + stage.transformSchema(inputSchema) + assert(stage.transform(data).head().getAs[Row]("error") == null && server.requests.size() == 1) + assert(server.requests.peek()._1.startsWith("/openai/deployments/selected/")) + } + } + } + + for (custom <- Seq(false, true); component <- Seq("path", "query")) { + test(s"endpoint Unicode rejects malformed $component with custom=$custom before HTTP") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val stages = AIRequestAdapter.SupportedApis.map(new AIPrompt().setApiType(_)) ++ + Seq(new AISystemOne(), new OpenAIDecisions()) + for (stage <- stages; malformed <- Seq("\uD800", "\uDC00")) { + val endpoint = server.url + (if (component == "path") s"/private-$malformed/v1" + else s"/v1?private=$malformed") + stage.setModel("model").setInputCol("text") + if (custom) stage.setCustomUrlRoot(endpoint) else stage.setUrl(endpoint) + if (AIRequestAdapter.NativeApis(stage.getApiType)) { + stage.setQuestions(if (stage.getApiType == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions) + } + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { invoke => + val error = intercept[IllegalArgumentException](invoke()) + assert(error.getMessage.contains("surrogate") && !error.getMessage.contains("private")) + assert(error.getMessage.contains(if (custom) "customUrlRoot" else "url")) + assert(error.getCause == null) + } + assert(server.requests.isEmpty) + } + } + } + } + + test("endpoint Unicode preserves valid path and query bytes and ignores unused URL defaults") { + val unicode = "caf\u00e9-\u96ea-\uD83D\uDE00-\uFFFD" + val encoded = "caf%C3%A9-%E9%9B%AA-%F0%9F%98%80-%EF%BF%BD" + val data = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val protocols = cases :+ ("systemone", "providers/microsoft/v1/systemone", NativeRequestFixtures.SystemResponse) + for ((api, route, reply) <- protocols; custom <- Seq(false, true)) { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val suffix = "/" + (if (api == "systemone") route else "v1/" + route) + val address = server.url + "/" + unicode + suffix + (if (custom) "?marker=" + unicode else "") + val stage = new AIPrompt().setApiType(api).setModel("model").setInputCol("text").setErrorCol("error") + if (custom) stage.setUrl(server.url + "/unused-\uD800").setCustomUrlRoot(address) else stage.setUrl(address) + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + else if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + assert(stage.transformSchema(inputSchema) == stage.transform(data).schema) + assert(stage.transform(data).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1) + assert(server.requests.peek()._1 == "/" + encoded + suffix + (if (custom) "?marker=" + encoded else "")) + } + } + } + + test("dated API versions reject malformed scalar and default Unicode before HTTP") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + try { + cases.filterNot(_._1 == "decisions").foreach { case (api, _, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + for (malformed <- Seq("private-\uD800", "private-\uDC00"); + binding <- Seq("explicit", "default", "global")) { + GlobalParams.resetGlobalParam(OpenAIApiVersion) + val stage = new TestableAIMutableJsonDefaultPrompt("version-unicode").setApiType(api) + .setUrl(server.url).setDeploymentName("model").setInputCol("text") + binding match { + case "explicit" => stage.setApiVersion(malformed) + case "default" => stage.replaceJsonDefault(stage.apiVersion, Left(malformed)) + case "global" => GlobalParams.setGlobalParam(OpenAIApiVersion, Left(malformed)) + } + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { invoke => + val error = intercept[IllegalArgumentException](invoke()) + assert(error.getMessage.contains("surrogate") && !error.getMessage.contains("private-")) + assert(error.getMessage.contains("apiVersion")) + assert(error.getCause == null) + } + assert(server.requests.isEmpty) + } + } + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("dated API version columns reject malformed Unicode and preserve valid URL encoding") { + val unicode = "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD &+?/" + val versions = Seq("\uD800", "\uDC00", unicode, "\\" + "ud800") + val schema = new StructType().add("text", StringType).add("version", StringType) + val data = spark.range(versions.size).map(index => Row("hello", versions(index.toInt)))(Encoders.row(schema)) + val reply = """{"data":[{"index":0,"embedding":[0.25,0.5]}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType("embeddings").setUrl(server.url).setDeploymentName("model") + .setInputCol("text").setApiVersionCol("version").setOutputCol("answer").setErrorCol("error") + assert(stage.transformSchema(schema).fieldNames.contains("answer")) + val rows = stage.transform(data).collect() + val errors = rows.filter(_.getAs[Row]("error") != null) + assert(rows.length == 4 && errors.length == 2 && server.requests.size() == 2) + assert(errors.forall(row => row.getAs[Row]("error").getString(0).contains("invalid_request") && + row.getAs[Row]("error").getString(0).contains("surrogate") && + row.getAs[Row]("error").getString(0).contains("apiVersion") && row.getAs[Any]("answer") == null)) + val queries = server.requests.asScala.map(_._1.split("\\?api-version=", 2).last).toSet + assert(queries == Set("caf%C3%A9+%E9%9B%AA+%F0%9F%98%80+%EF%BF%BD+%26%2B%3F%2F", "%5Cud800")) + assert(queries.map(value => URLDecoder.decode(value, "UTF-8")) == versions.drop(2).toSet) + } + } + + test("non-dated endpoints continue ignoring malformed API version defaults") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val protocols = cases :+ ("systemone", "providers/microsoft/v1/systemone", NativeRequestFixtures.SystemResponse) + protocols.foreach { case (api, _, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new TestableAIMutableJsonDefaultPrompt("ignored-version").setApiType(api) + .setUrl(if (api == "systemone") server.url else server.url + "/v1") + .setModel("model").setInputCol("text").setErrorCol("error") + stage.replaceJsonDefault(stage.apiVersion, Left("private-\uD800")) + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + else if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && !server.requests.peek()._1.contains("api-version")) + } + } + } + + Seq(false, true).foreach { explicit => + test(s"Responses forwards sampling options for GPT-5-named deployments with explicit=$explicit") { + val previousTemperature = GlobalParams.getGlobalParam(OpenAITemperatureKey) + val previousTopP = GlobalParams.getGlobalParam(OpenAITopPKey) + val reply = + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + try { + GlobalParams.setGlobalParam(OpenAITemperatureKey, Left(0.25)) + GlobalParams.setGlobalParam(OpenAITopPKey, Left(0.75)) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1") + .setDeploymentName("gpt-5-deployment").setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (explicit) stage.setTemperature(0.5).setTopP(0.9) + val data = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val schema = stage.transformSchema(data.schema) + assert(server.requests.isEmpty) + val result = stage.transform(data) + assert(result.schema == schema && server.requests.isEmpty) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + val body = AIJson.objectString(server.requests.peek()._2).fields + assert(body("model") == JsString("gpt-5-deployment")) + assert(body("temperature") == JsNumber(if (explicit) 0.5 else 0.25)) + assert(body("top_p") == JsNumber(if (explicit) 0.9 else 0.75)) + } + } finally { + previousTemperature match { + case Some(value) => GlobalParams.setGlobalParam(OpenAITemperatureKey, value) + case None => GlobalParams.resetGlobalParam(OpenAITemperatureKey) + } + previousTopP match { + case Some(value) => GlobalParams.setGlobalParam(OpenAITopPKey, value) + case None => GlobalParams.resetGlobalParam(OpenAITopPKey) + } + } + } + } + + test("Chat extraBody can select max_tokens for a dated route when typed token options are unset") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url).setDeploymentName("model").setApiVersion("2023-05-15") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setExtraBody(Map("max_tokens" -> 32)) + val row = stage.transform(input).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + val (url, body) = server.requests.peek() + assert(url == "/openai/deployments/model/chat/completions?api-version=2023-05-15") + val fields = AIJson.objectString(body).fields + assert(fields("max_tokens") == JsNumber(32)) + assert(!fields.contains("max_completion_tokens") && !fields.contains("max_output_tokens")) + } + } + + test("invalid endpoint ports fail preflight") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + Seq(0, 65536, Int.MaxValue).foreach { port => + val endpoint = s"http://127.0.0.1:$port/v1" + Seq(false, true).foreach { custom => + val stage = new AIPrompt().setModel("model").setInputCol("text") + if (custom) stage.setCustomUrlRoot(endpoint) else stage.setUrl(endpoint) + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("port") && !error.getMessage.contains(endpoint)) + } + } + } + Seq("", ":1", ":65535").foreach { port => + val stage = new AIPrompt().setUrl(s"http://127.0.0.1$port/v1").setModel("model").setInputCol("text") + stage.transformSchema(inputSchema) + } + } + + test("underscore endpoint hosts report the unsupported host before execution") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val stages = AIRequestAdapter.SupportedApis.map(new AIPrompt().setApiType(_)) ++ + Seq(new AISystemOne(), new OpenAIDecisions()) + stages.foreach { stage => + val endpoint = "http://llm_gateway.invalid:8000/v1" + stage.setUrl(endpoint).setModel("model").setInputCol("text") + if (AIRequestAdapter.NativeApis(stage.getApiType)) { + stage.setQuestions(if (stage.getApiType == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions) + } + Seq(false, true).foreach { custom => + if (custom) stage.setCustomUrlRoot(endpoint) + val calls = Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)) + calls.foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("host") && error.getMessage.contains("underscores")) + assert(!error.getMessage.contains("credentials") && !error.getMessage.contains(endpoint)) + } + } + } + } + + test("native endpoint guards recognize absolute DNS names and regional OpenAI hosts") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val endpoints = Seq( + "systemone" -> "https://API.OPENAI.COM.", + "systemone" -> "https://eu.api.openai.com", + "systemone" -> "https://EU.API.OPENAI.COM.", + "decisions" -> "https://resource.openai.azure.com./v1", + "decisions" -> "https://resource.services.ai.azure.com./v1") + endpoints.foreach { case (api, endpoint) => + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model").setInputCol("text") + .setQuestions(if (api == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions) + Seq(false, true).foreach { custom => + if (custom) stage.setCustomUrlRoot(endpoint) + val error = intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + assert(error.getMessage.contains("endpoint") && error.getCause == null) + intercept[IllegalArgumentException](stage.transform(input)) + } + } + } + + Seq("azure.com", "azure.us", "azure.cn").foreach { cloud => + test(s"Decisions rejects $cloud endpoints before execution") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + for { + prefix <- Seq("resource.openai", "resource.services.ai", "resource.cognitiveservices") + host <- Seq(s"$prefix.$cloud", s"$prefix.$cloud".toUpperCase(Locale.ROOT) + ".") + route <- Seq("", "/v1", "/v1/decisions") + stage <- Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions()) + custom <- Seq(false, true) + } { + stage.setModel("model").setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions) + val endpoint = s"https://$host$route" + stage.clear(stage.customUrlRoot) + if (custom) stage.setCustomUrlRoot(endpoint) else stage.setUrl(endpoint) + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("not a Foundry endpoint") && error.getCause == null) + assert(!error.getMessage.contains(host)) + } + } + } + } + + test("Decisions Azure host guards preserve explicit non-Azure proxy endpoints") { + Seq("gateway.azure.us.example.invalid", "gateway.azure.cn.example.invalid", "notazure.us").foreach { host => + val stage = new AIPrompt().setApiType("decisions").setUrl(s"https://$host/v1") + .setModel("model").setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions) + stage.transformSchema(inputSchema) + stage.setCustomUrlRoot(s"https://$host/proxy?version=custom").transformSchema(inputSchema) + } + } + + cases.foreach { case (api, route, reply) => + test(s"$api normalizes bare OpenAI roots without changing stored URLs or applying dated defaults") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val previousVersion = GlobalParams.getGlobalParam(OpenAIApiVersion) + try { + GlobalParams.setGlobalParam(OpenAIApiVersion, Right("unavailableVersion")) + Seq("https://api.openai.com", "https://eu.api.openai.com/", "https://US.API.OPENAI.COM.:443///") + .foreach { endpoint => + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setSubscriptionKey("synthetic-key") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val schema = stage.transformSchema(inputSchema) + val expected = endpoint.replaceAll("/+$", "") + s"/v1/$route" + stage.setHandler { (_, request) => + assert(request.requestLine.uri == expected) + val body = AIJson.objectString(new String(request.entity.get.content, StandardCharsets.UTF_8)) + assert(body.fields("model") == JsString("model")) + val headers = request.headers.map(h => h.name.toLowerCase(Locale.ROOT) -> h.value).toMap + assert(headers.get("authorization").contains("Bearer " + "synthetic-key") && + !headers.contains("api-key")) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + val result = stage.transform(input) + assert(result.schema == schema && stage.getUrl == endpoint) + val row = result.head() + assert(row.getAs[Row]("error") == null && row.getAs[Any]("answer") != null) + stage.setApiVersion("explicit-version") + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + stage.setApiVersionCol("text") + intercept[IllegalArgumentException](stage.transform(input)) + } + } finally { + previousVersion match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + } + + test("OpenAI root normalization leaves custom URL roots and other hosts unchanged") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + Seq("api.openai.com", "EU.API.OPENAI.COM.", "proxy.example.invalid", + "api.openai.com.example.invalid", "notapi.openai.com").foreach { host => + val custom = host == "api.openai.com" || host == "EU.API.OPENAI.COM." + val endpoint = s"https://$host/" + val stage = new AIPrompt().setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + if (custom) stage.setCustomUrlRoot(endpoint) + else stage.setUrl(endpoint).setApiVersion("explicit-version") + val expected = if (custom) endpoint else + endpoint + "openai/deployments/model/chat/completions?api-version=explicit-version" + stage.setHandler { (_, request) => + assert(request.requestLine.uri == expected) + HTTPSchema.stringToResponse( + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""", 200, "OK") + } + val row = stage.transform(input).head() + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer") + } + } + + cases.filterNot(_._1 == "decisions").foreach { case (api, _, reply) => + test(s"$api uses bearer keys for regional and absolute DNS OpenAI hosts only") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val hosts = Seq("api.openai.com." -> true, "EU.API.OPENAI.COM." -> true, + "us.api.openai.com" -> true, "api.openai.com.example.invalid" -> false, + "notapi.openai.com" -> false, "resource.openai.azure.com." -> false) + hosts.foreach { case (host, bearer) => + val stage = new AIPrompt().setApiType(api).setUrl(s"https://$host/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setSubscriptionKey("synthetic-key") + stage.setHandler { (_, request) => + val headers = request.headers.map(h => h.name.toLowerCase(Locale.ROOT) -> h.value).toMap + assert(headers.get("authorization") == (if (bearer) Some("Bearer synthetic-key") else None)) + assert(headers.get("api-key") == (if (bearer) None else Some("synthetic-key"))) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + val row = stage.transform(input).head() + assert(row.getAs[Row]("error") == null && row.getAs[Any]("answer") != null) + } + } + } + + test("complete Azure deployment routes named v1 retain dated API versions") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + try { + GlobalParams.setGlobalParam(OpenAIApiVersion, Left("global-version")) + cases.filter(entry => Set("chat_completions", "embeddings")(entry._1)).foreach { case (api, route, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq("v1", "V1").foreach { deployment => + val suffix = s"/openai/deployments/$deployment/$route" + val stage = new AIPrompt().setApiType(api).setUrl(server.url + suffix) + .setDeploymentName(deployment).setInputCol("text").setOutputCol("answer").setErrorCol("error") + Seq(None, Some("explicit-version"), Some("column-version")).foreach { version => + version.foreach { + case "column-version" => stage.setApiVersionCol("text") + case value => stage.setApiVersion(value) + } + assert(stage.transform(input).head().getAs[Row]("error") == null) + val expected = version.map { + case "column-version" => "hello" + case value => value + }.getOrElse("global-version") + assert(server.requests.asScala.last._1 == suffix + s"?api-version=$expected") + } + } + assert(server.requests.size() == 6) + } + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("incomplete Azure deployment bases reject before HTTP instead of matching operation names") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + Seq("embeddings", "responses", "v1", "model").foreach { deployment => + Seq("chat_completions", "responses", "embeddings").foreach { api => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + s"/openai/deployments/$deployment") + .setDeploymentName(deployment).setInputCol("text") + val error = intercept[IllegalArgumentException](stage.transform(input)) + assert(error.getMessage.contains("resource endpoint or a complete API route")) + } + } + assert(server.requests.isEmpty) + } + } + + test("v1 base routing preserves legacy case detection and normalizes redundant trailing slashes") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + cases.foreach { case (api, route, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq("/v1", "/openai/V1").foreach { base => + Seq("", "/", "///").foreach { trailing => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + base + trailing) + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val row = stage.transform(input).head() + assert(row.getAs[Row]("error") == null && row.getAs[Any]("answer") != null) + val (address, body) = server.requests.asScala.last + assert(address == s"$base/$route") + assert(AIJson.objectString(body).fields("model") == JsString("model")) + stage.setApiVersion("unused") + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + } + } + assert(server.requests.size() == 6) + } + } + } + + test("complete v1 overrides keep path and query bytes without adding routes or versions") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + cases.foreach { case (api, route, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val suffix = s"/openai/V1/$route///?value=A%2Fb%2F&case=A+B" + val stage = new AIPrompt().setApiType(api).setCustomUrlRoot(server.url + suffix) + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.requests.peek()._1 == suffix) + } + } + } + + test("complete mixed-case v1 URLs preserve their destination and reject dated versions") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + cases.foreach { case (api, route, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(route, route.toUpperCase(Locale.ROOT)).foreach { operation => + val suffix = s"/openai/V1/$operation" + val stage = new AIPrompt().setApiType(api).setUrl(server.url + suffix + "///") + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.asScala.last._1 == suffix) + stage.setApiVersion("unused") + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + stage.setApiVersionCol("text") + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + } + assert(server.requests.size() == 2) + } + } + } + + test("route classification does not decode encoded path separators") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + cases.filterNot(_._1 == "decisions").foreach { case (api, route, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val base = "/proxy%2Fv1%2F" + route.replace("/", "%2F") + val stage = new AIPrompt().setApiType(api).setUrl(server.url + base) + .setModel("Model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + assert(stage.transform(input).head().getAs[Row]("error") == null) + val suffix = if (api == "responses") "/openai/responses" + else s"/openai/deployments/Model/$route" + assert(server.requests.size() == 1) + assert(server.requests.peek()._1 == base + suffix + "?api-version=2025-04-01-preview") + } + } + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val decisions = new AIPrompt().setApiType("decisions").setUrl(server.url + "/proxy%2Fv1%2Fdecisions") + .setModel("model").setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions) + val systemone = new AISystemOne().setUrl(server.url + "/providers%2Fmicrosoft%2Fv1%2Fsystemone") + .setModel("model").setInputCol("text").setQuestions(NativeRequestFixtures.SystemQuestions) + Seq(decisions, systemone).foreach(stage => + intercept[IllegalArgumentException](stage.transform(input))) + assert(server.requests.isEmpty) + } + } + + test("mixed-case complete routes still validate deployment identity and API selection") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + Using.resource(new AIRecordingServer((_, _) => 200 -> cases.head._3)) { server => + val suffix = "/OPENAI/DEPLOYMENTS/Model/CHAT/COMPLETIONS" + val stage = new AIPrompt().setUrl(server.url + suffix).setDeploymentName("Model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.peek()._1 == suffix + "?api-version=2025-04-01-preview") + stage.setDeploymentName("model") + intercept[IllegalArgumentException](stage.transform(input)) + stage.setDeploymentName("Model").setApiType("embeddings") + intercept[IllegalArgumentException](stage.transform(input)) + stage.setApiType("chat_completions").setUrl(server.url + "/MODELS/CHAT/COMPLETIONS") + intercept[IllegalArgumentException](stage.transform(input)) + assert(server.requests.size() == 1) + } + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + val suffix = "/PROVIDERS/MICROSOFT/V1/SYSTEMONE" + val stage = new AIPrompt().setApiType("systemone").setUrl(server.url + suffix) + .setModel("model").setInputCol("text") + .setQuestions(NativeRequestFixtures.SystemQuestions).setOutputCol("answer").setErrorCol("error") + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.requests.peek()._1 == suffix) + stage.setApiType("decisions").setQuestions(NativeRequestFixtures.DecisionQuestions) + intercept[IllegalArgumentException](stage.transform(input)) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricProbe.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricProbe.scala new file mode 100644 index 00000000000..86b38b8b20b --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricProbe.scala @@ -0,0 +1,80 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HTTPSchema, URLKey} +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.OpenAISubscriptionKey +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.types.StructType +import spray.json.{JsArray, JsString} + +import java.util.Collections + +/** Test-jar bridge for managed-runtime routing without external inference or storage. */ +object AIPromptFabricProbe { + def verifyImplicitEndpoints(spark: SparkSession): String = verifyImplicitEndpoints(spark, () => new AIPrompt()) + + private[ai] def verifyImplicitEndpoints(spark: SparkSession, + create: () => AIPrompt): String = AIPrompt.withSession(spark) { + withoutGlobalDefaults { + val data = spark.createDataFrame(Collections.singletonList(Row("synthetic-model")), + new StructType().add("text", "string")) + val apis = Vector("chat_completions", "responses", "embeddings") + apis.foreach { api => + val reply = api match { + case "chat_completions" => + """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + case "responses" => + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}]}""" + case "embeddings" => """{"data":[{"index":0,"embedding":[0.5]}]}""" + } + val stage = create().setApiType(api).setModelCol("text").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + val effective = stage.effectiveCopy() + val root = effective.getUrl.stripSuffix("/") + "/" + require(effective.usingImplicitFabricEndpoint(root), "Expected an implicit Fabric route") + stage.setHandler { (_, request) => + require(request.requestLine.uri.startsWith(root), "Request did not use the implicit Fabric route") + val headers = request.headers.map(h => h.name.toLowerCase(java.util.Locale.ROOT) -> h.value).toMap + require(headers.get("authorization").exists(_.trim.nonEmpty), "Expected runtime authorization") + require(request.headers.exists(_.name == HTTPRequestData.FabricAuthMarkerHeader), + "Expected Fabric fallback authentication") + require(headers.get("x-taxonomy-traffictype").contains("Background") && + headers.get("x-llm-service-tier").contains("flex"), "Expected Fabric workload classification") + HTTPSchema.stringToResponse(reply, 200, "OK") + } + val output = stage.transform(data) + require(output.schema == stage.transformSchema(data.schema), "Implicit route schema changed") + val row = output.head() + require(row.getAs[Row]("error") == null && !row.isNullAt(row.fieldIndex("answer")), + "Implicit route did not preserve the synthetic answer") + require(!stage.isDefined(stage.url), "Execution must not persist the implicit URL") + } + require(GlobalParams.getGlobalParam(URLKey).isEmpty, "Execution must not change the global URL") + require(GlobalParams.getGlobalParam(OpenAISubscriptionKey).isEmpty, + "Execution must not change the global subscription key") + JsArray(apis.map(JsString(_))).compactPrint + } + } + + private def withoutGlobalDefaults[T](action: => T): T = { + val previousUrl = GlobalParams.getGlobalParam(URLKey) + val previousKey = GlobalParams.getGlobalParam(OpenAISubscriptionKey) + try { + GlobalParams.resetGlobalParam(URLKey) + GlobalParams.resetGlobalParam(OpenAISubscriptionKey) + action + } finally { + previousUrl match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + previousKey match { + case Some(value) => GlobalParams.setGlobalParam(OpenAISubscriptionKey, value) + case None => GlobalParams.resetGlobalParam(OpenAISubscriptionKey) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricSuite.scala new file mode 100644 index 00000000000..5822b776d73 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricSuite.scala @@ -0,0 +1,525 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HTTPSchema, HandlingUtils, URLKey} +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.{OpenAIApiVersion, OpenAISubscriptionKey} +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{MapType, StringType, StructType} +import spray.json._ + +import java.nio.file.Files +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.util.Using + +class TestableAIHeaderPrompt(override val uid: String) extends AIPrompt(uid) { + override protected val subscriptionKeyHeaderName: String = "x-model-key" + override protected val aadHeaderName: String = "x-model-token" +} + +class TestableAISharedHeaderPrompt(override val uid: String) extends AIPrompt(uid) { + override protected val subscriptionKeyHeaderName: String = "AUTHORIZATION" +} + +class AIPromptFabricSuite extends TestBase { + private val root = "https://fabric.example.invalid/cognitive/openai" + private val chat = """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + private val responses = """{"output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"ok"}]}]}""".stripMargin + private val embeddings = """{"data":[{"index":0,"embedding":[0.5]}]}""" + + test("JSON Content-Type replaces all caller casings in scalar and column header maps") { + val custom = Map("content-type" -> "text/plain", "Content-Type" -> "application/xml", "X-Custom" -> "retained") + val telemetry = Map("CONTENT-TYPE" -> "application/octet-stream", "X-Telemetry" -> "retained") + val schema = new StructType().add("text", StringType) + .add("custom", MapType(StringType, StringType)).add("telemetry", MapType(StringType, StringType)) + val data = spark.createDataFrame(Seq(Row("hello", custom, telemetry)).asJava, schema) + Seq(false, true).foreach { perRow => + Seq("chat_completions" -> chat, "decisions" -> NativeRequestFixtures.DecisionResponse) + .foreach { case (api, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + if (perRow) { + stage.setVectorParam(stage.customHeaders, "custom") + stage.setVectorParam(stage.telemHeaders, "telemetry") + } else stage.setCustomHeaders(custom).setTelemHeaders(telemetry) + val row = stage.transform(data).head() + assert(row.getAs[Row]("error") == null && row.getAs[Any]("answer") != null) + assert(server.requests.size() == 1) + val headers = server.headers.peek() + assert(headers("content-type") == "application/json") + assert(headers("x-custom") == "retained" && headers("x-telemetry") == "retained") + } + } + } + } + + test("HTTP header columns reject malformed names and values before sending valid siblings") { + val invalid = Seq( + "X-Caller: ignored\r\nX-Injected" -> "synthetic-value", + "X-Caller\nX-Injected" -> "synthetic-value", + "X-\u0000Caller" -> "synthetic-value", + "X Caller" -> "synthetic-value", + "X:Caller" -> "synthetic-value", + "" -> "synthetic-value", + "ap\u0131-key" -> "synthetic-value", + "author\u0131zation" -> "synthetic-value", + "X-Caller" -> "synthetic-value\r\nX-Injected: yes", + "X-Caller" -> "synthetic-\u0000value", + "X-Caller" -> "synthetic-\u010d\u010aX-Injected: yes", + "X-Caller" -> "synthetic-\u0085value", + "X-Caller" -> "synthetic-\u0159value") + val schema = new StructType().add("text", StringType).add("headers", MapType(StringType, StringType)) + val data = spark.createDataFrame((invalid.map { case (name, value) => Row("invalid", Map(name -> value)) } :+ + Row("valid", Map("X-Caller" -> "ordinary value", "X-Empty" -> "", "X-Latin" -> "caf\u00e9"))).asJava, schema) + Seq(false, true).foreach { telemetry => + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + stage.setVectorParam(if (telemetry) stage.telemHeaders else stage.customHeaders, "headers") + val rows = stage.transform(data).collect() + rows.filter(_.getAs[String]("text") == "invalid").foreach { row => + val diagnostic = row.getAs[Row]("error").getString(0).parseJson.asJsObject.fields + assert(diagnostic("code") == JsString("invalid_request")) + assert(Set[JsValue](JsString("requirement failed: Invalid HTTP header name"), + JsString("requirement failed: Invalid HTTP header value")) + .contains(diagnostic("message"))) + assert(row.getAs[String]("answer") == null) + } + assert(rows.find(_.getAs[String]("text") == "valid").get.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + assert(server.headers.peek()("x-caller") == "ordinary value") + assert(server.headers.peek()("x-empty") == "") + assert(server.headers.peek()("x-latin") == "caf\u00e9") + } + } + } + + test("HTTP header validation covers scalar maps and resolved native credentials without exposing values") { + val cases = Seq( + ("systemone", NativeRequestFixtures.SystemResponse, NativeRequestFixtures.SystemQuestions), + ("decisions", NativeRequestFixtures.DecisionResponse, NativeRequestFixtures.DecisionQuestions)) + cases.foreach { case (api, reply, questions) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + def stage: AIPrompt = new AIPrompt().setApiType(api).setQuestions(questions) + .setUrl(if (api == "systemone") server.url else server.url + "/v1") + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + val input = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", StringType)) + val invalid = Seq("synthetic-private-value\r\nX-Injected: yes", + "synthetic-private-value\u010d\u010aX-Injected: yes", "synthetic-private-value\u0085") + val scalarCases = invalid.flatMap { value => + Seq[AIPrompt => AIPrompt]( + _.setCustomHeaders(Map("X-Caller" -> value)), + _.setTelemHeaders(Map("X-Caller" -> value)), + _.setSubscriptionKey(value), + _.setAADToken(value), + _.setCustomAuthHeader(value)) + } ++ Seq("ap\u0131-key", "author\u0131zation").flatMap { name => + Seq[AIPrompt => AIPrompt]( + _.setCustomHeaders(Map(name -> "synthetic-private-value")), + _.setTelemHeaders(Map(name -> "synthetic-private-value"))) + } + scalarCases.foreach { configure => + val candidate = configure(stage) + Seq( + intercept[IllegalArgumentException](candidate.transformSchema(input.schema)), + intercept[IllegalArgumentException](candidate.transform(input)) + ).foreach { error => + assert(error.getMessage.contains("Invalid HTTP header") && Option(error.getCause).isEmpty) + assert(!error.getMessage.contains("synthetic-private-value") && !error.getMessage.contains("X-Injected")) + } + } + assert(server.requests.isEmpty) + val data = spark.createDataFrame((invalid.map(value => Row("bad", value)) :+ + Row("good", "synthetic-value")).asJava, + new StructType().add("text", StringType).add("credential", StringType)) + val columnCases: Seq[AIPrompt => AIPrompt] = Seq( + _.setSubscriptionKeyCol("credential"), _.setAADTokenCol("credential"), + _.setCustomAuthHeaderCol("credential")) + columnCases.foreach { configure => + val rows = configure(stage).transform(data).collect() + rows.filter(_.getAs[String]("text") == "bad").foreach { row => + val diagnostic = row.getAs[Row]("error").getString(0) + assert(diagnostic.parseJson.asJsObject.fields("message") == + JsString("requirement failed: Invalid HTTP header value")) + } + assert(rows.find(_.getAs[String]("text") == "good").get.getAs[Row]("error") == null) + } + assert(server.requests.size() == columnCases.size) + } + } + } + + test("scalar header preflight preserves ignored credential and reserved marker values") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", StringType)) + val malformed = "synthetic-ignored\r\nvalue" + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setErrorCol("error").setSubscriptionKey("synthetic-key") + .setAADToken(malformed).setCustomAuthHeader(malformed) + .setCustomHeaders(Map("Authorization" -> malformed, "X-Caller" -> malformed, + "content-type" -> malformed, HTTPRequestData.FabricAuthMarkerHeader -> malformed)) + .setTelemHeaders(Map("api-key" -> malformed, "X-Caller" -> "retained")) + assert(stage.transformSchema(data.schema).fieldNames.contains("error")) + assert(stage.transform(data).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.headers.peek()("x-caller") == "retained") + assert(server.headers.peek()("api-key") == "synthetic-key") + } + } + + test("mixed header preflight defers conditional scalar credentials to each row") { + val malformed = "synthetic-ignored\r\nvalue" + val data = spark.createDataFrame(Seq( + Row("valid", "synthetic-key"), Row("null", Option.empty[String].orNull), Row("blank", " ")).asJava, + new StructType().add("text", StringType).add("key", StringType)) + val cases: Seq[AIPrompt => AIPrompt] = Seq( + _.setSubscriptionKeyCol("key").setAADToken(malformed), + _.setAADTokenCol("key").setCustomAuthHeader(malformed), + _.setCustomAuthHeaderCol("key").setCustomHeaders(Map("api-key" -> malformed))) + cases.foreach { configure => + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + val stage = configure(new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error")) + assert(stage.transformSchema(data.schema).fieldNames.contains("answer")) + val rows = stage.transform(data).collect() + assert(rows.find(_.getAs[String]("text") == "valid").get.getAs[Row]("error") == null) + rows.filter(_.getAs[String]("text") != "valid").foreach { row => + val diagnostic = row.getAs[Row]("error").getString(0) + assert(diagnostic.contains("invalid_request") && !diagnostic.contains("synthetic-ignored")) + assert(row.getAs[String]("answer") == null) + } + assert(server.requests.size() == 1) + stage.setSubscriptionKey(malformed) + assert(intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + .getMessage.contains("Invalid HTTP header value")) + } + } + } + + test("mixed header preflight defers telemetry overrides but preserves guaranteed scalar authentication") { + val malformed = "synthetic-ignored\r\nvalue" + val data = spark.createDataFrame(Seq( + Row("valid", Map("X-Caller" -> "retained")), + Row("null", Option.empty[Map[String, String]].orNull), Row("missing", Map.empty[String, String])).asJava, + new StructType().add("text", StringType).add("telemetry", MapType(StringType, StringType))) + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + .setCustomHeaders(Map("X-Caller" -> malformed, "api-key" -> "synthetic-key")) + stage.setVectorParam(stage.telemHeaders, "telemetry") + val rows = stage.transform(data).collect() + assert(rows.find(_.getAs[String]("text") == "valid").get.getAs[Row]("error") == null) + assert(rows.filter(_.getAs[String]("text") != "valid").forall(row => + row.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(server.requests.size() == 1 && server.headers.peek()("x-caller") == "retained") + assert(server.headers.peek()("api-key") == "synthetic-key") + Seq(Map("api-key" -> malformed), Map("Invalid Header" -> "value")).foreach { headers => + stage.setCustomHeaders(headers) + assert(intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + .getMessage.contains("Invalid HTTP header")) + } + } + } + + test("native bearer conversion preserves credential sources when authentication header names coincide") { + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello", "synthetic-token")), + new StructType().add("text", "string").add("token", "string")) + val base = new TestableAISharedHeaderPrompt("shared-auth-header").setApiType("decisions") + .setUrl(server.url + "/v1").setModel("model").setInputCol("text").setErrorCol("error") + .setQuestions(NativeRequestFixtures.DecisionQuestions) + val cases: Seq[(AIPrompt => Unit, String)] = Seq( + (_.setAADToken("synthetic-token"), "Bearer synthetic-token"), + (_.setAADTokenCol("token"), "Bearer synthetic-token"), + (_.setCustomAuthHeader("Custom synthetic-token"), "Custom synthetic-token"), + (_.setCustomAuthHeaderCol("token"), "synthetic-token"), + (_.setSubscriptionKey("synthetic-key"), "Bearer synthetic-key"), + (_.setCustomHeaders(Map("authorization" -> "synthetic-key")), "Bearer synthetic-key"), + (stage => { + stage.setSubscriptionKey(" ").setAADToken("synthetic-token") + .setCustomAuthHeader("Custom ignored").setCustomHeaders(Map("authorization" -> "ignored")) + }, "Bearer synthetic-token")) + cases.foreach { case (configure, expected) => + val stage = base.copy(org.apache.spark.ml.param.ParamMap.empty) + configure(stage) + assert(stage.transform(data).head().getAs[Row]("error") == null) + assert(server.headers.remove().get("authorization").contains(expected)) + } + assert(server.requests.size() == cases.size) + } + } + + test("public inference honors inherited authentication header names including native bearer conversion") { + Seq("chat_completions" -> chat, "decisions" -> NativeRequestFixtures.DecisionResponse).foreach { + case (api, response) => + Using.resource(new AIRecordingServer((_, _) => 200 -> response)) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new TestableAIHeaderPrompt("custom-auth-headers").setApiType(api) + .setUrl(server.url + "/v1").setModel("model").setInputCol("text").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + Seq(false, true).foreach { useToken => + if (useToken) { + stage.clear(stage.subscriptionKey) + stage.setAADToken("synthetic-token") + } else stage.setSubscriptionKey("synthetic-key") + assert(stage.transform(data).head().getAs[Row]("error") == null) + val headers = server.headers.remove() + val bearer = useToken || api == "decisions" + val name = if (bearer) "x-model-token" else "x-model-key" + val credential = if (useToken) "synthetic-token" else "synthetic-key" + assert(headers.get(name).contains(if (bearer) "Bearer " + credential else credential)) + assert(!headers.contains("api-key") && !headers.contains("authorization")) + } + assert(server.requests.size() == 2) + } + } + } + + Seq(false, true).foreach { endpointConfigured => + test(s"native stages ignore global endpoint and API-key defaults with explicit endpoint=$endpointConfigured") { + val previousUrl = GlobalParams.getGlobalParam(URLKey) + val previousKey = GlobalParams.getGlobalParam(OpenAISubscriptionKey) + try { + Seq(new AIPrompt().setApiType("systemone"), new AISystemOne(), + new AIPrompt().setApiType("decisions"), new OpenAIDecisions()).foreach { stage => + val systemone = stage.getApiType == "systemone" + val reply = if (systemone) NativeRequestFixtures.SystemResponse else NativeRequestFixtures.DecisionResponse + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + GlobalParams.setGlobalParam(URLKey, server.url) + GlobalParams.setGlobalParam(OpenAISubscriptionKey, Left("synthetic-unrelated-key")) + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + stage.setModel("model").setInputCol("text").setErrorCol("error") + .setQuestions(if (systemone) NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions) + if (endpointConfigured) { + stage.setUrl(server.url) + assert(stage.transform(data).head().getAs[Row]("error") == null) + val anonymous = server.headers.peek() + assert(!anonymous.contains("authorization") && !anonymous.contains("api-key")) + server.headers.clear() + stage.setSubscriptionKey("synthetic-native-key") + assert(stage.transform(data).head().getAs[Row]("error") == null) + val name = if (systemone) "api-key" else "authorization" + val expected = if (systemone) "synthetic-native-key" else "Bearer synthetic-native-key" + assert(server.headers.peek()(name) == expected && server.requests.size() == 2) + } else { + val error = intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + assert(error.getMessage.contains("Configure an endpoint")) + intercept[IllegalArgumentException](stage.transform(data)) + assert(server.requests.isEmpty && !stage.isDefined(stage.url)) + } + } + } + } finally { + previousUrl match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + previousKey match { + case Some(value) => GlobalParams.setGlobalParam(OpenAISubscriptionKey, value) + case None => GlobalParams.resetGlobalParam(OpenAISubscriptionKey) + } + } + } + } + + private def withoutGlobalUrl(action: => Unit): Unit = { + val previous = GlobalParams.getGlobalParam(URLKey) + try { + GlobalParams.resetGlobalParam(URLKey) + action + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + + Seq(false, true).foreach { fail => + test(s"managed endpoint probe isolates caller global authentication and restores it with failure=$fail") { + val previousUrl = GlobalParams.getGlobalParam(URLKey) + val previousKey = GlobalParams.getGlobalParam(OpenAISubscriptionKey) + val callerUrl = "https://caller.example.invalid/v1" + val callerKey = Left("synthetic-caller-key") + val stopped = new IllegalStateException("synthetic probe interruption") + try { + GlobalParams.setGlobalParam(URLKey, callerUrl) + GlobalParams.setGlobalParam(OpenAISubscriptionKey, callerKey) + val create = () => { + if (fail) throw stopped + new TestableAIFabricAuthPrompt("managed-endpoint-probe") + } + if (fail) { + assert(intercept[IllegalStateException] { + AIPromptFabricProbe.verifyImplicitEndpoints(spark, create) + } eq stopped) + } else { + assert(AIPromptFabricProbe.verifyImplicitEndpoints(spark, create) == + """["chat_completions","responses","embeddings"]""") + } + assert(GlobalParams.getGlobalParam(URLKey).contains(callerUrl)) + assert(GlobalParams.getGlobalParam(OpenAISubscriptionKey).contains(callerKey)) + } finally { + previousUrl match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + previousKey match { + case Some(value) => GlobalParams.setGlobalParam(OpenAISubscriptionKey, value) + case None => GlobalParams.resetGlobalParam(OpenAISubscriptionKey) + } + } + } + } + + test("explicit Fabric credentials never evaluate an unavailable implicit fallback") { + withoutGlobalUrl { + val reply = chat + val data = spark.createDataFrame(Collections.singletonList( + Row("hello", "model", "synthetic-credential", Map("Authorization" -> "Embedded synthetic-credential"))), + new StructType().add("text", StringType).add("model", StringType).add("credential", StringType) + .add("headers", MapType(StringType, StringType))) + val cases: Seq[(AIPrompt => AIPrompt, String, String)] = Seq( + (_.setSubscriptionKey("synthetic-credential"), "api-key", "synthetic-credential"), + (_.setAADToken("synthetic-credential"), "authorization", "Bearer synthetic-credential"), + (_.setCustomAuthHeader("Custom synthetic-credential"), "authorization", "Custom synthetic-credential"), + (_.setSubscriptionKeyCol("credential"), "api-key", "synthetic-credential"), + (_.setAADTokenCol("credential"), "authorization", "Bearer synthetic-credential"), + (_.setCustomAuthHeaderCol("credential"), "authorization", "synthetic-credential"), + (_.setCustomHeaders(Map("API-KEY" -> "synthetic-credential")), "api-key", "synthetic-credential"), + (_.setCustomHeaders(Map("AUTHORIZATION" -> "Embedded synthetic-credential")), + "authorization", "Embedded synthetic-credential"), + (stage => stage.setVectorParam(stage.customHeaders, "headers"), "authorization", + "Embedded synthetic-credential")) + cases.foreach { case (configure, name, expected) => + val stage = configure(new TestableAIThrowingFabricAuthPrompt("unused-fabric-fallback") + .setModelCol("model").setInputCol("text").setOutputCol("answer").setErrorCol("error")) + stage.setHandler { (_, request) => + val headers = request.headers.map(h => h.name.toLowerCase(java.util.Locale.ROOT) -> h.value).toMap + assert(headers(name) == expected) + assert(!request.headers.exists(_.name == HTTPRequestData.FabricAuthMarkerHeader)) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + val row = stage.transform(data).head() + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "ok") + } + } + } + + Seq("chat_completions" -> chat, "responses" -> responses, "embeddings" -> embeddings) + .foreach { case (api, reply) => + test(s"$api automatically uses an unsaved Fabric endpoint when no URL is configured") { + withoutGlobalUrl { + val previousVersion = GlobalParams.getGlobalParam(OpenAIApiVersion) + val directory = Files.createTempDirectory("aiprompt-implicit-fabric-") + try { + GlobalParams.setGlobalParam(OpenAIApiVersion, Left("synthetic-version")) + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new TestableAIFabricAuthPrompt(s"implicit-$api").setApiType(api) + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + stage.setHandler { (_, request) => + assert(request.requestLine.uri.startsWith("https://fabric.example.invalid/cognitive/openai/")) + assert(request.requestLine.uri.contains("api-version=synthetic-version")) + val headers = request.headers.map(h => h.name.toLowerCase(java.util.Locale.ROOT) -> h.value).toMap + assert(headers("authorization") == "MwcToken synthetic") + assert(headers("x-taxonomy-traffictype") == "Background" && headers("x-llm-service-tier") == "flex") + assert(request.headers.exists(_.name == HTTPRequestData.FabricAuthMarkerHeader)) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + stage.transformSchema(data.schema) + assert(!stage.isDefined(stage.url) && !stage.isSet(stage.url)) + val row = stage.transform(data).head() + assert(row.getAs[Row]("error") == null && !row.isNullAt(row.fieldIndex("answer"))) + val path = directory.resolve("stage").toString + stage.write.save(path) + val loaded = AIPrompt.load(path) + assert(!loaded.isDefined(loaded.url) && !loaded.isSet(loaded.url)) + assert(loaded.effectiveCopy().getUrl == root) + assert(GlobalParams.getGlobalParam(URLKey).isEmpty) + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(directory.toFile) + previousVersion match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + } + } + + test("missing explicit endpoints fail clearly outside Fabric and for native Fabric APIs") { + withoutGlobalUrl { + val schema = new StructType().add("text", "string") + Seq(new AIPrompt().setInputCol("text"), + new TestableAIFabricAuthPrompt("no-native-route").setApiType("systemone") + .setInputCol("text").setQuestions(NativeRequestFixtures.SystemQuestions), + new TestableAIFabricAuthPrompt("no-decisions-route").setApiType("decisions") + .setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions)).foreach { stage => + stage.setModel("model") + val error = intercept[IllegalArgumentException](stage.transformSchema(schema)) + assert(error.getMessage.contains("endpoint") && error.getMessage.contains("setUrl")) + assert(!stage.isDefined(stage.url)) + } + + } + } + + test("injected bounded Fabric handler preserves rejected row diagnostics without authentication replay") { + val rejected = """{"error":{"code":"Unauthorized","message":"Synthetic rejection"}}""" + Using.resource(new AIRecordingServer( + (attempt, _) => if (attempt == 1) 401 -> rejected else 200 -> chat, + responseHeaders = Map("Retry-After" -> "120"))) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + stage.setHandler { (client, data) => + val request = data.toHTTPCore + try HandlingUtils.convertAndClose(HandlingUtils.sendWithRetryPolicy( + client, request, Array(0, 0), Set.empty, 0, boundedPolicy = true, quiet = true, + refreshAuth = Some(_ => "MwcToken fresh"))) + finally request.releaseConnection() + } + val row = stage.transform(data).head() + assert(row.isNullAt(row.fieldIndex("answer"))) + assert(row.getAs[Row]("error").getString(0).contains("http_error")) + assert(row.getAs[Row]("error").getString(0).contains("HTTP 401")) + assert(row.getAs[String]("raw") == rejected && server.requests.size() == 1) + } + } + + test("Fabric refresh failures become terminal HTTP row errors instead of Spark task failures") { + val rejected = """{"error":{"code":"Unauthorized","message":"Synthetic rejection"}}""" + Using.resource(new AIRecordingServer((_, _) => 401 -> rejected)) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + stage.setHandler { (client, data) => + val request = data.toHTTPCore + try HandlingUtils.convertAndClose(HandlingUtils.sendWithRetryPolicy( + client, request, Array(0, 0), Set.empty, 0, boundedPolicy = true, quiet = true, + refreshAuth = Some(_ => throw new NoSuchMethodException("Synthetic token refresh failure")))) + finally request.releaseConnection() + } + val row = stage.transform(data).head() + assert(row.isNullAt(row.fieldIndex("answer"))) + assert(row.getAs[Row]("error").getString(0).contains("http_error")) + assert(row.getAs[Row]("error").getString(0).contains("HTTP 401")) + assert(row.getAs[String]("raw") == rejected && server.requests.size() == 1) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptHttpPolicySuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptHttpPolicySuite.scala new file mode 100644 index 00000000000..9d45fa67196 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptHttpPolicySuite.scala @@ -0,0 +1,111 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.spark.ml.util.Identifiable +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{MapType, StringType, StructType} + +import scala.collection.JavaConverters._ +import scala.util.Using + +private[ai] class TestableFramingTokenAIPrompt(override val uid: String) extends AIPrompt(uid) { + def this() = this(Identifiable.randomUID("TestableFramingTokenAIPrompt")) + override protected val aadHeaderName: String = "cOnTeNt-LeNgTh" +} + +private[ai] class TestableFramingKeyAIPrompt(override val uid: String) extends AIPrompt(uid) { + def this() = this(Identifiable.randomUID("TestableFramingKeyAIPrompt")) + override protected val subscriptionKeyHeaderName: String = "tRaNsFeR-EnCoDiNg" +} + +class AIPromptHttpPolicySuite extends TestBase { + private val chat = """{"choices":[{"index":0,"message":{"role":"assistant","content":"caf\u00e9"}}]}""" + private def input = spark.createDataFrame(Seq(Row("hello")).asJava, new StructType().add("text", StringType)) + + private def stage(url: String): AIPrompt = new AIPrompt().setUrl(url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + + test("caller framing headers fail scalar preflight or individual rows without blocking other headers") { + val forbidden = Seq("Content-Length" -> "0", "cOnTeNt-LeNgTh" -> "123", + "Transfer-Encoding" -> "chunked", "tRaNsFeR-EnCoDiNg" -> "identity") + val preserved = Map("Host" -> "proxy.example.invalid", "Connection" -> "close", + "Expect" -> "100-continue", "TE" -> "trailers", "Upgrade" -> "example", "X-Proxy" -> "retained") + val schema = new StructType().add("text", StringType).add("headers", MapType(StringType, StringType)) + val data = spark.createDataFrame((forbidden.map(header => Row("bad", Map(header))) :+ + Row("good", preserved)).asJava, schema) + Seq(false, true).foreach { telemetry => + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + forbidden.foreach { header => + val configured = stage(server.url) + configured.setScalarParam(if (telemetry) configured.telemHeaders else configured.customHeaders, Map(header)) + assert(intercept[IllegalArgumentException](configured.transformSchema(input.schema)) + .getMessage.contains("HTTP framing headers")) + intercept[IllegalArgumentException](configured.transform(input)) + assert(server.requests.isEmpty) + } + val configured = stage(server.url) + configured.setVectorParam(if (telemetry) configured.telemHeaders else configured.customHeaders, "headers") + val rows = configured.transform(data).collect() + rows.filter(_.getAs[String]("text") == "bad").foreach { row => + val diagnostic = row.getAs[Row]("error").getString(0) + assert(diagnostic.contains("invalid_request") && diagnostic.contains("HTTP framing headers")) + } + assert(rows.find(_.getAs[String]("text") == "good").get.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + preserved.foreach { case (name, value) => + assert(server.headers.peek()(name.toLowerCase(java.util.Locale.ROOT)) == value) + } + } + } + } + + test("authentication framing hooks fail scalar preflight before HTTP") { + val cases: Seq[() => AIPrompt] = Seq( + () => new TestableFramingTokenAIPrompt().setAADToken("synthetic-token"), + () => new TestableFramingTokenAIPrompt().setCustomAuthHeader("synthetic-custom"), + () => new TestableFramingTokenAIPrompt().setApiType("decisions").setSubscriptionKey("synthetic-key") + .setQuestions(NativeRequestFixtures.DecisionQuestions), + () => new TestableFramingKeyAIPrompt().setSubscriptionKey("synthetic-key")) + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + cases.foreach { create => + val configured = create().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + Seq[() => Any](() => configured.transformSchema(input.schema), () => configured.transform(input)) + .foreach { invoke => + val error = intercept[IllegalArgumentException](invoke()) + assert(error.getMessage.contains("HTTP framing headers") && error.getCause == null) + assert(!error.getMessage.contains("synthetic-")) + assert(server.requests.isEmpty) + } + } + } + } + + test("authentication framing hooks are checked only after credential selection and bearer conversion") { + val data = spark.createDataFrame(Seq(Row("valid", "synthetic-key"), + Row("fallback", Option.empty[String].orNull)).asJava, + new StructType().add("text", StringType).add("key", StringType)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chat)) { server => + val configured = new TestableFramingTokenAIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setSubscriptionKeyCol("key").setAADToken("synthetic-token") + assert(configured.transformSchema(data.schema).fieldNames.contains("answer")) + val rows = configured.transform(data).collect() + assert(rows.find(_.getAs[String]("text") == "valid").get.getAs[Row]("error") == null) + val rejected = rows.find(_.getAs[String]("text") == "fallback").get + val diagnostic = rejected.getAs[Row]("error").getString(0) + assert(diagnostic.contains("invalid_request") && diagnostic.contains("HTTP framing headers")) + assert(!diagnostic.contains("synthetic-token") && rejected.getAs[String]("answer") == null) + assert(server.requests.size() == 1 && server.headers.peek()("api-key") == "synthetic-key") + } + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val converted = new TestableFramingKeyAIPrompt().setUrl(server.url + "/v1").setApiType("decisions") + .setModel("model").setInputCol("text").setErrorCol("error").setSubscriptionKey("synthetic-key") + .setQuestions(NativeRequestFixtures.DecisionQuestions) + assert(converted.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.headers.peek()("authorization") == "Bearer synthetic-key") + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptInputSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptInputSuite.scala new file mode 100644 index 00000000000..2f7223b5325 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptInputSuite.scala @@ -0,0 +1,349 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.io.http.{HTTPSchema, URLKey} +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.OpenAIApiVersion +import com.microsoft.azure.synapse.ml.services.aifoundry.AIFoundryModel +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIDecisions, OpenAIEmbeddingDeploymentNameKey} +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.{Encoders, Row} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import spray.json._ + +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptInputSuite extends AIPromptTestBase { + private val chatReply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + + private def contents(server: AIRecordingServer): Vector[String] = + server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + }.sorted + + test("inherited optional parameters use fully initialized response format parameters") { + Seq(new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + stage.setResponseFormat("json_object") + assert(stage.getOptionalParams(Row.empty)("response_format") == Map("type" -> "json_object")) + assert(stage.params.count(_.name == "responseFormat") == 1) + assert(stage.copy(ParamMap.empty).getOptionalParams(Row.empty)("response_format") == + Map("type" -> "json_object")) + } + } + + test("null bound API versions produce row errors without requests on dated routes") { + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + val data = spark.createDataFrame( + Seq(Row("valid", "selected-version"), Row("missing", Option.empty[String].orNull)).asJava, + new StructType().add("text", StringType).add("version", StringType)) + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq(false, true).foreach { global => + val stage = new AIPrompt().setUrl(server.url).setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + if (global) GlobalParams.setGlobalParam(OpenAIApiVersion, Right("version")) + else stage.setApiVersionCol("version") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect().map(row => row.getAs[String]("text") -> row).toMap + assert(rows("valid").getAs[String]("answer") == "answer") + assert(rows("missing").getAs[String]("answer") == null) + assert(AIJson.objectString(rows("missing").getAs[Row]("error").getString(0)) + .fields("code") == JsString("missing_input")) + } + assert(server.requests.size() == 2) + assert(server.requests.asScala.forall(_._1.endsWith("?api-version=selected-version"))) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("implicit embedding deployment excludes unused model defaults from planning") { + val previousModel = GlobalParams.getGlobalParam(AIFoundryModel) + val previousEmbedding = GlobalParams.getGlobalParam(OpenAIEmbeddingDeploymentNameKey) + val reply = """{"data":[{"index":0,"embedding":[0.25,0.5]}],"model":"embedding-model"}""" + val input = spark.createDataFrame(Collections.singletonList(Row("hello")), new StructType().add("text", StringType)) + try { + GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, Left("embedding-model")) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType("embeddings").setUrl(server.url + "/v1").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + GlobalParams.setGlobalParam(AIFoundryModel, Left("unused-model")) + val expected = stage.transformSchema(input.schema) + Seq[Either[String, String]](Right("absent-model"), Left("changed-unused-model")).foreach { value => + GlobalParams.setGlobalParam(AIFoundryModel, value) + assert(stage.transformSchema(input.schema) == expected) + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(AIJson.objectString(server.requests.asScala.last._2).fields("model") == JsString("embedding-model")) + } + assert(server.requests.size() == 2) + } + } finally { + previousModel match { + case Some(value) => GlobalParams.setGlobalParam(AIFoundryModel, value) + case None => GlobalParams.resetGlobalParam(AIFoundryModel) + } + previousEmbedding match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, value) + case None => GlobalParams.resetGlobalParam(OpenAIEmbeddingDeploymentNameKey) + } + } + } + + test("wide column validation does not repeatedly traverse input prefixes") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + try { + Seq(false, true).foreach { sensitive => + spark.conf.set("spark.sql.caseSensitive", sensitive.toString) + var reads = 0 + val names = new IndexedSeq[String] { + override def length: Int = 10000 + override def apply(index: Int): String = { + reads += 1 + assert(reads <= 5 * length, "Column validation repeatedly traversed the input names") + s"feature_$index" + } + } + AIColumnPlan.distinct(names, "input columns") + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + + test("column ambiguity preserves Spark Unicode case comparison in public schema planning") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + try { + Seq(false, true).foreach { sensitive => + spark.conf.set("spark.sql.caseSensitive", sensitive.toString) + Seq("text" -> "TEXT", "i" -> "\u0131", "\u03c3" -> "\u03c2", "\u00e9" -> "e").foreach { + case (left, right) => + val schema = new StructType().add(left, StringType).add(right, StringType) + val stage = new AIPrompt().setUrl("http://example.invalid/v1").setModel("model") + .setInputCol(left).setOutputCol("answer").setErrorCol("error") + if (org.apache.spark.sql.internal.SQLConf.get.resolver(left, right)) { + assert(intercept[IllegalArgumentException](stage.transformSchema(schema)) + .getMessage.contains("Ambiguous input columns")) + } else assert(stage.transformSchema(schema).fieldNames.take(2).toSeq == Seq(left, right)) + } + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + + test("non-row-local templates reject before inference") { + val data = spark.range(2).selectExpr("id", + "CASE WHEN id = 0 THEN CAST('2020-01-01 00:00:00' AS TIMESTAMP) END AS ts") + withServer { server => + Seq("{cast(window(ts, '1 hour').start as string)}", + "{cast(window(ts, '1 hour', '30 minutes').start as string)}", + "{cast(session_window(ts, '1 hour').start as string)}", + "{first(cast(id as string))}", "{row_number() over (order by id)}").foreach { template => + val stage = prompt(server).setPromptTemplate(template) + Seq(intercept[IllegalArgumentException](stage.transformSchema(data.schema)), + intercept[IllegalArgumentException](stage.transform(data))).foreach { error => + assert(error.getMessage.contains("row-local") && Option(error.getCause).isEmpty) + } + assert(server.requests.get() == 0) + } + } + } + + Seq("SHOW TABLES" -> "tableName", "SHOW COLUMNS" -> "col_name").foreach { case (command, column) => + test(s"templates preserve command-result input from $command") { + val view = "aiprompt_command_" + java.util.UUID.randomUUID().toString.replace("-", "") + spark.range(1).selectExpr("id AS first", "'value' AS second").createOrReplaceTempView(view) + try { + val data = spark.sql(if (command == "SHOW TABLES") s"$command LIKE '$view'" else s"$command IN $view") + val expected = data.collect().toVector + assert(expected.nonEmpty) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate(s"Describe {$column}").setOutputCol("answer").setErrorCol("error") + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema) && server.requests.isEmpty) + val rows = output.collect() + assert(rows.map(_.toSeq.take(data.schema.length)).toSet == expected.map(_.toSeq).toSet) + assert(rows.forall(row => row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + assert(contents(server) == expected.map(row => s"Describe ${row.getAs[String](column)}").sorted) + assert(server.requests.size() == expected.size) + } + } finally spark.catalog.dropTempView(view) + } + } + + Seq( + ("to_json(map(*))", Seq("concat('key-', id) AS first", "cast(id as string) AS second")), + ("nvl(*)", Seq("CASE WHEN id = 0 THEN 'first' END AS first", "cast(id as string) AS second")), + ("cast(coalesce(*) as string)", Seq( + "CASE WHEN id = 0 THEN DATE '2020-01-01' END AS first", "DATE '2020-01-02' AS second")), + ("to_json(struct(*))", Seq("'duplicate' AS text", "'public' AS _aiprompt_row_public")) + ).foreach { case (expression, columns) => + test(s"wildcard template $expression uses only public columns") { + val data = spark.range(3).selectExpr(columns: _*) + val expected = data.selectExpr(expression).collect().map(_.getString(0)).toVector.sorted + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate(s"{$expression}").setOutputCol("answer").setErrorCol("error") + val output = stage.transform(data) + val rows = output.collect() + assert(output.schema == stage.transformSchema(data.schema) && rows.length == 3) + assert(rows.forall(row => row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + assert(rows.map(row => row.toSeq.take(2)).toSet == data.collect().map(_.toSeq).toSet) + assert(contents(server) == expected && server.requests.size() == 3) + } + } + } + + test("templates preserve precomputed window columns and null input rows") { + val data = spark.range(2).selectExpr("id", "row_number() over (order by id) AS ordinal", + "CASE WHEN id = 0 THEN CAST('2020-01-01 00:00:00' AS TIMESTAMP) END AS ts") + withServer { server => + val stage = prompt(server).setPromptTemplate( + "{coalesce(cast(ts as string), 'missing')}:{cast(ordinal as string)}") + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val rows = output.collect().sortBy(_.getAs[Long]("id")) + assert(rows.map(_.getAs[Int]("ordinal")).toSeq == Seq(1, 2)) + assert(rows(1).isNullAt(rows(1).fieldIndex("ts"))) + assert(rows.forall(row => row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + assert(server.requests.get() == 2) + } + } + + test("nested message schema attachment validates the whole array before returning") { + val schema = new StructType().add("messages", ArrayType(new StructType().add("role", StringType))) + intercept[IndexOutOfBoundsException](AIColumnPlan.withSchema(Row(Seq(Row("user"), Row.empty)), schema)) + } + + test("fused JVM message arrays preserve retained values and invalid row errors") { + val messageType = new StructType().add("role", StringType) + .add("content", ArrayType(MapType(StringType, StringType))) + val schema = new StructType().add("id", "long").add("messages", ArrayType(messageType)) + val input = spark.range(2).map { id => + val message = if (id == 0) Row("user", Array(Map("type" -> "text", "text" -> "hello"))) + else Option.empty[Row].orNull + Row(id, Array(message)) + }(Encoders.row(schema)) + val expected = input.repartition(1).collect().map(row => row.getLong(0) -> row).toMap + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setMessagesCol("messages") + .setOutputCol("answer").setErrorCol("error") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(schema)) + val rows = output.collect().map(row => row.getLong(0) -> row).toMap + assert(rows(0L).getAs[String]("answer") == "answer" && rows(0L).getAs[Row]("error") == null) + assert(rows(1L).getAs[Row]("error").getString(0).contains("invalid_request")) + rows.foreach { case (id, row) => + assert(row.getAs[Seq[Row]]("messages") == expected(id).getAs[Seq[Row]]("messages")) + } + assert(contents(server) == Vector("hello")) + } + } + + test("templates support UDFs subqueries and random IV encryption without early evaluation") { + val session = spark.newSession() + import session.implicits._ + session.udf.register("aiprompt_render_ticket", (text: String) => "old:" + text) + // Keep plaintext as an input attribute so Spark does not constant-fold encryption. + val data = Seq.fill(2)("hello").toDF("text") + AIPrompt.withSession(session) { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq("{aiprompt_render_ticket(text)}" -> "old:hello", + "{concat((SELECT 'old'), ':', text)}" -> "old:hello", + "{(SELECT max('old'))}:{text}" -> "old:hello").foreach { case (template, expected) => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate(template).setOutputCol("answer").setErrorCol("error") + val output = stage.transform(data.limit(1)) + assert(server.requests.isEmpty && output.schema == stage.transformSchema(data.schema)) + assert(output.head().getAs[String]("answer") == "answer") + assert(contents(server) == Vector(expected)) + server.requests.clear() + } + val encrypted = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{base64(aes_encrypt(text, '0000111122223333'))}") + .setOutputCol("answer").setErrorCol("error") + assert(encrypted.transform(data).collect().forall(_.getAs[Row]("error") == null)) + assert(contents(server).size == 2 && contents(server).distinct.size == 2) + } + } + } + + test("schema planning allows row credentials extensions handlers and runtime templates without jobs") { + val session = spark.newSession() + session.udf.register("aiprompt_schema_udf", (value: String) => value) + val schema = new StructType().add("text", StringType).add("credential", StringType) + .add("headers", MapType(StringType, StringType)) + AIPrompt.withSession(session) { + withServer { server => + val settings: Seq[AIPrompt => AIPrompt] = Seq( + _.setSubscriptionKeyCol("credential"), _.setAADTokenCol("credential"), + _.setCustomAuthHeaderCol("credential"), stage => stage.setVectorParam(stage.customHeaders, "headers"), + stage => stage.setVectorParam(stage.telemHeaders, "headers"), stage => stage.setHandler(stage.getHandler), + _.setPreviousResponseId("provider-response"), _.setPreviousResponseIdCol("credential"), + _.setExtraBody(Map("file_id" -> "provider-file")), + _.setExtraBody(Map("extension" -> Map("file_url" -> "https://example.invalid/media")))) + val group = java.util.UUID.randomUUID().toString + session.sparkContext.setJobGroup(group, "Schema-only configuration validation") + try { + settings.foreach(configure => + configure(prompt(server).setApiType("responses").setUrl(server.url + "/v1")).transformSchema(schema)) + Seq("{current_timestamp()}", "{aiprompt_schema_udf(text)}", "{(select 'external')}", "{upper(text)}:{1 / 0}") + .foreach(template => prompt(server).setPromptTemplate(template).transformSchema(schema)) + assert(session.sparkContext.statusTracker.getJobIdsForGroup(group).isEmpty) + assert(server.requests.get() == 0) + } finally session.sparkContext.clearJobGroup() + } + } + } + + test("telemetry cannot override implicit Fabric classification but explicit endpoints retain caller headers") { + val previous = GlobalParams.getGlobalParam(URLKey) + val root = "https://fabric.example.invalid/cognitive/openai" + val headers = Map("X-Llm-Service-Tier" -> "standard", "x-llm-service-tier" -> "other", + "X-Taxonomy-TrafficType" -> "Interactive", "x-taxonomy-traffictype" -> "Other", + "X-Taxonomy-ExtendedProperties" -> "caller", "x-taxonomy-extendedproperties" -> "other", "x-trace" -> "retained") + try { + GlobalParams.setGlobalParam(URLKey, root) + for (perRow <- Seq(false, true); explicitKey <- Seq(false, true); explicitUrl <- Seq(false, true)) { + val data = spark.range(1).withColumn("text", lit("hello")).withColumn("headers", typedLit(headers)) + val stage = new TestableAIFabricAuthPrompt("fabric-telemetry").setModel("model") + .setInputCol("text").setOutputCol("answer") + if (explicitKey) stage.setSubscriptionKey("synthetic-explicit-key") + if (explicitUrl) stage.setUrl(root) + if (perRow) stage.setVectorParam(stage.telemHeaders, "headers") else stage.setTelemHeaders(headers) + val reply = chatReply + stage.setHandler { (_, request) => + val grouped = request.headers.groupBy(_.name.toLowerCase(java.util.Locale.ROOT)) + assert(grouped("x-trace").map(_.value).toSeq == Seq("retained")) + if (explicitUrl) { + assert(grouped("x-llm-service-tier").map(_.value).toSet == Set("standard", "other")) + } else { + Seq("x-llm-service-tier", "x-taxonomy-traffictype", "x-taxonomy-extendedproperties").foreach { key => + assert(grouped(key).length == 1, s"Duplicate owned header $key") + } + assert(grouped("x-llm-service-tier").head.value == "flex") + assert(grouped("x-taxonomy-traffictype").head.value == "Background") + assert(AIJson.objectString(grouped("x-taxonomy-extendedproperties").head.value) + .fields("feature") == JsString("synapseml")) + } + if (explicitKey) assert(grouped("api-key").head.value == "synthetic-explicit-key") + HTTPSchema.stringToResponse(reply, 200, "OK") + } + assert(stage.transform(data).head().getAs[String]("answer") == "answer") + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLoggingSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLoggingSuite.scala new file mode 100644 index 00000000000..5a61b0fae62 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLoggingSuite.scala @@ -0,0 +1,553 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.io.http.{HTTPTransformer, HandlingUtils} +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIChatCompletion, OpenAIDecisions, OpenAIResponses} +import com.sun.net.httpserver.{HttpExchange, HttpServer} +import org.apache.commons.io.IOUtils +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.{Filter, LogEvent, LoggerContext} +import org.apache.logging.log4j.core.appender.AbstractAppender +import org.apache.logging.log4j.core.config.Property +import org.apache.logging.log4j.core.layout.PatternLayout +import org.apache.spark.ml.{Pipeline, PipelineModel, Serializer} +import org.apache.spark.ml.param.{Param, ParamMap} +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.types.{StringType, StructType} + +import java.io.{PrintWriter, StringWriter} +import java.net.InetSocketAddress +import java.nio.file.Files +import java.util.concurrent.ConcurrentLinkedQueue +import scala.collection.JavaConverters._ +import scala.util.Using + +private[ai] class TestableAIMutableDefaultHandlerPrompt(override val uid: String) extends AIPrompt(uid) { + def replaceHandlerDefault(): this.type = setDefault(handler -> HandlingUtils.advancedUDF(0)) +} + +private[ai] class TestableAIAlternateDefaultHandlerPrompt(override val uid: String) extends AIPrompt(uid) { + setDefault(handler -> HandlingUtils.advancedUDF(0)) +} + +private[ai] class TestableAIDefaultTimeoutPrompt(override val uid: String) extends AIPrompt(uid) { + def setConcurrentTimeoutDefault(value: Double): this.type = setDefault(concurrentTimeout -> value) +} + +private[ai] class TestableAIMutableJsonDefaultPrompt(override val uid: String) extends AIPrompt(uid) { + def replaceJsonDefault[T](param: Param[T], value: T): this.type = setDefault(param -> value) +} + +class AIPromptLoggingSuite extends AIPromptTestBase { + private val chatReply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + + private def input: DataFrame = spark.createDataFrame( + Seq(Row("hello")).asJava, new StructType().add("text", StringType)) + + private def withLogs(action: ConcurrentLinkedQueue[String] => Unit): Unit = { + val logged = new ConcurrentLinkedQueue[String]() + val layout = PatternLayout.newBuilder().withPattern("%m%n%throwable").build() + val noFilter = Option.empty[Filter].orNull + val appender = new AbstractAppender("aiprompt-test", noFilter, layout, false, Array.empty[Property]) { + override def append(event: LogEvent): Unit = logged.add(layout.toSerializable(event)) + } + val root = LogManager.getContext(false).asInstanceOf[LoggerContext].getRootLogger + val loggers = Seq(classOf[AIPrompt], classOf[OpenAIChatCompletion], classOf[OpenAIResponses], + classOf[HTTPTransformer]) + .map(cls => LogManager.getLogger(cls).asInstanceOf[org.apache.logging.log4j.core.Logger]) + val levels = loggers.map(logger => logger -> logger.getLevel) + appender.start() + root.addAppender(appender) + loggers.foreach(_.setLevel(Level.INFO)) + try action(logged) finally { + levels.foreach { case (logger, level) => logger.setLevel(level) } + root.removeAppender(appender) + appender.stop() + } + } + + test("request planning uses pure encoders without constructing legacy OpenAI stages") { + val data = input + Seq("chat_completions", "responses", "systemone", "decisions", "embeddings").foreach { api => + val endpoint = if (api == "systemone") "https://example.services.ai.azure.com" else "https://example.invalid/v1" + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model").setInputCol("text") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + withLogs { logged => + stage.transformSchema(data.schema) + stage.transform(data) + assert(logged.asScala.exists(_.contains("constructor")), "Constructor logging must be observable") + Seq("OpenAIChatCompletion", "OpenAIResponses").foreach { legacy => + assert(!logged.asScala.exists(_.contains(s".$legacy")), s"Planning constructed $legacy for $api") + } + } + } + } + + Seq(false, true).foreach { succeeds => + test(s"public transform logging uses the execution session with succeeds=$succeeds") { + val original = SparkSession.getActiveSession + try { + val previous = spark.newSession() + val execution = spark.newSession() + previous.conf.set("trident.artifact.id", "synthetic-previous-artifact") + execution.conf.set("trident.artifact.id", "synthetic-execution-artifact") + val data = execution.createDataFrame( + Seq(Row("hello")).asJava, new StructType().add("text", StringType)) + withServer { server => + Seq(Some(previous), None).foreach { prior => + val stage = prompt(server) + if (!succeeds) stage.setInputCol("missing") + prior match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + withLogs { logged => + def invoke(): DataFrame = stage.transform(data) + if (succeeds) { + val result = invoke() + assert(SparkSession.getActiveSession == prior && result.sparkSession == execution) + assert(result.count() == 1) + } else { + intercept[IllegalArgumentException](invoke()) + assert(SparkSession.getActiveSession == prior) + } + val messages = logged.asScala.filter(message => + message.contains(s""""modelUid":"${stage.uid}"""") && message.contains(""""method":"transform"""")) + assert(messages.size == 1) + assert(messages.head.contains(""""artifactId":"synthetic-execution-artifact"""")) + } + } + } + } finally { + original match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } + } + + Seq(1, 2).foreach { concurrency => + test(s"bounded transport never replays response cookies across row credentials at concurrency $concurrency") { + val keys = (1 to 4).map(number => s"synthetic-key-$number") + val data = spark.createDataFrame(keys.map(key => Row("hello", key)).asJava, + new StructType().add("text", StringType).add("credential", StringType)).coalesce(1) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply, + responseHeaders = Map("Set-Cookie" -> "provider-session=synthetic; Path=/"))) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setSubscriptionKeyCol("credential") + .setConcurrency(concurrency) + val rows = stage.transform(data).collect() + assert(rows.length == 4 && rows.forall(row => + row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + val headers = server.headers.asScala.toVector + assert(headers.size == 4 && headers.map(_("api-key")).sorted == keys.sorted) + assert(headers.forall(!_.contains("cookie"))) + } + } + } + + test("bounded transport preserves caller cookies while explicit handlers retain legacy cookie management") { + Seq(false, true).foreach { legacy => + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply, + responseHeaders = Map("Set-Cookie" -> "provider-session=synthetic; Path=/"))) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setConcurrency(1) + if (legacy) stage.setHandler(HandlingUtils.advancedUDF(0)) + else stage.setCustomHeaders(Map("Cookie" -> "caller-session=explicit")) + val rows = stage.transform(input.union(input).coalesce(1)).collect() + assert(rows.length == 2 && rows.forall(row => + row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + val cookies = server.headers.asScala.toVector.map(_.get("cookie")) + if (legacy) assert(cookies == Vector(None, Some("provider-session=synthetic"))) + else assert(cookies == Vector.fill(2)(Some("caller-session=explicit"))) + } + } + } + + test("redirects retain original HTTP diagnostics without forwarding or logging provider credentials") { + Seq(false, true).foreach { truncated => + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { target => + val source = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + val received = new ConcurrentLinkedQueue[String]() + source.createContext("/", (exchange: HttpExchange) => { + received.add(exchange.getRequestMethod) + Using.resource(exchange.getRequestBody)(IOUtils.toByteArray) + exchange.getResponseHeaders.add("Location", target.url + "/v1/chat/completions") + exchange.sendResponseHeaders(303, if (truncated) 100 else -1) + if (truncated) exchange.getResponseBody.write('{'.toInt) + exchange.close() + }) + source.start() + try { + withLogs { logged => + val key = "synthetic-redirect-key" + Seq(1, 2).foreach { concurrency => + val stage = new AIPrompt().setUrl(s"http://127.0.0.1:${source.getAddress.getPort}/v1") + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setRawResponseCol("raw").setSubscriptionKey(key).setConcurrency(concurrency).setTimeout(2) + val row = stage.transform(input).head() + assert(target.requests.isEmpty && target.headers.isEmpty) + val error = row.getAs[Row]("error") + assert(row.isNullAt(row.fieldIndex("answer")) && error.getString(0).contains("http_error")) + assert(error.getStruct(1).getAs[Int]("statusCode") == 303) + if (truncated) assert(row.getAs[String]("raw").contains("Could not read")) + } + assert(received.asScala.toSeq == Seq("POST", "POST")) + assert(!logged.isEmpty && logged.asScala.forall(!_.contains(key))) + } + } finally source.stop(0) + } + } + } + + for (binding <- Seq("constructor", "runtime", "explicit"); concurrency <- Seq(1, 2)) { + test(s"custom handlers retain legacy redirects with $binding binding at concurrency $concurrency") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { target => + val source = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + val received = new ConcurrentLinkedQueue[String]() + source.createContext("/", (exchange: HttpExchange) => { + received.add(exchange.getRequestMethod) + Using.resource(exchange.getRequestBody)(IOUtils.toByteArray) + exchange.getResponseHeaders.add("Location", target.url + "/v1/chat/completions") + exchange.sendResponseHeaders(303, -1) + exchange.close() + }) + source.start() + try { + val stage = (binding match { + case "explicit" => new AIPrompt().setHandler(HandlingUtils.advancedUDF(0)) + case "runtime" => new TestableAIMutableDefaultHandlerPrompt("runtime-redirect").replaceHandlerDefault() + case _ => new TestableAIAlternateDefaultHandlerPrompt("alternate-redirect") + }) + .setUrl(s"http://127.0.0.1:${source.getAddress.getPort}/v1") + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setConcurrency(concurrency) + val restored = if (binding == "explicit") { + val saved = path() + stage.write.session(spark).save(saved) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + try { + spark.conf.unset(config) + val error = intercept[SecurityException](AIPrompt.read.session(spark).load(saved)) + assert(error.getMessage.contains(config)) + spark.conf.set(config, "true") + AIPrompt.read.session(spark).load(saved) + } finally previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } else stage.copy(ParamMap.empty) + assert(restored.isSet(restored.handler) == (binding == "explicit")) + val row = restored.transform(input).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(received.asScala.toSeq == Seq("POST") && target.requests.size() == 1) + } finally source.stop(0) + } + } + } + + test("bounded default handlers load without unsafe Java deserialization opt-in") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val saved = path() + new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text").setOutputCol("answer") + .write.session(spark).save(saved) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + try { + spark.conf.unset(config) + val restored = AIPrompt.read.session(spark).load(saved) + assert(restored.transform(input).head().getAs[String]("answer") == "answer") + assert(server.requests.size() == 1) + } finally previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + + Seq(false, true).foreach { allowed => + test(s"custom handlers use their assigned Spark session deserialization policy with allowed=$allowed") { + val config = Serializer.LegacyObjectDeserializationConfig + val active = spark.newSession() + active.conf.set(config, (!allowed).toString) + val readerSession = spark.newSession() + readerSession.conf.set(config, allowed.toString) + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(active) + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val saved = path() + new AIPrompt().setHandler(HandlingUtils.advancedUDF(0)) + .setUrl(server.url + "/v1").setModel("model").setInputCol("text").setOutputCol("answer") + .write.session(readerSession).save(saved) + if (allowed) { + val restored = AIPrompt.read.session(readerSession).load(saved) + assert(restored.transform(input).head().getAs[String]("answer") == "answer") + assert(server.requests.size() == 1) + } else { + val error = intercept[SecurityException](AIPrompt.read.session(readerSession).load(saved)) + assert(error.getMessage.contains(config) && server.requests.isEmpty) + } + } + } finally previous.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + + for (binding <- Seq("explicit", "default", "default-copy", "default-load"); concurrency <- Seq(1, 2)) { + test(s"concurrent timeout cancels bounded retry workers with $binding binding at concurrency $concurrency") { + Using.resource(new AIRecordingServer((number, _) => + if (number <= concurrency) 429 -> "{}" else 200 -> chatReply, retryAfter = "3")) { server => + val configured = (if (binding == "explicit") new AIPrompt().setConcurrentTimeout(1) + else new TestableAIDefaultTimeoutPrompt("default-timeout").setConcurrentTimeoutDefault(1)) + .setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setConcurrency(concurrency).setTimeout(5) + val stage = binding match { + case "default-copy" => configured.copy(ParamMap.empty) + case "default-load" => + val saved = tmpDir.resolve(s"timeout-$concurrency").toString + configured.write.session(spark).save(saved) + AIPrompt.read.session(spark).load(saved) + case _ => configured + } + assert(stage.isSet(stage.concurrentTimeout) == (binding == "explicit")) + intercept[org.apache.spark.SparkException] { + stage.transform(input.union(input).coalesce(1)).collect() + } + assert(server.requests.size() == concurrency) + Thread.sleep(4000) + assert(server.requests.size() == concurrency, "Timed-out workers issued another paid request") + } + } + } + + Seq("original", "copy", "deferred", "constructor").foreach { mode => + test(s"saving replaced handler defaults fails before overwriting existing files with mode=$mode") { + val stage = new TestableAIMutableDefaultHandlerPrompt("runtime-default") + val deferred = stage.write.session(spark) + stage.replaceHandlerDefault() + val candidate = mode match { + case "copy" => stage.copy(ParamMap.empty) + case "constructor" => new TestableAIAlternateDefaultHandlerPrompt("constructor-default") + case _ => stage + } + val saved = Files.createTempDirectory(tmpDir, "handler-default") + val marker = saved.resolve("keep") + Files.write(marker, Array[Byte](1, 2, 3)) + val failure = intercept[IllegalArgumentException] { + val writer = if (mode == "deferred") deferred else candidate.write.session(spark) + writer.overwrite().save(saved.toString) + } + assert(failure.getMessage.contains("setHandler")) + assert(Files.readAllBytes(marker).toSeq == Seq[Byte](1, 2, 3)) + } + } + + test("public native writers reject invalid parsers before replacing saved stages") { + Seq[() => AIPrompt](() => new AIPrompt(), () => new AISystemOne(), () => new OpenAIDecisions()) + .foreach { create => + Seq("csv", "json").foreach { parser => + Seq(false, true).foreach { deferred => + val saved = path() + create().setPostProcessingOptions(Map("delimiter" -> "|")).write.session(spark).save(saved) + val stage = create() + if (parser == "csv") stage.setPostProcessingOptions(Map("delimiter" -> ";")) + val writer = if (deferred) Some(stage.write.session(spark).overwrite()) else None + if (parser == "csv") stage.clear(stage.postProcessing) else stage.set(stage.postProcessing, "json") + val error = intercept[IllegalArgumentException] { + writer.getOrElse(stage.write.session(spark).overwrite()).save(saved) + } + assert(error.getMessage == (if (parser == "csv") "postProcessing must be 'csv'" + else "requirement failed: jsonSchema must be specified with json postProcessing")) + val restored = stage match { + case _: AISystemOne => AISystemOne.read.session(spark).load(saved) + case _: OpenAIDecisions => OpenAIDecisions.read.session(spark).load(saved) + case _ => AIPrompt.read.session(spark).load(saved) + } + assert(restored.getClass == stage.getClass) + assert(restored.getPostProcessing == "csv" && restored.getPostProcessingOptions == Map("delimiter" -> "|")) + } + } + } + } + + test("unencodable generic JSON bindings fail before direct deferred and copied stage overwrites") { + val saved = path() + new AIPrompt().setExtraBody(Map("keep" -> 1)).write.session(spark).save(saved) + def unchanged(): Unit = assert(AIPrompt.read.session(spark).load(saved).getExtraBody == Map("keep" -> 1)) + val invalid = Map[String, Any]("x" -> Double.NaN) + Seq[() => AIPrompt](() => new AIPrompt(), () => new AISystemOne(), () => new OpenAIDecisions()) + .foreach { create => + val stage = create() + def verify[T](param: Param[T], value: T): Unit = { + val deferred = stage.write.session(spark).overwrite() + val copied = stage.copy(ParamMap(param -> value)) + stage.set(param, value) + Seq(() => stage.write.session(spark).overwrite().save(saved), () => deferred.save(saved), + () => copied.write.session(spark).overwrite().save(saved)).foreach { save => + assert(intercept[IllegalArgumentException](save()).getMessage.contains("finite")) + unchanged() + } + stage.clear(param) + } + verify(stage.extraBody, Left(invalid)) + verify(stage.responseFormat, Left(invalid)) + verify(stage.questions, Left(invalid)) + } + } + + test("shadowed unencodable JSON defaults fail before overwriting a saved stage") { + val saved = path() + new AIPrompt().setExtraBody(Map("keep" -> 1)).write.session(spark).save(saved) + val invalid = Map[String, Any]("x" -> Double.PositiveInfinity) + Seq("extraBody", "responseFormat", "questions").foreach { name => + val stage = new TestableAIMutableJsonDefaultPrompt("invalid-json-default") + .setExtraBody(Map("valid" -> 1)).setQuestions("""{"valid":1}""") + .setResponseSchema(Map("type" -> "object")) + val writer = stage.write.session(spark).overwrite() + name match { + case "extraBody" => stage.replaceJsonDefault(stage.extraBody, Left(invalid)) + case "responseFormat" => stage.replaceJsonDefault(stage.responseFormat, Left(invalid)) + case "questions" => stage.replaceJsonDefault(stage.questions, Left(invalid)) + } + assert(intercept[IllegalArgumentException](writer.save(saved)).getMessage.contains("finite")) + assert(AIPrompt.read.session(spark).load(saved).getExtraBody == Map("keep" -> 1)) + } + } + + test("oversized native JSON fails before explicit default copied and deferred stage overwrites") { + val large = Map[String, Any]("description" -> ("x" * (32 * 1024 * 1024))) + val small = Map[String, Any]("keep" -> 1) + Seq("explicit" -> "extraBody", "default" -> "responseFormat", "shadowed" -> "questions", + "copied" -> "extraBody", "deferred" -> "responseFormat").foreach { case (mode, name) => + val saved = path() + new AIPrompt().setExtraBody(small).write.session(spark).save(saved) + val stage = new TestableAIMutableJsonDefaultPrompt(s"large-json-$mode") + val deferred = stage.write.session(spark).overwrite() + val param = name match { + case "extraBody" => stage.extraBody + case "responseFormat" => stage.responseFormat + case "questions" => stage.questions + } + if (mode == "default" || mode == "shadowed") { + stage.replaceJsonDefault(param, Left(large)) + if (mode == "shadowed") stage.set(param, Left(small)) + } else { + stage.set(param, Left(large)) + } + val candidate = if (mode == "copied") stage.copy(ParamMap.empty) else stage + val writer = if (mode == "deferred") deferred else candidate.write.session(spark).overwrite() + val error = intercept[IllegalArgumentException](writer.save(saved)) + assert(error.getMessage.contains("Spark JSON parsing limits") && error.getCause == null) + val restored = AIPrompt.read.session(spark).load(saved) + assert(restored.getExtraBody == small) + } + } + + test("large ordinary string params fail before replacing saved stages with unreadable metadata") { + val large = "x" * (32 * 1024 * 1024) + val saved = path() + new AIPrompt().setSystemPrompt("keep").write.session(spark).save(saved) + val stage = new AIPrompt().setSystemPrompt(large) + intercept[com.fasterxml.jackson.core.exc.StreamConstraintsException] { + stage.systemPrompt.jsonDecode(stage.systemPrompt.jsonEncode(large)) + } + val error = intercept[IllegalArgumentException] { + stage.write.session(spark).overwrite().save(saved) + } + assert(error.getMessage.contains("Spark JSON parsing limits") && error.getCause == null) + assert(AIPrompt.read.session(spark).load(saved).getSystemPrompt == "keep") + } + + Seq(false, true).foreach { fitted => + test(s"Spark pipeline writers reject oversized native JSON with fitted=$fitted") { + val large = Map[String, Any]("description" -> ("x" * (32 * 1024 * 1024))) + val stage = new AIPrompt().setExtraBody(large).setInputCol("text") + .setUrl("http://localhost/v1").setModel("model") + val pipeline = new Pipeline().setStages(Array(stage)) + val saved = path() + val error = intercept[IllegalArgumentException] { + if (fitted) pipeline.fit(input).write.session(spark).save(saved) + else pipeline.write.session(spark).save(saved) + } + assert(error.getMessage.contains("Spark JSON parsing limits") && error.getCause == null) + } + + test(s"Spark pipelines preserve supported native JSON and exact numbers with fitted=$fitted") { + val body = Map[String, Any]("description" -> ("x" * (32 * 1024)), "number" -> BigDecimal("1e400")) + val stage = new AIPrompt().setExtraBody(body).setInputCol("text") + .setUrl("http://localhost/v1").setModel("model") + val pipeline = new Pipeline().setStages(Array(stage)) + val saved = path() + val restored = (if (fitted) { + pipeline.fit(input).write.session(spark).save(saved) + PipelineModel.read.session(spark).load(saved).stages.head + } else { + pipeline.write.session(spark).save(saved) + Pipeline.read.session(spark).load(saved).getStages.head + }).asInstanceOf[AIPrompt] + assert(restored.getExtraBody == body && restored.uid == stage.uid) + assert(restored.isSet(restored.extraBody) && !restored.isSet(restored.temperature)) + } + } + + test("null ordinary and service params fail before direct deferred and copied stage overwrites") { + val saved = path() + val missing = Option.empty[String].orNull + new AIPrompt().setExtraBody(Map("keep" -> 1)).write.session(spark).save(saved) + Seq[() => AIPrompt](() => new AIPrompt(), () => new AISystemOne(), () => new OpenAIDecisions()) + .foreach { create => + val stage = create() + def verify[T](param: Param[T], value: T, setter: () => Unit): Unit = { + val deferred = stage.write.session(spark).overwrite() + val copied = stage.copy(ParamMap(param -> value)) + setter() + Seq(() => stage.write.session(spark).overwrite().save(saved), () => deferred.save(saved), + () => copied.write.session(spark).overwrite().save(saved)).foreach { save => + if (param == stage.user) intercept[IllegalArgumentException](save()) + else intercept[NullPointerException](save()) + assert(AIPrompt.read.session(spark).load(saved).getExtraBody == Map("keep" -> 1)) + } + stage.clear(param) + } + verify(stage.systemPrompt, missing, () => { stage.setSystemPrompt(missing); () }) + verify(stage.user, Left(missing), () => { stage.setUser(missing); () }) + } + } + + test("shadowed null ordinary and service defaults fail before overwriting a saved stage") { + val saved = path() + val missing = Option.empty[String].orNull + new AIPrompt().setExtraBody(Map("keep" -> 1)).write.session(spark).save(saved) + Seq("systemPrompt", "user").foreach { name => + val stage = new TestableAIMutableJsonDefaultPrompt("invalid-scalar-default") + .setSystemPrompt("valid").setUser("valid") + val writer = stage.write.session(spark).overwrite() + name match { + case "systemPrompt" => stage.replaceJsonDefault(stage.systemPrompt, missing) + case "user" => stage.replaceJsonDefault(stage.user, Left(missing)) + } + if (name == "user") intercept[IllegalArgumentException](writer.save(saved)) + else intercept[NullPointerException](writer.save(saved)) + assert(AIPrompt.read.session(spark).load(saved).getExtraBody == Map("keep" -> 1)) + } + } + + test("prompt template failures keep literal text out of logs and throwable chains") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val marker = "SYNTHETIC_PRIVATE_PROMPT_MARKER" + withLogs { logged => + Seq(s"$marker {", s"$marker {text", s"$marker {missing_column}").foreach { template => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setPromptTemplate(template) + val failure = intercept[IllegalArgumentException](stage.transform(input)) + val trace = new StringWriter() + failure.printStackTrace(new PrintWriter(trace)) + assert(!trace.toString.contains(marker), "Template text escaped through a throwable") + assert(failure.getCause == null && failure.getSuppressed.isEmpty) + } + assert(!logged.isEmpty, "The regression must observe the transform failure logger") + assert(logged.asScala.forall(!_.contains(marker)), "Template text escaped through a log event") + assert(server.requests.isEmpty) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMediaInputSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMediaInputSuite.scala new file mode 100644 index 00000000000..39916ca31a2 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMediaInputSuite.scala @@ -0,0 +1,265 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import spray.json._ + +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.util.Using + +class AIPromptMediaInputSuite extends AIPromptTestBase { + private val validImage = "data:image/png;base64,AQID" + private val malformedImages = Seq("", "data:image/png,AQID", "data:image/png") + private val chatReply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + private val responsesReply = """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + + private def mediaInput: DataFrame = { + import spark.implicits._ + (malformedImages :+ validImage).toDF("asset").withColumn("text", lit("describe")) + } + + private def messageInput(api: String): DataFrame = { + val partType = if (api == "chat_completions") "image_url" else "input_image" + mediaInput.withColumn("messages", array(struct(lit("user").as("role"), + array(map(lit("type"), lit(partType), lit("image_url"), col("asset"))).as("content")))) + } + + private def verifyMedia(api: String, input: DataFrame, expectedCodes: Seq[String], sends: Int) + (configure: AIPrompt => AIPrompt): Unit = { + val reply = api match { + case "decisions" => NativeRequestFixtures.DecisionResponse + case "responses" => responsesReply + case _ => chatReply + } + Using.resource(new AIRecordingServer((_, body) => { + if (body.contains(validImage)) 200 -> reply + else 400 -> """{"error":{"code":"invalid_media","message":"synthetic malformed media"}}""" + })) { server => + val stage = configure(new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setOutputCol("answer").setErrorCol("error")) + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect() + val byAsset = rows.map(row => row.getAs[String]("asset") -> row).toMap + assert(rows.length == 4 && byAsset.size == 4) + malformedImages.zip(expectedCodes).foreach { case (asset, code) => + val error = byAsset(asset).getAs[Row]("error") + assert(error != null && AIJson.objectString(error.getString(0)).fields("code") == JsString(code)) + } + assert(byAsset(validImage).getAs[Row]("error") == null && byAsset(validImage).getAs[Any]("answer") != null) + assert(server.requests.size() == sends) + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"malformed inline $api paths retain local row errors") { + verifyMedia(api, mediaInput, Seq("missing_input", "invalid_request", "invalid_request"), sends = 1) { + _.setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")) + } + } + + test(s"malformed inline $api messages retain local and provider row errors") { + verifyMedia(api, messageInput(api), Seq("invalid_request", "http_error", "http_error"), sends = 3) { + _.setMessagesCol("messages") + } + } + + test(s"$api attachments follow input column order through copies and persistence") { + import spark.implicits._ + val left = validImage + val right = "data:image/png;base64,BAUG" + val data = Seq((left, right)).toDF("left", "right").select("right", "left") + val reply = if (api == "chat_completions") chatReply else responsesReply + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setApiType(api) + .setOutputCol("answer").setErrorCol("error") + .setPromptTemplate("Compare {left} and {`RiGhT`}; literal {{LEFT}}") + .setColumnTypes(ListMap("LEFT" -> "path", "RIGHT" -> "path")) + val reordered = stage.copy(ParamMap.empty).setColumnTypes(ListMap("RIGHT" -> "path", "LEFT" -> "path")) + val saved = path() + reordered.write.session(spark).save(saved) + Seq(stage, reordered, AIPrompt.load(saved)).foreach { configured => + assert(configured.transform(data).head().getAs[String]("answer") == "answer") + } + assert(server.requests.size() == 3) + server.requests.asScala.foreach { case (_, body) => + assert(body.indexOf(right) >= 0 && body.indexOf(right) < body.indexOf(left)) + assert(body.contains("Compare attachment.png and attachment.png; literal {LEFT}")) + } + } + } + } + + test("malformed inline Decisions JSON rejects invalid rows without sending") { + val input = messageInput("decisions").withColumn("input", to_json(col("messages"))).drop("messages") + verifyMedia("decisions", input, Seq.fill(3)("invalid_request"), sends = 1) { + _.setInputCol("input").setInputMode("json").setQuestions(NativeRequestFixtures.DecisionQuestions) + } + } + + test("Decisions image URL casing and MIME parameters preserve payloads") { + import spark.implicits._ + val accepted = Seq(validImage, "DATA:image/png;base64,AQID", "data:IMAGE/PNG;base64,AQID", + "data:image/png;BASE64,AQID", "data:image/png;name=a.png;base64,AQID", + "DATA:IMAGE/PNG;NAME=a%20b.png;charset=UTF-8;BASE64,AQID") + val rejected = Seq("data:text/plain;base64,AQID", "data:image/png;name;base64,AQID", + "data:image/png;=a.png;base64,AQID", "data:image/png;name=bad%xy;base64,AQID", + "data:image/png;base64;name=a.png,AQID", "data:image/png;name=a.png;base64,@@@", + "data:image/png;name=a.png;base64,", "data:image/png;name=a.png,AQID") + val messages = (accepted ++ rejected).toDF("asset").withColumn("messages", array(struct( + lit("user").as("role"), array(map(lit("type"), lit("input_image"), + lit("image_url"), col("asset"))).as("content")))) + for (json <- Seq(false, true); stage <- Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions())) { + val input = if (json) messages.withColumn("input", to_json(col("messages"))).drop("messages") else messages + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + stage.setUrl(server.url + "/v1").setModel("model").setQuestions(NativeRequestFixtures.DecisionQuestions) + .setOutputCol("answer").setErrorCol("error") + if (json) stage.setInputCol("input").setInputMode("json") else stage.setMessagesCol("messages") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect() + val byAsset = rows.map(row => row.getAs[String]("asset") -> row).toMap + assert(rows.length == accepted.size + rejected.size && byAsset.size == rows.length) + accepted.foreach { asset => + assert(byAsset(asset).getAs[Row]("error") == null && byAsset(asset).getAs[Row]("answer") != null, asset) + } + rejected.foreach { asset => + assert(byAsset(asset).getAs[Row]("error").getString(0).contains("invalid_request"), asset) + assert(byAsset(asset).getAs[Row]("answer") == null) + } + assert(server.requests.size() == accepted.size) + } + } + } + + test("Decisions null image options normalize through direct and facade stages") { + import spark.implicits._ + Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions()).foreach { stage => + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val image = Map("type" -> JsString("input_image"), "image_url" -> JsString(validImage), + "provider_extension" -> JsNull) + val variants = Seq(Map("file_id" -> JsNull), Map("detail" -> JsNull), + Map("file_id" -> JsNull, "detail" -> JsNull)) + val data = variants.map { fields => + JsArray(JsObject("role" -> JsString("user"), + "content" -> JsArray(JsObject(image ++ fields)))).compactPrint + }.toDF("text") + stage.setUrl(server.url + "/v1").setModel("model").setInputCol("text").setInputMode("json") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setOutputCol("answer").setErrorCol("error") + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val rows = output.collect() + assert(rows.length == 3 && rows.forall(row => + row.getAs[Row]("error") == null && row.getAs[Row]("answer") != null)) + assert(server.requests.size() == 3) + server.requests.asScala.foreach { case (_, body) => + val message = AIJson.objectString(body).fields("input").asInstanceOf[JsArray].elements.head.asJsObject + val part = message.fields("content").asInstanceOf[JsArray].elements.head.asJsObject + assert(!part.fields.contains("file_id") && !part.fields.contains("detail")) + assert(part.fields("provider_extension") == JsNull) + } + } + } + } + + test("attachment wildcard templates use only public columns") { + val data = spark.range(3).select(col("id").cast(StringType).as("ticket"), + lit(validImage).as("asset"), lit("user-owned").as("asset_filename")) + val expected = data.toJSON.collect().toVector.sorted + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset};{to_json(struct(*))}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + val rows = result.collect() + assert(result.schema == stage.transformSchema(data.schema)) + assert(rows.length == 3 && rows.forall(_.getAs[Row]("error") == null)) + val contents = server.requests.asScala.toVector.map { case (_, body) => + val parts = AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsArray].elements + val texts = parts.flatMap(_.asJsObject.fields.get("text")).collect { + case JsString(text) if text.startsWith("attachment.png;") => text.stripPrefix("attachment.png;") + } + assert(texts.size == 1) + texts.head + }.sorted + assert(contents == expected && server.requests.size() == 3) + } + } + + test("path normalization preserves missing-input errors and inline attachments") { + import spark.implicits._ + val data = Seq(None, Some(""), Some(" \t "), Some("data:text/plain;base64,aGVsbG8="), + Some(" \tDATA:text/plain;BASE64,aGVsbG8= \t")).toDF("asset") + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val rows = stage.transform(data).collect() + assert(rows.length == 5) + val errors = rows.flatMap(row => Option(row.getAs[Row]("error"))) + assert(errors.length == 3) + errors.foreach { row => + val error = AIJson.objectString(row.getString(0)) + assert(error.fields("code") == JsString("missing_input")) + assert(error.fields("message") == JsString("attachment input is missing")) + } + assert(server.requests.size() == 2) + server.requests.asScala.foreach { case (_, body) => + val message = AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last.asJsObject + assert(message.fields("content").asInstanceOf[JsString].value.takeWhile(_ != '\n') == "attachment.txt") + } + } + } + + test("large inline messages remain valid") { + val inline = "data:image/png;base64," + ("A" * 20000004) + val message = new StructType().add("role", StringType) + .add("content", ArrayType(MapType(StringType, StringType))) + val data = spark.createDataFrame(Collections.singletonList( + Row(Seq(Row("user", Seq(Map("type" -> "input_image", "image_url" -> inline)))))), + new StructType().add("messages", ArrayType(message))) + Using.resource(new AIRecordingServer((_, _) => 200 -> responsesReply)) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + val row = stage.transform(data).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.requests.peek()._2.contains(inline)) + } + } + + test("Responses forwards literal file_data and preserves provider validation failures") { + val raw = "JVBERi0xLjQK" + val inline = "data:application/pdf;base64," + raw + val contents = Seq(raw, inline, "not base64", "https://example.invalid/not-a-location") + val messageType = new StructType().add("role", StringType) + .add("content", ArrayType(MapType(StringType, StringType))) + val input = spark.createDataFrame(contents.map { text => + Row(text, Seq(Row("user", Seq(Map("type" -> "input_file", "filename" -> "note.pdf", "file_data" -> text))))) + }.asJava, new StructType().add("text", StringType).add("messages", ArrayType(messageType))).coalesce(1) + def fileData(body: String): String = { + val message = AIJson.objectString(body).fields("input").asInstanceOf[JsArray].elements.head.asJsObject + message.fields("content").asInstanceOf[JsArray].elements.head.asJsObject + .fields("file_data").asInstanceOf[JsString].value + } + Using.resource(new AIRecordingServer((_, body) => { + if (Set(raw, inline)(fileData(body))) 200 -> responsesReply + else 400 -> """{"error":{"code":"invalid_file","message":"Unsupported inline file content"}}""" + })) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + val rows = stage.transform(input).collect() + assert(rows.length == 4 && rows.count(_.getAs[Row]("error") != null) == 2) + assert(server.requests.asScala.map(entry => fileData(entry._2)).toSeq.sorted == contents.sorted) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMessageSchemaSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMessageSchemaSuite.scala new file mode 100644 index 00000000000..996494b9c10 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptMessageSchemaSuite.scala @@ -0,0 +1,63 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.sql.Row +import org.apache.spark.sql.types._ + +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptMessageSchemaSuite extends TestBase { + private val chatReply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + + test("duplicate protocol message fields fail schema planning even for empty inputs") { + def message(content: DataType): StructType = + new StructType().add("role", StringType).add("content", content) + val duplicateText = new StructType().add("type", StringType) + .add("text", StringType).add("text", StringType) + val duplicateUrl = new StructType().add("url", StringType).add("url", StringType) + val schemas = Seq( + message(StringType).add("content", StringType), + message(StringType).add("role", StringType), + message(ArrayType(duplicateText)), + message(ArrayType(new StructType().add("type", StringType).add("image_url", duplicateUrl))), + message(MapType(StringType, ArrayType(duplicateText)))) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stages = Seq(new AIPrompt(), new AIPrompt().setApiType("responses"), + new AIPrompt().setApiType("decisions").setQuestions(NativeRequestFixtures.DecisionQuestions), + new OpenAIDecisions().setQuestions(NativeRequestFixtures.DecisionQuestions)) + stages.foreach { stage => + stage.setUrl(server.url + "/v1").setModel("model").setMessagesCol("messages") + schemas.foreach { fields => + val schema = new StructType().add("messages", ArrayType(fields)) + val data = spark.createDataFrame(java.util.Collections.emptyList[Row](), schema) + assert(intercept[IllegalArgumentException](stage.transformSchema(schema)) + .getMessage.contains("duplicate")) + assert(intercept[IllegalArgumentException](stage.transform(data)).getMessage.contains("duplicate")) + } + } + assert(server.requests.isEmpty) + } + } + + test("message validation leaves duplicate fields in unrelated retained structs unchanged") { + val message = new StructType().add("role", StringType).add("content", StringType) + val retained = new StructType().add("text", StringType).add("text", StringType) + val data = spark.createDataFrame(Seq(Row(Seq(Row("user", "hello")), Row("one", "two"))).asJava, + new StructType().add("messages", ArrayType(message)).add("retained", retained)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setMessagesCol("messages") + .setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + assert(server.requests.isEmpty && result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("retained") == Row("one", "two")) + assert(server.requests.size() == 1) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPersistenceSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPersistenceSuite.scala new file mode 100644 index 00000000000..bbd5d0ad0ad --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPersistenceSuite.scala @@ -0,0 +1,129 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.OpenAIApiVersion +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.functions.lit +import spray.json.{JsNumber, JsObject, JsString} + +import scala.util.Using + +class AIPromptPersistenceSuite extends AIPromptTestBase { + test("public requests preserve rendered JSON number kinds") { + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val data = spark.range(1).select(lit("hello").as("text")) + val cases: Seq[(Either[String, Double], String)] = + Seq(Left("1e0") -> "1", Left("3.2e1") -> "32", Right(12345678.0) -> "12345678") + cases.foreach { case (value, rendered) => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + value.fold(text => stage.setExtraBody(s"""{"x":$text}"""), number => stage.setExtraBody(Map("x" -> number))) + assert(stage.transform(data).head().getAs[String]("answer") == "answer") + assert(AIJson.objectString(server.requests.poll()._2).fields("x").compactPrint == rendered) + } + assert(server.requests.isEmpty) + } + } + + test("numeric JSON setters reject undecodable rendered exponents without changing parameters") { + val cases = Seq("100e2147483647" -> false, "-100e2147483647" -> false, "100e2147483646" -> false, + "100e2147483645" -> true, "-100e2147483645" -> true, "0e2147483647" -> true, "1e-2147483647" -> true) + cases.foreach { case (token, supported) => + val value = BigDecimal(token) + val typed = JsObject("x" -> JsNumber(value)) + val text = s"""{"x":$token}""" + val stage = new AIPrompt().setExtraBody("""{"safe":1}""").setQuestions("""{"safe":1}""") + .setResponseSchema(Map("type" -> "object")) + val original = stage.extractParamMap().toSeq + val setters: Seq[() => AIPrompt] = Seq( + () => stage.setExtraBody(text), () => stage.setExtraBody(typed), () => stage.setExtraBody(Map("x" -> value)), + () => stage.setQuestions(text), () => stage.setQuestions(typed), () => stage.setQuestions(Map("x" -> value)), + () => stage.setResponseSchema(Map("type" -> "object", "minimum" -> value))) + setters.foreach { set => + if (supported) set() else { + intercept[IllegalArgumentException](set()) + assert(stage.extractParamMap().toSeq.toSet == original.toSet) + } + } + if (supported) { + assert(AIJson.parse(AIJson.parse(token).compactPrint) == JsNumber(value)) + assert(stage.extraBody.jsonDecode(stage.extraBody.jsonEncode(Left(stage.getExtraBody))) == + Left(stage.getExtraBody)) + assert(stage.questions.jsonDecode(stage.questions.jsonEncode(Left(stage.getQuestions))) == + Left(stage.getQuestions)) + assert(stage.responseFormat.jsonDecode(stage.responseFormat.jsonEncode(Left(stage.getResponseFormat))) == + Left(stage.getResponseFormat)) + } else intercept[IllegalArgumentException](AIJson.parse(token)) + } + } + + test("typed readers do not initialize incompatible metadata-selected classes") { + import spark.implicits._ + val initialized = "synapseml.aiprompt.test.incompatible.initialized" + val constructed = "synapseml.aiprompt.test.incompatible.constructed" + val before = Seq(initialized, constructed).map(key => key -> Option(System.getProperty(key))) + assert(before.forall(_._2.isEmpty)) + try { + val saved = path() + new AIPrompt().write.session(spark).save(saved) + val metadataPath = new Path(saved, "metadata").toString + val metadata = AIJson.objectString(spark.read.text(metadataPath).head().getString(0)) + val changed = JsObject(metadata.fields.updated("class", + JsString("synapsemltest.ai.AIIncompatibleStageProbe"))).compactPrint + Seq(changed).toDF("value").write.mode("overwrite").text(metadataPath) + val readers = Seq(() => AIPrompt.load(saved), () => AISystemOne.load(saved), () => OpenAIDecisions.load(saved)) + readers.foreach { load => + val error = intercept[IllegalArgumentException](load()) + assert(error.getMessage.contains("incompatible") && error.getCause == null) + } + assert(System.getProperty(initialized) == null && System.getProperty(constructed) == null) + } finally { + before.foreach { + case (key, Some(value)) => System.setProperty(key, value) + case (key, None) => System.clearProperty(key) + } + } + } + + test("dated fallback is an effective default while original and v1 native stages remain unmodified") { + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + try { + GlobalParams.resetGlobalParam(OpenAIApiVersion) + val dated = new AIPrompt().setUrl("https://example.invalid").setDeploymentName("model").setInputCol("text") + val effective = dated.effectiveCopy() + assert(!dated.isDefined(dated.apiVersion) && !effective.isSet(effective.apiVersion)) + assert(effective.getDefault(effective.apiVersion).contains(Left("2025-04-01-preview"))) + Seq(new AIPrompt().setUrl("https://example.invalid/v1"), + new AISystemOne().setUrl("https://example.invalid"), + new OpenAIDecisions().setUrl("https://example.invalid/v1")).foreach { stage => + assert(!stage.effectiveCopy().isDefined(stage.apiVersion)) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("leaf readers reject incompatible saved classes while the facade accepts its subclasses") { + Seq(new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { original => + val saved = path() + original.write.session(spark).save(saved) + assert(AIPrompt.load(saved).getClass == original.getClass) + if (!original.isInstanceOf[AISystemOne]) { + val error = intercept[IllegalArgumentException] { AISystemOne.load(saved); () } + assert(error.getMessage.contains("incompatible") && error.getCause == null) + } + if (!original.isInstanceOf[OpenAIDecisions]) { + val error = intercept[IllegalArgumentException] { OpenAIDecisions.load(saved); () } + assert(error.getMessage.contains("incompatible") && error.getCause == null) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPostProcessingSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPostProcessingSuite.scala new file mode 100644 index 00000000000..51647907d17 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPostProcessingSuite.scala @@ -0,0 +1,46 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{StringType, StructType} + +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptPostProcessingSuite extends AIPromptTestBase { + Seq("csv", "regex").foreach { mode => + test(s"lazy $mode postprocessing keeps planned schemas and values") { + Seq("chat_completions", "responses").foreach { api => + val reply = if (api == "chat_completions") { + """{"choices":[{"index":0,"message":{"role":"assistant","content":"one;two"},"finish_reason":"stop"}]}""" + } else { + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"one;two"}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val options = if (mode == "csv") Map("delimiter" -> ";") else { + Map("regex" -> "one;(two)", "regexGroup" -> "1") + } + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setSubscriptionKey("synthetic-key").setInputCol("text").setOutputCol("answer") + .setErrorCol("error").setRawResponseCol("raw").setPostProcessingOptions(options) + val data = spark.createDataFrame(Seq(Row("hello"), Row(Option.empty[String].orNull)).asJava, + new StructType().add("text", StringType)) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + val answered = rows.find(_.getAs[String]("text") == "hello").get + if (mode == "csv") assert(answered.getAs[Seq[String]]("answer") == Seq("one", "two")) + else assert(answered.getAs[String]("answer") == "two") + assert(answered.getAs[Row]("error") == null) + val missing = rows.find(_.getAs[String]("text") == null).get + assert(missing.isNullAt(missing.fieldIndex("answer"))) + assert(missing.getAs[Row]("error").getString(0).contains("missing_input")) + assert(server.requests.size() == 1) + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPreflightSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPreflightSuite.scala new file mode 100644 index 00000000000..2929db33273 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPreflightSuite.scala @@ -0,0 +1,778 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.contracts.HasInputCol +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.{HTTPSchema, URLKey} +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.OpenAIApiVersion +import com.microsoft.azure.synapse.ml.services.aifoundry.AIFoundryModel +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIDecisions, OpenAIDeploymentNameKey, + OpenAIPrompt, OpenAITemperatureKey} +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.types.{ArrayType, StringType, StructType} +import spray.json._ + +import java.nio.file.Files +import java.util.concurrent.{Callable, CountDownLatch, ExecutionException, Executors, TimeUnit} +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptPreflightSuite extends TestBase { + private val chatReply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + + private def input: DataFrame = spark.createDataFrame( + Seq(Row("hello")).asJava, new StructType().add("text", StringType)) + + test("rendered template types must be strings before lazy inference") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setErrorCol("error") + Seq("{array(text)}", "{encode(text, 'UTF-8')}").foreach { template => + stage.setPromptTemplate(template) + val schemaError = intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + val transformError = intercept[IllegalArgumentException](stage.transform(input)) + Seq(schemaError, transformError).foreach { error => + assert(error.getMessage.contains("Prompt templates must produce strings") && error.getCause == null) + } + assert(server.requests.isEmpty) + } + stage.setPromptTemplate("{cast(array(text) as string)}") + assert(stage.transform(input).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1 && server.requests.peek()._2.contains("[hello]")) + } + } + + test("column type validation cannot be bypassed by setting the Param directly") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setPromptTemplate("{text}") + Seq("image", "Path", Option.empty[String].orNull).foreach { kind => + stage.set(stage.columnTypes, Map("text" -> kind)) + intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + intercept[IllegalArgumentException](stage.transform(input)) + assert(server.requests.isEmpty) + } + } + } + + Seq("chat_completions", "responses", "decisions").foreach { api => + test(s"$api explicit messages reject unused construction settings before HTTP") { + val data = input.selectExpr("text", "array(named_struct('role', 'user', 'content', text)) AS messages") + val settings: Seq[(String, AIPrompt => AIPrompt)] = if (api == "decisions") { + Seq("inputMode" -> (_.setInputMode("text")), "inputMode" -> (_.setInputMode("json"))) + } else Seq("systemPrompt" -> (_.setSystemPrompt("Do not ignore this")), + "columnTypes" -> (_.setColumnTypes(Map("text" -> "path"))), + "columnTypes" -> (_.setColumnTypes(Map("text" -> "text"))), + "fileSizeLimitMB" -> (_.setFileSizeLimitMB(1.0)), + "dropPrompt" -> (_.setDropPrompt(false)), "dropPrompt" -> (_.setDropPrompt(true))) + val stages = if (api == "decisions") Seq(new AIPrompt().setApiType(api), new OpenAIDecisions()) + else Seq(new AIPrompt().setApiType(api)) + val reply = if (api == "responses") { + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + } else chatReply + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + stages.foreach { base => + base.setUrl(server.url + "/v1").setModel("model").setMessagesCol("messages").setErrorCol("error") + if (api == "decisions") base.setQuestions(NativeRequestFixtures.DecisionQuestions) + settings.foreach { case (name, configure) => + val stage = configure(base.copy(ParamMap.empty)) + assert(intercept[IllegalArgumentException](stage.transformSchema(data.schema)).getMessage.contains(name)) + assert(intercept[IllegalArgumentException](stage.transform(data)).getMessage.contains(name)) + assert(server.requests.isEmpty) + } + if (api != "decisions") { + assert(base.setColumnTypes(Map.empty[String, String]).transform(data).head().getAs[Row]("error") == null) + assert(server.requests.size() == 1) + server.requests.clear() + } + } + } + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"$api filename interpolation resolves column aliases without rewriting escaped literals") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + val reply = if (api == "chat_completions") chatReply else { + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}]}""" + } + try { + spark.conf.set("spark.sql.caseSensitive", "false") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val file = "data:text/plain;base64,aGVsbG8=" + val data = spark.createDataFrame(Seq(Row(file, file, file, "keep")).asJava, + new StructType().add("aSsEt", StringType).add("file.name", StringType) + .add("file`name", StringType).add("tag", StringType)) + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("Inspect {asset} / {`FILE.NAME`} / {`file``NAME`}; " + + "literal {{ASSET}}; SQL {'asset'}; expression {upper(tag)}") + .setColumnTypes(Map("ASSET" -> "path", "FiLe.NaMe" -> "path", "FILE`NAME" -> "path")) + .setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.peek()._2.contains( + "Inspect attachment.txt / attachment.txt / attachment.txt; literal {ASSET}; SQL asset; expression KEEP")) + assert(row.getAs[String]("aSsEt") == file && row.getAs[String]("file`name") == file) + spark.conf.set("spark.sql.caseSensitive", "true") + intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + assert(server.requests.size() == 1) + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + } + + test("saved stages preserve explicitly disabled postprocessing") { + val session = spark + Seq(new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + val root = Files.createTempDirectory("aiprompt-postprocessing") + try { + stage.setPostProcessing("") + stage.write.session(session).save(root.resolve("explicit").toString) + val loaded = AIPrompt.read.session(session).load(root.resolve("explicit").toString) + Seq(stage, loaded, loaded.copy(org.apache.spark.ml.param.ParamMap.empty)).foreach { candidate => + intercept[IllegalArgumentException](candidate.setPostProcessingOptions(Map("delimiter" -> ";"))) + } + loaded.clear(loaded.postProcessing) + loaded.setPostProcessingOptions(Map("delimiter" -> ";")) + assert(loaded.getPostProcessing == "csv") + val implicitStage = new AIPrompt() + implicitStage.write.session(session).save(root.resolve("implicit").toString) + val implicitLoaded = AIPrompt.read.session(session).load(root.resolve("implicit").toString) + assert(!implicitLoaded.isSet(implicitLoaded.postProcessing)) + implicitLoaded.setPostProcessingOptions(Map("delimiter" -> ";")) + assert(implicitLoaded.getPostProcessing == "csv") + } finally org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } + + test("Responses incomplete details preserve partial text and diagnostics without a status") { + Seq("", """"status":"completed",""").foreach { status => + val reply = s"""{$status"incomplete_details":{"reason":"max_output_tokens"}, + |"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"partial"}]}], + |"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setUsageCol("usage") + val row = stage.transform(input).head() + assert(row.getAs[String]("answer") == "partial") + assert(Option(row.getAs[Row]("error")).exists(_.getString(0).contains("incomplete"))) + assert(row.getAs[Row]("usage").getLong(2) == 3L && row.getAs[String]("raw") == reply) + } + } + } + + test("Foundry routes use the same implicit model binding as the request body") { + val stage = new AIPrompt().setUrl("https://example.services.ai.azure.com").setInputCol("text") + val previousModel = GlobalParams.getGlobalParam(AIFoundryModel) + val previousDeployment = GlobalParams.getGlobalParam(OpenAIDeploymentNameKey) + def route(candidate: AIPrompt, expectedPath: String, expectedModel: String): Unit = { + candidate.transformSchema(input.schema) + val adapter = new AIRequestAdapter(candidate.effectiveCopy(), input.schema) + val (url, body) = adapter.request(Row("hello"), None, + Some(Seq(new org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema( + Array[Any]("user", "hello"), new StructType().add("role", StringType).add("content", StringType))))) + assert(new java.net.URI(url).getPath == expectedPath) + assert(AIJson.objectString(body).fields("model") == JsString(expectedModel)) + } + try { + GlobalParams.setGlobalParam(AIFoundryModel, Left("global-model")) + GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, Left("global-deployment")) + route(stage, "/openai/deployments/global-deployment/chat/completions", "global-deployment") + route(stage.copy(org.apache.spark.ml.param.ParamMap.empty).setModel("explicit-model"), + "/models/chat/completions", "explicit-model") + route(stage.copy(org.apache.spark.ml.param.ParamMap.empty).setDeploymentName("explicit-deployment"), + "/openai/deployments/explicit-deployment/chat/completions", "explicit-deployment") + } finally { + previousModel match { + case Some(value) => GlobalParams.setGlobalParam(AIFoundryModel, value) + case None => GlobalParams.resetGlobalParam(AIFoundryModel) + } + previousDeployment match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, value) + case None => GlobalParams.resetGlobalParam(OpenAIDeploymentNameKey) + } + } + } + + Seq("systemone", "decisions", "embeddings").foreach { api => + test(s"$api ignores inactive global temperature in schema and requests") { + val previous = GlobalParams.getGlobalParam(OpenAITemperatureKey) + val reply = api match { + case "systemone" => NativeRequestFixtures.SystemResponse + case "decisions" => NativeRequestFixtures.DecisionResponse + case _ => """{"data":[{"index":0,"embedding":[0.25,0.5]}],"model":"model"}""" + } + try { + GlobalParams.resetGlobalParam(OpenAITemperatureKey) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = (api match { + case "systemone" => new AISystemOne().setStateCol("text") + .setQuestions(NativeRequestFixtures.SystemQuestions) + case "decisions" => new OpenAIDecisions().setInputCol("text") + .setQuestions(NativeRequestFixtures.DecisionQuestions) + case _ => new AIPrompt().setApiType(api).setInputCol("text") + }).setUrl(if (api == "systemone") server.url else server.url + "/v1") + .setModel("model").setOutputCol("answer").setErrorCol("error") + val original = stage.effectiveCopy() + val expectedSchema = stage.transformSchema(input.schema) + val defaults = Seq[Either[Double, String]](Right("absent-global-column"), Right("text"), Left(0.25)) + defaults.foreach { default => + GlobalParams.setGlobalParam(OpenAITemperatureKey, default) + assert(stage.transformSchema(input.schema) == expectedSchema) + val row = stage.transform(input).head() + assert(row.getAs[Row]("error") == null && !row.isNullAt(row.fieldIndex("answer"))) + assert(!AIJson.objectString(server.requests.asScala.last._2).fields.contains("temperature")) + assert(GlobalParams.getGlobalParam(OpenAITemperatureKey).contains(default)) + } + val requestsBefore = server.requests.size() + Seq(stage.copy(ParamMap.empty).setTemperature(0.25), + stage.copy(ParamMap.empty).setTemperatureCol("text")).foreach { explicit => + intercept[IllegalArgumentException](explicit.transformSchema(input.schema)) + } + assert(server.requests.size() == requestsBefore) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAITemperatureKey, value) + case None => GlobalParams.resetGlobalParam(OpenAITemperatureKey) + } + } + } + } + + test("native stages require explicit models despite global chat and Foundry defaults") { + val previousModel = GlobalParams.getGlobalParam(AIFoundryModel) + val previousDeployment = GlobalParams.getGlobalParam(OpenAIDeploymentNameKey) + try { + Seq("systemone", "decisions").foreach { api => + val questions = if (api == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions + val direct = if (api == "systemone") new AISystemOne() else new OpenAIDecisions() + Seq(direct, new AIPrompt().setApiType(api)).foreach { stage => + stage.setUrl(if (api == "systemone") "https://example.services.ai.azure.com" + else "https://example.invalid/v1").setInputCol("text").setQuestions(questions) + GlobalParams.resetGlobalParam(AIFoundryModel) + GlobalParams.resetGlobalParam(OpenAIDeploymentNameKey) + val explicit = stage.copy(ParamMap.empty).setModel("native-model") + val schema = explicit.transformSchema(input.schema) + Seq[Either[String, String]](Left("chat-model"), Right("absent-model-column")).foreach { value => + GlobalParams.setGlobalParam(AIFoundryModel, value) + GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, value) + val error = intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + assert(error.getMessage.contains("Set deploymentName/model")) + assert(explicit.transformSchema(input.schema) == schema) + assert(GlobalParams.getGlobalParam(AIFoundryModel).contains(value)) + } + } + } + } finally { + previousModel match { + case Some(value) => GlobalParams.setGlobalParam(AIFoundryModel, value) + case None => GlobalParams.resetGlobalParam(AIFoundryModel) + } + previousDeployment match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, value) + case None => GlobalParams.resetGlobalParam(OpenAIDeploymentNameKey) + } + } + } + + test("explicit message rows keep their own schema when intermediate encoders are materialized") { + val session = spark.newSession() + session.conf.set("spark.sql.optimizer.excludedRules", + "org.apache.spark.sql.catalyst.optimizer.EliminateSerialization") + Seq(("chat_completions", false), ("chat_completions", true), ("responses", false)).foreach { + case (api, withName) => + val base = new StructType().add("role", StringType).add("content", StringType) + val message = if (withName) base.add("name", StringType) else base + val schema = new StructType().add("messages", ArrayType(message)) + val value = if (withName) Row("user", "hello", "speaker") else Row("user", "hello") + val data = session.createDataFrame(Seq(Row(Seq(value))).asJava, schema) + val reply = if (api == "chat_completions") chatReply else { + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(schema)) + val row = result.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(row.getSeq[Row](row.fieldIndex("messages")) == Seq(value)) + assert(server.requests.size() == 1) + assert(server.requests.peek()._2.contains("hello")) + } + } + } + + Seq("systemone", "decisions", "chat_completions", "responses", "embeddings", "custom", "global-v1") + .foreach { route => + test(s"global API versions are inactive for $route in schema and requests") { + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + val previousUrl = GlobalParams.getGlobalParam(URLKey) + val data = input.withColumn("wrongVersion", org.apache.spark.sql.functions.lit(1)) + val api = if (Set("custom", "global-v1")(route)) "chat_completions" else route + val reply = api match { + case "systemone" => NativeRequestFixtures.SystemResponse + case "decisions" => NativeRequestFixtures.DecisionResponse + case "embeddings" => """{"data":[{"index":0,"embedding":[0.25,0.5]}],"model":"model"}""" + case "responses" => + """{"output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + case _ => chatReply + } + try { + GlobalParams.resetGlobalParam(OpenAIApiVersion) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = (api match { + case "systemone" => new AISystemOne().setStateCol("text") + .setQuestions(NativeRequestFixtures.SystemQuestions) + case "decisions" => new OpenAIDecisions().setInputCol("text") + .setQuestions(NativeRequestFixtures.DecisionQuestions) + case _ => new AIPrompt().setApiType(api).setInputCol("text") + }).setModel("model").setOutputCol("answer").setErrorCol("error") + route match { + case "systemone" => stage.setUrl(server.url + "/providers/microsoft/v1/systemone") + case "custom" => stage.setCustomUrlRoot(server.url + "/chat/completions?api-version=owned") + case "global-v1" => GlobalParams.setGlobalParam(URLKey, server.url + "/v1/chat/completions") + case _ => stage.setUrl(server.url + "/v1") + } + val expectedSchema = stage.transformSchema(data.schema) + val original = stage.transform(data).head() + assert(original.getAs[Row]("error") == null && !original.isNullAt(original.fieldIndex("answer"))) + val expectedRequest = server.requests.asScala.last + val defaults = Seq[Either[String, String]]( + Right("absentVersion"), Right("wrongVersion"), Left("dated-one"), Left("dated-two")) + defaults.foreach { default => + GlobalParams.setGlobalParam(OpenAIApiVersion, default) + assert(stage.transformSchema(data.schema) == expectedSchema) + val row = stage.transform(data).head() + assert(row.getAs[Row]("error") == null && !row.isNullAt(row.fieldIndex("answer"))) + assert(server.requests.asScala.last == expectedRequest) + assert(GlobalParams.getGlobalParam(OpenAIApiVersion).contains(default)) + } + Seq(stage.copy(ParamMap.empty).setApiVersion("explicit"), + stage.copy(ParamMap.empty).setApiVersionCol("text")).foreach { explicit => + val error = intercept[IllegalArgumentException](explicit.transformSchema(data.schema)) + assert(error.getMessage.contains("apiVersion is not supported")) + if (route == "custom") assert(error.getMessage.contains("customUrlRoot")) + val transformError = intercept[IllegalArgumentException](explicit.transform(data)) + assert(transformError.getMessage == error.getMessage) + } + assert(server.requests.size() == 5) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + previousUrl match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + } + + test("global API versions remain active for dated routes") { + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + val data = input.withColumn("version", org.apache.spark.sql.functions.lit("dated-column")) + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url).setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + GlobalParams.setGlobalParam(OpenAIApiVersion, Right("absentVersion")) + intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + GlobalParams.setGlobalParam(OpenAIApiVersion, Right("version")) + val columnVersion = stage.transform(data).head() + assert(columnVersion.getAs[String]("answer") == "answer" && columnVersion.getAs[Row]("error") == null) + assert(server.requests.asScala.last._1.endsWith("?api-version=dated-column")) + GlobalParams.setGlobalParam(OpenAIApiVersion, Left("dated-one")) + val scalarVersion = stage.transform(data).head() + assert(scalarVersion.getAs[String]("answer") == "answer" && scalarVersion.getAs[Row]("error") == null) + assert(server.requests.asScala.last._1.endsWith("?api-version=dated-one")) + GlobalParams.setGlobalParam(OpenAIApiVersion, Left("dated-two")) + assert(stage.transform(data).head().getAs[String]("answer") == "answer") + assert(server.requests.asScala.last._1.endsWith("?api-version=dated-two")) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("public input binding follows the repository contract without introducing a default") { + Seq(new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + assert(classOf[HasInputCol].isInstance(stage)) + assert(stage.getDefault(stage.inputCol).isEmpty) + assert(stage.setInputCol("text").getInputCol == "text") + } + } + + test("API selector documentation matches the accepted protocol domain on every public stage") { + val protocols = Seq("chat_completions", "responses", "systemone", "decisions", "embeddings") + Seq(new OpenAIPrompt(), new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + protocols.foreach { protocol => + assert(stage.apiType.doc.contains(s"'$protocol'") == stage.apiType.isValid(protocol)) + } + } + assert(new OpenAIPrompt().apiType.doc == + "The OpenAI API type to use: 'chat_completions' or 'responses'") + } + + test("Python template hooks validate JSON replacements and render class-specific provenance") { + val setter = "value = self._service_param_value_to_java(value)" + Seq("extraBody", "questions").foreach { name => + Seq("", setter + "\n" + setter).foreach { source => + assert(intercept[IllegalArgumentException](AIPromptPython.paramSetter(source, name)) + .getMessage.contains(s"$name setter template")) + } + assert(AIPromptPython.paramSetter(setter, name).contains("self._ai_json")) + } + Seq(new OpenAIPrompt(), new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + val methods = stage.pyAdditionalMethods + val explicit = if (stage.isInstanceOf[AIPrompt]) "self.isSet(self.postProcessing)" + else "self._post_processing_explicitly_set" + assert(methods.contains(s"def _post_processing_is_explicit(self):\n return $explicit")) + Seq("clear(param)", "copy(extra)").foreach { call => + assert(methods.contains(s"super(${stage.getClass.getSimpleName}, self).$call")) + } + } + } + + test("struct JSON postprocessing diagnoses malformed and incompatible fields without rejecting null objects") { + val cases = Seq( + ("not json", true), ("""{"value":"wrong"}""", true), + ("""{"value":1},"extra":2""", true), ("""{"value":1} {"value":2}""", true), + ("""{"value":1,"nested":{"count":"wrong"}}""", true), + ("""{"value":null,"nested":null}""", false), ("{}", false), + ("```json\n{\"value\":2}\n```", false), + ("""{"value":3,"__synapseml_corrupt_record":"user field"}""", false)) + Seq("chat_completions", "responses").foreach { api => + cases.foreach { case (text, invalid) => + val quoted = JsString(text).compactPrint + val reply = if (api == "chat_completions") { + s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$quoted},"finish_reason":"stop"}]}""" + } else { + s"""{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":$quoted}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> + "value INT, nested STRUCT, __synapseml_corrupt_record STRING")) + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(Option(row.getAs[Row]("error")).exists(_.getString(0).contains("parse_error")) == invalid, + s"$api postprocessing incorrectly classified $text") + assert(row.getAs[String]("raw") == reply) + if (text.contains("user field")) { + assert(row.getAs[Row]("answer").getAs[String]("__synapseml_corrupt_record") == "user field") + } + assert(server.requests.size() == 1) + } + } + } + } + + test("sub-nanosecond concurrent timeouts fail preflight instead of becoming immediate") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + Seq(Double.MinPositiveValue, 0.5e-9, 0.999e-9).foreach { timeout => + stage.setConcurrentTimeout(timeout) + val error = intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + assert(error.getMessage.contains("concurrentTimeout")) + intercept[IllegalArgumentException](stage.transform(input)) + } + stage.setConcurrentTimeout(1e-9).transformSchema(input.schema) + assert(server.requests.isEmpty) + } + } + + test("sub-millisecond HTTP timeouts fail preflight instead of becoming infinite") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + Seq(Double.MinPositiveValue, 0.0005, 0.000999).foreach { timeout => + val error = intercept[IllegalArgumentException](stage.setTimeout(timeout).transformSchema(input.schema)) + assert(error.getMessage.contains("timeout")) + intercept[IllegalArgumentException](stage.transform(input)) + } + stage.setTimeout(0.001).transformSchema(input.schema) + assert(server.requests.isEmpty) + } + } + + test("native scalar null questions fail preflight while null question cells remain row errors") { + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val data = input.withColumn("questions", org.apache.spark.sql.functions.lit( + Option.empty[String].orNull).cast(StringType)) + Seq(new AISystemOne(), new AIPrompt().setApiType("systemone"), + new OpenAIDecisions(), new AIPrompt().setApiType("decisions")).foreach { stage => + stage.setUrl(if (stage.getApiType == "systemone") server.url else server.url + "/v1") + .setModel("model").setInputCol("text").setQuestions("null").setErrorCol("error") + intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + intercept[IllegalArgumentException](stage.transform(data)) + val row = stage.setQuestionsCol("questions").transform(data).head() + assert(row.getAs[Row]("error").getString(0).contains("missing_input")) + } + assert(server.requests.isEmpty) + } + } + + test("struct JSON postprocessing preserves an existing incomplete-response diagnostic") { + val reply = chatReply.replace("\"answer\"", "\"not json\"").replace("\"stop\"", "\"length\"") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> "value INT")) + val row = stage.transform(input).head() + assert(row.getAs[Row]("error").getString(0).contains("incomplete")) + assert(row.getAs[String]("raw") == reply) + } + } + + test("container JSON postprocessing diagnoses nested conversion failures but preserves valid nulls") { + val cases = Seq( + ("ARRAY>", """[{"value":"wrong"}]""", true), + ("ARRAY>", """{"value":"wrong"}""", true), + ("ARRAY>", """{"value":1}""", false), + ("ARRAY>", """[{"value":null},null,{}]""", false), + ("ARRAY>", "```json\n[{\"value\":2}]\n```", false), + ("MAP>", """{"entry":{"value":"wrong"}}""", true), + ("MAP>", """{"entry":{"value":null},"empty":null}""", false), + ("MAP>", """{"entry":[1,"wrong"]}""", true), + ("MAP>", """{"entry":[1,null]}""", false)) + Seq("chat_completions", "responses").foreach { api => + cases.foreach { case (schema, text, invalid) => + val quoted = JsString(text).compactPrint + val reply = if (api == "chat_completions") { + s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$quoted},"finish_reason":"stop"}]}""" + } else { + s"""{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":$quoted}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> schema)) + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(Option(row.getAs[Row]("error")).exists(_.getString(0).contains("parse_error")) == invalid, + s"$api postprocessing incorrectly classified $text as $schema") + if (!invalid) assert(!row.isNullAt(row.fieldIndex("answer"))) + if (text == """{"value":1}""") assert(row.getAs[Seq[Row]]("answer") == Seq(Row(1))) + assert(row.getAs[String]("raw") == reply && server.requests.size() == 1) + } + } + } + } + + test("complete deployment URLs reject literal and encoded dot segments for scalar and column models") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq(".", "..", "%2e", "%2E%2e", "..%2fother", "foo%2F..%2Fadmin", + "foo%5c..%5cadmin", "%252e%252e%252fadmin").foreach { segment => + Seq("chat_completions", "embeddings").foreach { api => + Seq(false, true).foreach { custom => + val route = if (api == "embeddings") "embeddings" else "chat/completions" + val address = server.url + s"/openai/deployments/$segment/$route" + val stage = new AIPrompt().setApiType(api).setInputCol("text") + .setDeploymentName(java.net.URLDecoder.decode(segment, "UTF-8")) + if (custom) stage.setCustomUrlRoot(address) else stage.setUrl(address) + Seq(stage, stage.copy(ParamMap.empty).setDeploymentNameCol("text")).foreach { candidate => + val error = intercept[IllegalArgumentException](candidate.transformSchema(input.schema)) + assert(error.getMessage.contains("dot segment")) + } + } + } + } + Seq(".", "..").foreach { model => + new AIPrompt().setUrl(server.url + "/v1").setModel(model).setInputCol("text").transform(input).collect() + assert(AIJson.objectString(server.requests.asScala.last._2).fields("model") == JsString(model)) + } + assert(server.requests.size() == 2) + } + } + + test("complete URL overrides enforce deployment identity and models-route selection before HTTP") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val address = server.url + "/openai/deployments/model/chat/completions/?version=custom" + def stage: AIPrompt = new AIPrompt().setCustomUrlRoot(address).setDeploymentName("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + intercept[IllegalArgumentException](stage.setDeploymentName("other").transformSchema(input.schema)) + val row = stage.setDeploymentNameCol("text").transform(input).head() + assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + val models = server.url + "/models/chat/completions?version=custom" + intercept[IllegalArgumentException](stage.setCustomUrlRoot(models).transformSchema(input.schema)) + assert(server.requests.isEmpty) + stage.transform(input).collect() + assert(server.requests.peek()._1 == "/openai/deployments/model/chat/completions/?version=custom") + val valid = stage.setCustomUrlRoot(models) + valid.clear(valid.deploymentName) + valid.setModel("model").transform(input).collect() + assert(server.requests.asScala.last._1 == "/models/chat/completions?version=custom") + } + } + + test("complete URL overrides preserve trailing path slashes and query bytes") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq("/proxy/", "/proxy?value=a/", "/proxy/?value=a%2Fb%2F&case=A+B").foreach { suffix => + new AIPrompt().setCustomUrlRoot(server.url + suffix).setModel("model").setInputCol("text") + .transform(input).collect() + assert(server.requests.asScala.last._1 == suffix) + } + } + } + + test("invalid concurrent timeout defaults fail before HTTP") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq(0.0, -1.0, 1e-10, Double.NaN, Double.PositiveInfinity).foreach { timeout => + val stage = new TestableAIDefaultTimeoutPrompt("invalid-default-timeout").setConcurrentTimeoutDefault(timeout) + .setUrl(server.url + "/v1").setModel("model").setInputCol("text") + assert(intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + .getMessage.contains("concurrentTimeout")) + assert(intercept[IllegalArgumentException](stage.transform(input)).getMessage.contains("concurrentTimeout")) + assert(server.requests.isEmpty) + } + } + } + + test("Spark job cancellation stops bounded retry workers at both concurrency settings") { + Seq((1, true), (2, true), (2, false)).foreach { case (concurrency, interrupt) => + val received = new CountDownLatch(concurrency) + Using.resource(new AIRecordingServer((number, _) => { + received.countDown() + if (number <= concurrency) 429 -> "{}" else 200 -> chatReply + }, retryAfter = "3")) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setConcurrency(concurrency).setTimeout(5) + val data = input.union(input).coalesce(1) + val context = spark.sparkContext + val group = java.util.UUID.randomUUID().toString + val executor = Executors.newSingleThreadExecutor() + val job = executor.submit(new Callable[Unit] { + override def call(): Unit = { + context.setJobGroup(group, "AIPrompt cancellation regression", interruptOnCancel = interrupt) + try stage.transform(data).collect() finally context.clearJobGroup() + } + }) + try { + assert(received.await(30, TimeUnit.SECONDS), "The requests never reached retry backoff") + context.cancelJobGroup(group) + intercept[ExecutionException](job.get(15, TimeUnit.SECONDS)) + Thread.sleep(4000) + assert(server.requests.size() == concurrency, "Canceled workers issued another paid request") + } finally { + context.cancelJobGroup(group) + job.cancel(true) + executor.shutdownNow() + assert(executor.awaitTermination(15, TimeUnit.SECONDS)) + } + } + } + } + + test("timeout and early consumer completion abort in-flight HTTP without leaving retry workers") { + Seq(false, true).foreach { earlyCompletion => + val blocked = new CountDownLatch(1) + val release = new CountDownLatch(1) + Using.resource(new AIRecordingServer((_, body) => { + if (body.contains("blocked") || !earlyCompletion) { + blocked.countDown() + release.await(15, TimeUnit.SECONDS) + } else { + assert(blocked.await(15, TimeUnit.SECONDS)) + } + 200 -> chatReply + })) { server => + try { + val data = spark.createDataFrame(Seq(Row("ready"), Row("blocked")).asJava, + new StructType().add("text", StringType)).coalesce(1) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setConcurrency(2).setConcurrentTimeout(2).setTimeout(60).setOutputCol("answer") + val result = stage.transform(data) + if (earlyCompletion) assert(result.limit(1).head().getAs[String]("answer") == "answer") + else intercept[org.apache.spark.SparkException](result.collect()) + assert(blocked.getCount == 0, "The regression did not reach an in-flight HTTP request") + assert(!Thread.getAllStackTraces.keySet().asScala.exists(thread => + thread.isAlive && thread.getName.startsWith("synapseml-bounded-http-"))) + } finally release.countDown() + } + } + } + + test("implicit Fabric requests preserve workload classification with explicit credentials and header precedence") { + val previous = GlobalParams.getGlobalParam(URLKey) + val root = "https://fabric.example.invalid/cognitive/openai" + val reply = chatReply + try { + GlobalParams.setGlobalParam(URLKey, root) + Seq(false, true).foreach { explicitKey => + val stage = new TestableAIFabricAuthPrompt("fabric-classification").setDeploymentName("model") + .setInputCol("text").setOutputCol("answer") + .setCustomHeaders(Map("x-taxonomy-traffictype" -> "Interactive", "X-Llm-Service-Tier" -> "standard", + "X-Taxonomy-ExtendedProperties" -> "caller", "x-caller" -> "preserved")) + if (explicitKey) stage.setSubscriptionKey("synthetic-explicit-key") + stage.setHandler { (_, request) => + val headers = request.headers.map(h => h.name.toLowerCase(java.util.Locale.ROOT) -> h.value).toMap + assert(headers("x-taxonomy-traffictype") == "Background") + assert(headers("x-llm-service-tier") == "flex") + val properties = headers("x-taxonomy-extendedproperties").parseJson.asJsObject.fields + assert(properties("feature") == JsString("synapseml") && properties.contains("runtime")) + assert(headers("x-caller") == "preserved") + HTTPSchema.stringToResponse(reply, 200, "OK") + } + assert(stage.transform(input).head().getAs[String]("answer") == "answer") + } + Seq("explicit", "other").foreach { destination => + val stage = new TestableAIFabricAuthPrompt("fabric-not-classified").setDeploymentName("model") + .setInputCol("text").setOutputCol("answer") + if (destination == "explicit") stage.setUrl(root) + else GlobalParams.setGlobalParam(URLKey, "https://other.example.invalid/v1") + stage.setHandler { (_, request) => + assert(!request.headers.exists(_.name.equalsIgnoreCase("X-Taxonomy-TrafficType"))) + assert(!request.headers.exists(_.name.equalsIgnoreCase("X-Llm-Service-Tier"))) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + assert(stage.transform(input).head().getAs[String]("answer") == "answer") + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + + test("implicit Fabric scalar model preflight rejects unavailable models only during transform") { + val previous = GlobalParams.getGlobalParam(URLKey) + try { + GlobalParams.setGlobalParam(URLKey, "https://fabric.example.invalid/cognitive/openai") + val stage = new TestableAIFabricAuthPrompt("fabric-preflight").setModel("blocked-model").setInputCol("text") + stage.transformSchema(input.schema) + val error = intercept[IllegalArgumentException](stage.transform(input)) + assert(error.getMessage.contains("Synthetic Fabric model is unavailable")) + val reply = chatReply + val column = new TestableAIFabricAuthPrompt("fabric-column-model").setModelCol("text") + .setInputCol("text").setOutputCol("answer") + column.setHandler((_, _) => HTTPSchema.stringToResponse(reply, 200, "OK")) + assert(column.transform(input).head().getAs[String]("answer") == "answer") + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseFormatSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseFormatSuite.scala new file mode 100644 index 00000000000..7f343fc01a5 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseFormatSuite.scala @@ -0,0 +1,339 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.{ArrayType, BinaryType, IntegerType, MapType, StringType, StructType, TimestampType} +import spray.json._ + +import java.nio.file.Files +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.util.Using + +class AIPromptResponseFormatSuite extends AIPromptTestBase { + private val replies = Seq( + "chat_completions" -> """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""", + "responses" -> """{"output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin) + + private def format(body: String, api: String): JsObject = { + val fields = AIJson.objectString(body).fields + if (api == "responses") fields("text").asJsObject.fields("format").asJsObject + else fields("response_format").asJsObject + } + + private def responseWithEcho(body: String): String = JsObject( + replies.toMap.apply("responses").parseJson.asJsObject.fields ++ Map( + "text" -> body.parseJson.asJsObject.fields("text"), + "id" -> JsString("echoed-schema"), + "usage" -> JsObject("total_tokens" -> JsNumber(7)))).compactPrint + + private def verifyEcho(row: Row, reply: String): Unit = { + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer") + assert(row.getAs[String]("response_id") == "echoed-schema") + assert(row.getAs[Row]("usage").getLong(2) == 7L && row.getAs[String]("raw") == reply) + } + + test("response format map schemas reject non-JSON SQL types before HTTP") { + val struct = new StructType().add("value", StringType) + val unsupported = Seq(MapType(IntegerType, StringType), MapType(StringType, struct), + MapType(StringType, ArrayType(struct)), MapType(StringType, MapType(IntegerType, StringType)), + MapType(StringType, TimestampType), MapType(StringType, BinaryType), + MapType(StringType, ArrayType(BinaryType)), MapType(StringType, MapType(StringType, BinaryType))) + Using.resource(new AIRecordingServer((_, _) => 200 -> replies.head._2)) { server => + unsupported.foreach { formatType => + val input = spark.createDataFrame(Collections.emptyList[Row](), + new StructType().add("text", StringType).add("format", formatType)) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setResponseFormatCol("format") + Seq[() => Any](() => stage.transformSchema(input.schema), () => stage.transform(input)).foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("responseFormatCol has an unsupported data type")) + } + assert(server.requests.isEmpty) + } + } + } + + test("schema convenience rejects invalid values without changing scalar or column bindings") { + val stage = new AIPrompt().setResponseFormat(Map("type" -> "json_object")) + val configured = stage.get(stage.responseFormat) + Seq[Any](new Object(), Double.NaN, Double.PositiveInfinity).foreach { value => + intercept[IllegalArgumentException](stage.setResponseSchema(Map("maximum" -> value))) + assert(stage.get(stage.responseFormat) == configured) + } + stage.setResponseFormatCol("format") + intercept[IllegalArgumentException](stage.setResponseSchema(Map.empty[String, Any])) + intercept[IllegalArgumentException](stage.setResponseSchema(Map("type" -> "object"), "invalid name")) + intercept[IllegalArgumentException](stage.setResponseSchema(Map("maximum" -> new Object()))) + assert(stage.getResponseFormatCol == "format") + } + + private def deepSchema(height: Int): Map[String, Any] = + (1 until height).foldLeft(Map[String, Any]("type" -> "string")) { (value, _) => + Map("type" -> "array", "items" -> value) + } + + private def bindSchema(stage: AIPrompt, schema: Map[String, Any], binding: String): AIPrompt = { + val flat = Map("type" -> "json_schema", "name" -> "response_schema", "schema" -> schema) + binding match { + case "convenience" => stage.setResponseSchema(schema) + case "raw" => stage.setResponseFormat(schema) + case "flat" => stage.setResponseFormat(flat) + case "nested" => stage.setResponseFormat(Map("type" -> "json_schema", "json_schema" -> (flat - "type"))) + } + } + + test("normalized response format depth is rejected before changing scalar or column bindings") { + for { + api <- replies.map(_._1) + binding <- Seq("convenience", "raw", "flat", "nested") + height <- Seq(AIJson.MaxDepth - 1, AIJson.MaxDepth) + perRow <- Seq(false, true) + } { + val stage = new AIPrompt().setApiType(api).setResponseFormat(Map("type" -> "json_object")) + if (perRow) stage.setResponseFormatCol("format") + val original = stage.get(stage.responseFormat) + withClue(s"$api $binding height=$height perRow=$perRow: ") { + val error = intercept[IllegalArgumentException](bindSchema(stage, deepSchema(height), binding)) + assert(error.getMessage.contains("nesting")) + assert(stage.get(stage.responseFormat) == original) + } + } + } + + test("responses column schema depth boundary retains echoed metadata") { + import spark.implicits._ + def flat(height: Int): Map[String, Any] = Map("type" -> "json_schema", + "name" -> "response_schema", "schema" -> deepSchema(height)) + val formats = Seq( + deepSchema(AIJson.MaxDepth - 1), + flat(AIJson.MaxDepth - 1), + Map("type" -> "json_schema", "json_schema" -> (flat(AIJson.MaxDepth - 2) - "type")), + deepSchema(AIJson.MaxDepth)) + val input = formats.zipWithIndex.map { case (value, index) => + (index, "hello", AIJson.fromAny(value).compactPrint) + }.toDF("index", "text", "format") + Using.resource(new AIRecordingServer((_, body) => 200 -> responseWithEcho(body))) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setResponseFormatCol("format").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("response_id").setRawResponseCol("raw") + def verify(rows: Array[Row]): Unit = { + val ordered = rows.sortBy(_.getAs[Int]("index")) + assert(ordered.length == 4) + ordered.take(3).foreach { row => + val expected = AIRequestAdapter.responseFormat(row.getAs[String]("format"), "responses") + val reply = responseWithEcho(JsObject("text" -> JsObject("format" -> expected)).compactPrint) + verifyEcho(row, reply) + } + assert(ordered.last.getAs[String]("answer") == null) + assert(ordered.last.getAs[Row]("error").getString(0).contains("invalid_request")) + } + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + verify(output.collect()) + assert(server.requests.size() == 3) + } + } + + replies.foreach { case (api, reply) => + test(s"$api response format depth boundary survives persistence and HTTP") { + val schema = deepSchema(AIJson.MaxDepth - 2) + Using.resource(new AIRecordingServer((_, body) => + 200 -> (if (api == "responses") responseWithEcho(body) else reply))) { server => + Seq("convenience", "raw", "flat", "nested").foreach { binding => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("response_id").setRawResponseCol("raw") + bindSchema(stage, schema, binding) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-depth").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val loaded = AIPrompt.load(saved) + val input = spark.range(1).withColumn("text", lit("hello")) + val result = loaded.transform(input) + assert(result.schema == loaded.transformSchema(input.schema)) + val row = result.head() + val sent = server.requests.remove()._2 + if (api == "responses") verifyEcho(row, responseWithEcho(sent)) + else assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer") + val body = sent.parseJson.asJsObject + val details = if (api == "responses") body.fields("text").asJsObject.fields("format").asJsObject + else body.fields("response_format").asJsObject.fields("json_schema").asJsObject + assert(details.fields("schema") == AIJson.fromAny(schema)) + } + assert(server.requests.isEmpty) + } + } + + test(s"$api schema convenience preserves exact numbers and JSON containers through persistence and HTTP") { + val decimal = BigDecimal("0.12345678901234567890123456789") + val integer = new java.math.BigInteger("1234567890123456789012345678901234567890") + val javaProperty = new java.util.LinkedHashMap[String, Any]() + javaProperty.put("type", "number") + javaProperty.put("minimum", decimal.bigDecimal) + val properties = scala.collection.immutable.ListMap[String, Any]( + "scala_decimal" -> Map("type" -> "number", "minimum" -> decimal), + "large_integer" -> Map("type" -> "integer", "maximum" -> integer), + "java_decimal" -> javaProperty, + "json_value" -> JsObject("type" -> JsString("string"), + "enum" -> JsArray(JsString("yes"), JsString("no")))) + val schema = scala.collection.immutable.ListMap[String, Any]( + "type" -> "object", "properties" -> properties, + "required" -> properties.keys.toArray, "additionalProperties" -> false) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setResponseSchema(schema, "precise", strict = false) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-schema").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val row = AIPrompt.load(saved).transform(spark.range(1).withColumn("text", lit("hello"))).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(server.requests.size() == 1) + val sent = format(server.requests.peek()._2, api) + val details = if (api == "responses") sent else sent.fields("json_schema").asJsObject + assert(details.fields("schema") == AIJson.fromAny(schema)) + assert(details.fields("schema").asJsObject.fields("properties").asJsObject.fields.keys.toSeq == + properties.keys.toSeq) + assert(details.fields("name") == JsString("precise") && details.fields("strict") == JsBoolean(false)) + } + } + + Seq(false, true).foreach { perRow => + test(s"$api preserves bare response format order through persistence and HTTP with perRow=$perRow") { + val fields = Map[String, Any]( + "name" -> "ordered", "description" -> "Extract the answer", "strict" -> true, + "schema" -> Map("type" -> "object", "properties" -> Map("answer" -> Map("type" -> "string")))) + val names = Vector("description", "strict", "schema", "name") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + Seq(names, names.reverse).foreach { order => + val bare = ListMap(order.map(name => name -> fields(name)): _*) + val input = spark.range(1).withColumn("text", lit("hello")) + .withColumn("format", lit(AIJson.fromAny(bare).compactPrint)) + if (perRow) stage.setResponseFormatCol("format") else stage.setResponseFormat(bare) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-bare-order").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val loaded = AIPrompt.load(saved) + val outputSchema = stage.transformSchema(input.schema) + assert(loaded.transformSchema(input.schema) == outputSchema) + val output = loaded.transform(input) + assert(output.schema == outputSchema) + val row = output.head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + val sent = format(server.requests.asScala.last._2, api) + val details = if (api == "responses") sent else sent.fields("json_schema").asJsObject + assert(details.fields.keys.toVector.filterNot(_ == "type") == order) + assert(details.fields - "type" == AIJson.fromAny(bare).asJsObject.fields) + } + assert(server.requests.size() == 2) + } + } + } + + test(s"$api preserves string and nested-array response format map bindings") { + val cases = Seq( + MapType(StringType, StringType) -> Map("type" -> "json_object"), + MapType(StringType, MapType(StringType, ArrayType(StringType))) -> + Map("items" -> Map("enum" -> Array("yes", "no")))) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + cases.zipWithIndex.foreach { case ((formatType, value), index) => + val input = spark.createDataFrame(Collections.singletonList(Row("hello", value)), + new StructType().add("text", StringType).add("format", formatType)) + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setResponseFormatCol("format").setOutputCol("answer").setErrorCol("error") + val row = stage.transform(input).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + val sent = format(server.requests.remove()._2, api) + if (index == 0) assert(sent.fields("type") == JsString("json_object")) + else { + val details = if (api == "responses") sent else sent.fields("json_schema").asJsObject + assert(details.fields("schema") == + JsObject("items" -> JsObject("enum" -> JsArray(JsString("yes"), JsString("no"))))) + } + } + } + } + + test(s"$api preserves response format descriptions and extensions") { + val details = Map[String, Any]("name" -> "answer", "description" -> "Extract the answer", + "schema" -> Map("type" -> "object", "description" -> "Inner schema", + "properties" -> Map("answer" -> Map("type" -> "string"))), + "strict" -> true, "provider_extension" -> Map("weight" -> BigDecimal("0.1234567890123456789"))) + val nested = Map[String, Any]("type" -> "json_schema", "json_schema" -> details) + val flat = details + ("type" -> "json_schema") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(false, true).foreach { perRow => + server.requests.clear() + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + val input = spark.range(1).withColumn("text", lit("hello")) + .withColumn("format", lit(AIJson.fromAny(flat).compactPrint)) + if (perRow) stage.setResponseFormatCol("format") else stage.setResponseFormat(nested) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-format").resolve("stage").toString + stage.copy(ParamMap.empty).write.session(spark).save(saved) + val loaded = AIPrompt.load(saved) + val result = loaded.transform(input) + assert(result.schema == loaded.transformSchema(input.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer") + assert(server.requests.size() == 1) + assert(format(server.requests.peek()._2, api) == AIJson.fromAny( + if (api == "responses") flat else nested)) + } + } + } + test(s"$api preserves text and JSON format extensions without relocating raw schema descriptions") { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq( + Map[String, Any]("type" -> "text", "provider_extension" -> "value"), + Map[String, Any]("type" -> "json_object", "provider_extension" -> "value"), + Map[String, Any]("type" -> "object", "description" -> "Inner schema", "properties" -> Map.empty), + Map[String, Any]("description" -> "Inner schema", "properties" -> Map.empty) + ).foreach { configured => + Seq(false, true).foreach { perRow => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + val input = spark.range(1).withColumn("text", lit("hello")) + .withColumn("format", lit(AIJson.fromAny(configured).compactPrint)) + if (perRow) stage.setResponseFormatCol("format") else stage.setResponseFormat(configured) + assert(stage.transform(input).head().getAs[Row]("error") == null) + val sent = format(server.requests.asScala.last._2, api) + if (configured.get("type").forall(_ == "object")) { + val details = if (api == "responses") sent else sent.fields("json_schema").asJsObject + assert(details.fields("schema") == AIJson.fromAny(configured)) + assert(!details.fields.contains("description")) + } else assert(sent == AIJson.fromAny(configured)) + } + } + assert(server.requests.size() == 8) + } + } + + test(s"$api rejects ambiguous response format wrappers instead of dropping their fields") { + import spark.implicits._ + val details = Map[String, Any]("name" -> "answer", "schema" -> Map("type" -> "object")) + val nested = Map[String, Any]("type" -> "json_schema", "json_schema" -> details) + val invalid = Seq(nested + ("description" -> "Outside the schema wrapper"), + nested + ("json_schema" -> (details + ("type" -> "text"))), + nested ++ details) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + invalid.foreach(value => intercept[IllegalArgumentException](stage.setResponseFormat(value))) + val input = (invalid :+ nested).zipWithIndex.map { case (value, index) => + (index, "hello", AIJson.fromAny(value).compactPrint) + }.toDF("id", "text", "format") + val rows = stage.setResponseFormatCol("format").transform(input).collect().sortBy(_.getAs[Int]("id")) + assert(rows.take(invalid.size).forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(rows.last.getAs[Row]("error") == null && rows.last.getAs[String]("answer") == "answer") + assert(server.requests.size() == 1) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseShapeSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseShapeSuite.scala new file mode 100644 index 00000000000..cee5a7285bf --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptResponseShapeSuite.scala @@ -0,0 +1,163 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.ml.linalg.Vectors +import org.apache.spark.sql.Row +import org.apache.spark.sql.types._ +import spray.json._ + +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptResponseShapeSuite extends AIPromptTestBase { + Seq("chat_completions", "responses").foreach { api => + test(s"$api JSON NOT NULL fields have nullable output schemas") { + val text = """{"value":7,"nested":{},"items":[{},null],"lookup":{"key":{}}}""" + val encoded = JsString(text).compactPrint + val reply = if (api == "chat_completions") { + s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$encoded}}]}""" + } else { + s"""{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":$encoded}]}]}""".stripMargin + } + val nested = new StructType().add("n", IntegerType) + val expected = new StructType().add("value", IntegerType).add("nested", nested) + .add("items", ArrayType(nested)).add("lookup", MapType(StringType, nested)) + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", StringType)) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setPostProcessingOptions(Map("jsonSchema" -> + """value INT NOT NULL, nested STRUCT NOT NULL, + |items ARRAY>, lookup MAP>""".stripMargin)) + assert(stage.transformSchema(data.schema)("answer").dataType == expected) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + val answer = row.getAs[Row]("answer") + assert(row.getAs[Row]("error") == null && answer.getInt(0) == 7) + assert(answer.getStruct(1).isNullAt(0)) + assert(answer.getSeq[Row](2).head.isNullAt(0) && answer.getSeq[Row](2).last == null) + assert(answer.getMap[String, Row](3)("key").isNullAt(0)) + assert(server.requests.size() == 1) + } + } + + test(s"$api failed projections retain partial observations and provider diagnostics") { + import spark.implicits._ + val key = if (api == "chat_completions") "choices" else "output" + val valid = if (api == "chat_completions") { + """{"index":0,"message":{"role":"assistant","content":"partial"}}""" + } else { + """{"type":"message","role":"assistant","content":[{"type":"output_text","text":"partial"}, + |{"type":"output_text","text":42},{"type":"output_text","text":" tail"}]}""".stripMargin + } + val projections = Map("NULL_CASE" -> "null", "WRONG_CASE" -> "42", + "MIXED_CASE" -> s"[$valid,null]", "FIRST_BAD_CASE" -> s"[null,$valid]") + val input = projections.keys.toSeq.toDF("text") + Using.resource(new AIRecordingServer((_, body) => { + val name = projections.keys.find(body.contains).get + 200 -> s"""{"id":"failed-id","error":{"code":"failed","message":"private provider detail"}, + |"$key":${projections(name)},"usage":{"total_tokens":7}}""".stripMargin + })) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setUsageCol("usage") + .setResponseIdCol("response_id").setRawResponseCol("raw") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect() + assert(rows.length == projections.size) + rows.foreach { row => + val name = row.getAs[String]("text") + val expected = if (name == "MIXED_CASE" || (api == "responses" && name == "FIRST_BAD_CASE")) { + if (api == "chat_completions") "partial" else "partial tail" + } else None.orNull + assert(row.getAs[String]("answer") == expected) + assert(row.getAs[String]("response_id") == "failed-id" && row.getAs[Row]("usage").getLong(2) == 7L) + val error = row.getAs[Row]("error") + assert(error != null && error.getString(0).contains("provider_error")) + assert(!error.getString(0).contains("private provider detail")) + assert(row.getAs[String]("raw").contains("private provider detail")) + } + assert(server.requests.size() == 4) + } + } + } + + Seq("responses", "decisions").foreach { api => + test(s"$api message type tags validate before HTTP") { + val reply = if (api == "decisions") NativeRequestFixtures.DecisionResponse else + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val message = new StructType().add("role", StringType).add("content", StringType).add("type", StringType) + val schema = new StructType().add("tag", StringType).add("messages", ArrayType(message)) + val data = spark.createDataFrame(Seq("message", "unsupported", "", Option.empty[String].orNull) + .map(tag => Row(tag, Seq(Row("user", "hello", tag)))).asJava, schema) + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + result.collect().foreach { row => + if (row.isNullAt(0) || row.getString(0) == "message") { + assert(row.getAs[Row]("error") == null && !row.isNullAt(row.fieldIndex("answer"))) + } else assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + } + val invalidType = spark.createDataFrame(Collections.singletonList(Row(Seq(Row("user", "hello", 1)))), + new StructType().add("messages", ArrayType( + new StructType().add("role", StringType).add("content", StringType).add("type", IntegerType)))) + assert(stage.transform(invalidType).head().getAs[Row]("error").getString(0).contains("invalid_request")) + assert(server.requests.size() == 2) + } + } + } + + test("nullable embedding object tags preserve vectors and diagnostics") { + import spark.implicits._ + val input = Seq("hello").toDF("text") + val reply = """{"id":"embedding-id","object":null,"data":[{"object":null,"index":0,"embedding":[0.25,-0.5]}], + |"usage":{"prompt_tokens":2,"total_tokens":2}}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType("embeddings").setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setUsageCol("usage") + .setResponseIdCol("response_id").setRawResponseCol("raw") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val row = output.head() + assert(row.getAs[org.apache.spark.ml.linalg.Vector]("answer") == Vectors.dense(0.25, -0.5)) + assert(row.getAs[Row]("error") == null && row.getAs[Row]("usage").getLong(2) == 2L) + assert(row.getAs[String]("response_id") == "embedding-id" && row.getAs[String]("raw") == reply) + assert(server.requests.size() == 1) + } + } + + test("Responses in-progress messages retain partial observations and incomplete errors") { + import spark.implicits._ + val input = Seq("absent", "completed").toDF("text") + Using.resource(new AIRecordingServer((_, body) => { + val status = if (body.contains("completed")) """"status":"completed",""" else "" + 200 -> s"""{$status"id":"partial-id","output":[{"type":"message","role":"assistant", + |"status":"in_progress","content":[{"type":"output_text","text":"partial"}]}], + |"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}""".stripMargin + })) { server => + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setUsageCol("usage") + .setResponseIdCol("response_id").setRawResponseCol("raw") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect() + assert(rows.length == 2) + rows.foreach { row => + assert(row.getAs[String]("answer") == "partial" && row.getAs[String]("response_id") == "partial-id") + assert(row.getAs[Row]("usage").getLong(2) == 3L && row.getAs[String]("raw").contains("in_progress")) + assert(Option(row.getAs[Row]("error")).exists(_.getString(0).contains("incomplete"))) + } + assert(server.requests.size() == 2) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptSuite.scala new file mode 100644 index 00000000000..003d03f624d --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptSuite.scala @@ -0,0 +1,792 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.fuzzing.TestObject +import com.microsoft.azure.synapse.ml.core.utils.ModelEquality +import com.microsoft.azure.synapse.ml.io.http.{HTTPSchema, URLKey} +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.microsoft.azure.synapse.ml.services.openai.{OpenAIDecisions, OpenAIDeploymentNameKey, + OpenAIEmbeddingDeploymentNameKey, OpenAIPrompt} +import org.apache.spark.ml.linalg.Vector +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.ml.util.MLReadable +import org.apache.spark.sql.{DataFrame, Encoders, Row} +import org.apache.spark.sql.types._ +import spray.json._ + +import java.nio.file.Files +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptSuite extends AIPromptFuzzing[AIPrompt] { + override def reader: MLReadable[_] = AIPrompt + + override def testObjects(): Seq[TestObject[AIPrompt]] = Seq(new TestObject( + new AIPrompt("aiprompt-fuzzing").setUrl("https://example.invalid/v1").setDeploymentName("synthetic-model") + .setSubscriptionKey("synthetic-key").setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setExtraBody(Map("metadata" -> Map("source" -> "fuzzing"))), + spark.createDataFrame(Seq(Row(Option.empty[String].orNull)).asJava, new StructType().add("text", StringType)))) + + private val chatReply = + """{"id":"chat-1","choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}], + |"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3},"unknown":{"retained":true}}""".stripMargin + private val responseReply = + """{"id":"response-1","status":"completed","output":[ + |{"type":"message","role":"assistant","content":[{"type":"output_text","text":"one"}, + |{"type":"output_text","text":"two"}]}, + |{"type":"function_call","arguments":"not prose"}, + |{"type":"message","role":"assistant","content":[{"type":"output_text","text":"three"}]}], + |"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3},"unknown":true}""".stripMargin + private val embeddingReply = + """{"data":[{"index":0,"embedding":[0.25,0.5]}],"model":"returned-model", + |"usage":{"prompt_tokens":2,"total_tokens":2}}""".stripMargin + + private def input: DataFrame = spark.createDataFrame( + Seq(Row("hello"), Row("world")).asJava, new StructType().add("text", StringType)) + + private def prompt(server: AIRecordingServer, api: String = "chat_completions"): AIPrompt = + new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setDeploymentName("explicit-model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setRawResponseCol("raw").setUsageCol("usage") + + test("R constructors preserve nested JSON values as single escaped arguments and retain column bindings") { + val questions = NativeRequestFixtures.DecisionQuestions + val body = """{"items":[{"quoted":"a \"quote\" and \\ path\nline","value":true},{"value":null}]}""" + val cases = Seq( + (new OpenAIDecisions().setQuestions(questions), "questions", questions), + (new AISystemOne().setQuestions(NativeRequestFixtures.SystemQuestions), + "questions", NativeRequestFixtures.SystemQuestions), + (new AIPrompt().setExtraBody(body), "extraBody", body)) + def argument(stage: AIPrompt, name: String): String = { + val prefix = s",$name=" + rTestInstantiateModel(stage, 0).split("\n").map(_.trim) + .find(_.startsWith(prefix)).getOrElse(fail(s"Missing R argument $name")).stripPrefix(prefix) + } + cases.foreach { case (stage, name, expected) => + val json = AIJson.parse(argument(stage, name)) match { + case JsString(value) => value + case _ => fail(s"$name must be one JSON string argument") + } + assert(AIJson.parse(json) == AIJson.parse(expected)) + name match { + case "questions" => stage.setQuestions(json).setQuestionsCol("json_column") + case "extraBody" => stage.setExtraBody(json).setExtraBodyCol("json_column") + } + assert(argument(stage, name + "Col") == "\"json_column\"") + } + } + + test("R fuzzing fixtures preserve implicit UID-derived defaults for all three public stages") { + val constructors = Seq[String => AIPrompt]( + uid => new AIPrompt(uid), uid => new AISystemOne(uid), uid => new OpenAIDecisions(uid)) + constructors.foreach { create => + def configured(uid: String): AIPrompt = + create(uid).setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + val original = configured("r-constructor-source") + val uidArgument = rTestInstantiateModel(original, 0).split("\n").map(_.trim) + .find(_.startsWith(",uid = ")).getOrElse(fail("Missing generated R UID")) + .stripPrefix(",uid = ").stripSuffix(")") + val uid = AIJson.parse(uidArgument) match { + case JsString(value) => value + case _ => fail("Expected a string UID") + } + val recreated = configured(uid) + assert(!recreated.isSet(recreated.messagesCol) && !recreated.isSet(recreated.responseIdCol)) + ModelEquality.assertEqual(original, recreated) + } + } + + test("all new stages are lazy and retain original columns, raw JSON and usage") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = prompt(server) + val result = stage.transform(input) + assert(server.requests.isEmpty) + assert(result.schema == stage.transformSchema(input.schema)) + val rows = result.collect() + assert(rows.forall(_.getAs[Row]("error") == null)) + assert(rows.map(_.getAs[String]("answer")).toSeq == Seq("answer", "answer")) + assert(server.requests.size() == 2) + assert(result.collect().forall(_.getAs[String]("raw").contains("retained"))) + assert(server.requests.size() == 4) + assert(server.requests.asScala.forall { case (path, body) => + path == "/v1/chat/completions" && AIJson.objectString(body).fields("model") == JsString("explicit-model") + }) + } + } + + test("Responses selector is explicit, concatenates assistant text, and preserves non-text items") { + Using.resource(new AIRecordingServer((_, _) => 200 -> responseReply)) { server => + val stage = prompt(server, "responses").setStore(true).setResponseIdCol("response_id") + .setExtraBody(Map("text" -> Map("verbosity" -> "low"), "provider_extension" -> Seq(1, 2))) + val df = stage.transform(input.limit(1)) + assert(df.schema == stage.transformSchema(input.limit(1).schema)) + val row = df.head() + assert(row.getAs[String]("answer") == "onetwothree") + assert(row.getAs[String]("response_id") == "response-1") + assert(row.getAs[String]("raw").contains("not prose")) + val (path, body) = server.requests.peek() + assert(path == "/v1/responses") + val payload = AIJson.objectString(body) + assert(payload.fields("text").asJsObject.fields("verbosity") == JsString("low")) + assert(payload.fields("provider_extension") == JsArray(JsNumber(1), JsNumber(2))) + } + } + + test("fused Dataset maps do not require input Rows to carry schemas") { + val original = input + val mapped = original.map(row => Row.fromSeq(row.toSeq))(Encoders.row(original.schema)) + val cases = Seq( + ("chat_completions", chatReply, None), + ("responses", responseReply, None), + ("embeddings", embeddingReply, None), + ("systemone", NativeRequestFixtures.SystemResponse, Some(NativeRequestFixtures.SystemQuestions)), + ("decisions", NativeRequestFixtures.DecisionResponse, Some(NativeRequestFixtures.DecisionQuestions))) + cases.foreach { case (api, reply, questions) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = prompt(server, api) + if (api == "systemone") stage.setUrl(server.url) + questions.foreach(stage.setQuestions) + if (Set("systemone", "decisions", "embeddings")(api)) { + val adapter = new AIRequestAdapter(stage, original.schema) + val (_, body) = adapter.request(Row("schema-less"), None, None) + val field = if (api == "systemone") "state" else "input" + assert(AIJson.objectString(body).fields(field) == JsString("schema-less")) + } + val result = stage.transform(mapped) + assert(result.schema == stage.transformSchema(original.schema)) + val rows = result.collect() + assert(rows.forall(_.getAs[Row]("error") == null)) + assert(rows.forall(!_.isNullAt(result.schema.fieldIndex("answer")))) + assert(server.requests.size() == 2) + } + } + } + + test("complete routes are not appended twice and contradictory API endpoints fail preflight") { + Using.resource(new AIRecordingServer((_, _) => 200 -> responseReply)) { server => + prompt(server, "responses").setUrl(server.url + "/v1/responses").transform(input.limit(1)).collect() + assert(server.requests.peek()._1 == "/v1/responses") + prompt(server, "responses").setCustomUrlRoot(server.url + "/proxy?version=custom") + .transform(input.limit(1)).collect() + assert(server.requests.asScala.last._1 == "/proxy?version=custom") + intercept[IllegalArgumentException] { + prompt(server, "responses").setCustomUrlRoot(server.url + "/v1/chat/completions") + .transformSchema(input.schema) + } + assert(server.requests.size() == 2) + } + } + + test("complete deployment routes reject scalar and row mismatches without sending") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val address = server.url + "/openai/deployments/explicit-model/chat/completions" + prompt(server).setUrl(address).transform(input.limit(1)).collect() + assert(server.requests.peek()._1.startsWith("/openai/deployments/explicit-model/chat/completions?")) + assert(server.requests.peek()._1.count(_ == '/') == 5) + intercept[IllegalArgumentException] { + prompt(server).setUrl(address).setDeploymentName("different").transformSchema(input.schema) + } + val row = prompt(server).setUrl(address).setDeploymentNameCol("text").transform(input.limit(1)).head() + assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + intercept[IllegalArgumentException] { + prompt(server).setUrl(server.url + "/models/chat/completions").transformSchema(input.schema) + } + assert(server.requests.size() == 1) + val models = prompt(server).setUrl(server.url + "/models/chat/completions") + models.clear(models.deploymentName) + models.setModel("model").transform(input.limit(1)).collect() + assert(server.requests.asScala.last._1.startsWith("/models/chat/completions?")) + assert(server.requests.size() == 2) + } + } + + test("map-backed chat content rejects empty arrays unsupported types and malformed text before HTTP") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + val rows = stage.transform(AIChatContentFixtures.mixed(spark)).collect() + val errors = rows.filter(_.getAs[Row]("error") != null) + assert(rows.length == 6 && errors.length == 5) + assert(errors.forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(rows.filter(_.getAs[Row]("error") == null).map(_.getAs[String]("answer")).toSeq == Seq("answer")) + assert(server.requests.size() == 1 && server.requests.peek()._2.contains("valid")) + } + } + + test("planning uses the input session resolver and restores the active session") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val session = spark.newSession() + session.conf.set("spark.sql.caseSensitive", "true") + val previous = org.apache.spark.sql.SparkSession.getActiveSession + org.apache.spark.sql.SparkSession.setActiveSession(spark) + try { + val schema = new StructType().add("text", StringType).add("TEXT", StringType) + val data = session.createDataFrame(Seq(Row("lower", "upper")).asJava, schema) + val stage = prompt(server).setPromptTemplate("{TEXT}") + stage.clear(stage.inputCol) + val result = stage.transform(data) + assert(result.head().getAs[String]("answer") == "answer") + assert(server.requests.peek()._2.contains("upper")) + assert(org.apache.spark.sql.SparkSession.getActiveSession.contains(spark)) + } finally { + previous match { + case Some(active) => org.apache.spark.sql.SparkSession.setActiveSession(active) + case None => org.apache.spark.sql.SparkSession.clearActiveSession() + } + } + } + } + + test("embedding user scalar and column reach HTTP and explicit deployment beats the embedding default") { + Using.resource(new AIRecordingServer((_, _) => 200 -> embeddingReply)) { server => + val previous = GlobalParams.getGlobalParam(OpenAIEmbeddingDeploymentNameKey) + GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, Left("wrong-global")) + try { + val stage = prompt(server, "embeddings").setUser("scalar-user") + val result = stage.transform(input.limit(1)) + assert(result.schema == stage.transformSchema(input.limit(1).schema)) + assert(result.head().getAs[Vector]("answer").toArray.toSeq == Seq(0.25, 0.5)) + assert(AIJson.objectString(server.requests.peek()._2).fields("user") == JsString("scalar-user")) + stage.setUserCol("text").transform(input.limit(1)).collect() + assert(AIJson.objectString(server.requests.asScala.last._2).fields("user") == JsString("hello")) + assert(server.requests.asScala.forall(r => + AIJson.objectString(r._2).fields("model") == JsString("explicit-model"))) + assert(GlobalParams.getGlobalParam(OpenAIEmbeddingDeploymentNameKey).contains(Left("wrong-global"))) + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, value) + case None => GlobalParams.resetGlobalParam(OpenAIEmbeddingDeploymentNameKey) + } + } + } + } + + test("endpoint routing uses case-insensitive hosts without overriding explicit API selection") { + val root = "https://RESOURCE.SERVICES.AI.AZURE.COM" + val stage = new AIPrompt().setUrl(root).setModel("model").setInputCol("text") + stage.transformSchema(input.schema) + val chat = new AIRequestAdapter(stage, input.schema) + val (chatUrl, _) = chat.request(Row("hello"), None, + Some(Seq(new org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema( + Array[Any]("user", "hello"), new StructType().add("role", StringType).add("content", StringType))))) + assert(chatUrl == root + "/models/chat/completions?api-version=2025-04-01-preview") + stage.setApiType("responses") + stage.transformSchema(input.schema) + val responses = new AIRequestAdapter(stage, input.schema) + val (responsesUrl, _) = responses.request(Row("hello"), None, + Some(Seq(new org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema( + Array[Any]("user", "hello"), new StructType().add("role", StringType).add("content", StringType))))) + assert(responsesUrl == root + "/openai/responses?api-version=2025-04-01-preview") + intercept[IllegalArgumentException] { + new AISystemOne().setUrl("https://API.OPENAI.COM").setDeploymentName("model") + .setStateCol("text").setQuestions(NativeRequestFixtures.SystemQuestions).transformSchema(input.schema) + } + intercept[IllegalArgumentException] { + new com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions().setUrl(root) + .setModel("model").setInputCol("text").setQuestions(NativeRequestFixtures.DecisionQuestions) + .transformSchema(input.schema) + } + } + + test("implicit Fabric credentials follow the effective destination and never a non-Fabric global URL") { + val previous = GlobalParams.getGlobalParam(URLKey) + val root = "https://fabric.example.invalid/cognitive/openai" + val reply = chatReply + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + GlobalParams.setGlobalParam(URLKey, server.url + "/v1") + val stage = new TestableAIFabricAuthPrompt("fabric-global-test").setDeploymentName("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + val row = stage.transform(input.limit(1)).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(!server.headers.peek().contains("authorization")) + } + val cases = Seq(root -> true, (root + "-other") -> false, + "https://other.example.invalid/cognitive/openai" -> false, + root.replace("https:", "http:") -> false, + root.replace(".invalid/", ".invalid:444/") -> false, + (root + "/../outside") -> false, (root + "/%2e%2e/outside") -> false) + cases.foreach { case (url, expectedFabric) => + GlobalParams.setGlobalParam(URLKey, url) + val stage = new TestableAIFabricAuthPrompt("fabric-destination-test").setDeploymentName("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + stage.setHandler { (_, request) => + assert(request.usesFabricAuth == expectedFabric, s"Unexpected Fabric fallback for $url") + assert(request.authorizationHeader == (if (expectedFabric) Some("MwcToken synthetic") else None)) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + assert(stage.transform(input.limit(1)).head().getAs[String]("answer") == "answer") + } + GlobalParams.setGlobalParam(URLKey, root) + Seq("url", "customRoot", "credential").foreach { explicit => + val stage = new TestableAIFabricAuthPrompt("fabric-explicit-test").setDeploymentName("model") + .setInputCol("text").setOutputCol("answer") + explicit match { + case "url" => stage.setUrl(root) + case "customRoot" => stage.setCustomUrlRoot(root + "/v1/chat/completions") + case "credential" => stage.setSubscriptionKey("synthetic-explicit-key") + } + stage.setHandler { (_, request) => + assert(!request.usesFabricAuth && request.authorizationHeader.isEmpty) + HTTPSchema.stringToResponse(reply, 200, "OK") + } + assert(stage.transform(input.limit(1)).head().getAs[String]("answer") == "answer") + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + + test("explicit model bindings ignore unused global deployment columns and scalar values") { + Using.resource(new AIRecordingServer((_, _) => 200 -> embeddingReply)) { server => + val previous = GlobalParams.getGlobalParam(OpenAIDeploymentNameKey) + val stage = prompt(server, "embeddings") + stage.clear(stage.deploymentName) + stage.setModel("explicit-model") + try { + GlobalParams.resetGlobalParam(OpenAIDeploymentNameKey) + val originalSchema = stage.transformSchema(input.schema) + Seq(Right("absent-global-column"), Left("")).foreach { default => + GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, default) + val effective = stage.effectiveCopy() + assert(effective.transformSchema(input.schema) == originalSchema) + val row = stage.transform(input.limit(1)).head() + assert(row.getAs[Row]("error") == null) + assert(AIJson.objectString(server.requests.asScala.last._2).fields("model") == JsString("explicit-model")) + stage.copy(ParamMap.empty).setModelCol("text").transform(input.limit(1)).collect() + assert(AIJson.objectString(server.requests.asScala.last._2).fields("model") == JsString("hello")) + assert(GlobalParams.getGlobalParam(OpenAIDeploymentNameKey).contains(default)) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, value) + case None => GlobalParams.resetGlobalParam(OpenAIDeploymentNameKey) + } + } + } + } + + test("malformed response format columns remain row errors without replaying valid requests") { + Seq("chat_completions" -> chatReply, "responses" -> responseReply).foreach { case (api, reply) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val data = spark.createDataFrame(Seq( + Row("valid", """{"type":"json_object"}"""), + Row("invalid number", """{"type":"json_schema","json_schema":42}"""), + Row("invalid null", """{"type":"json_schema","json_schema":null}"""), + Row("invalid array", """{"type":"json_schema","json_schema":[]}"""), + Row("invalid type", """{"type":null}""")).asJava, + new StructType().add("text", StringType).add("format", StringType)) + val rows = prompt(server, api).setResponseFormatCol("format").transform(data.coalesce(1)).collect() + assert(rows.head.getAs[Row]("error") == null) + assert(rows.tail.forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(server.requests.size() == 1) + } + } + } + + test("malformed row extensions and null explicit model produce terminal row errors without HTTP") { + Using.resource(new AIRecordingServer((_, _) => 200 -> embeddingReply)) { server => + val schema = new StructType().add("text", StringType).add("body", StringType).add("deployment", StringType) + val data = spark.createDataFrame(Seq( + Row("ok", """{"encoding_format":"float"}""", "row-model"), + Row("bad", """{"encoding_format":"base64"}""", "row-model"), + Row("duplicate", """{"x":1,"x":2}""", "row-model"), + Row("missing", "{}", Option.empty[String].orNull)).asJava, schema) + val stage = prompt(server, "embeddings").setDeploymentNameCol("deployment").setExtraBodyCol("body") + val result = stage.transform(data).collect() + assert(result.length == 4) + assert(result.count(_.getAs[Row]("error") != null) == 3) + assert(server.requests.size() == 1) + assert(result.last.getAs[Row]("error").getString(0).contains("missing_input")) + } + } + + test("deployment path segments reject dot routes and preserve encoded model names") { + Using.resource(new AIRecordingServer((_, _) => 200 -> embeddingReply)) { server => + val stage = prompt(server, "embeddings").setUrl(server.url) + val invalid = Seq(".", "..", "../other", "foo/../../admin", ".\\admin", "x\\..\\admin", + "%2e%2e%2fsecret", "..%252fsecret") + invalid.foreach { model => + intercept[IllegalArgumentException](stage.setDeploymentName(model).transformSchema(input.schema)) + } + val model = "a b/c+\u00e9" + stage.setDeploymentName(model).transform(input.limit(1)).collect() + assert(server.requests.peek()._1.startsWith("/openai/deployments/a%20b%2Fc%2B%C3%A9/embeddings?")) + val data = spark.createDataFrame((invalid :+ model).map(Row(_)).asJava, + new StructType().add("text", StringType)) + val rows = stage.setDeploymentNameCol("text").transform(data.coalesce(1)).collect() + assert(rows.dropRight(1).forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + assert(rows.last.getAs[Row]("error") == null && server.requests.size() == 2) + stage.setUrl(server.url + "/v1").setDeploymentName(".").transform(input.limit(1)).collect() + assert(AIJson.objectString(server.requests.asScala.last._2).fields("model") == JsString(".")) + } + } + + test("HTTP failures retain usage and identity without replacing the authoritative status diagnostic") { + val bodies = Seq( + """{"id":"failed-id","usage":{"total_tokens":3},"error":{"code":"rejected"}}""", + """{"id":42,"usage":{"total_tokens":3},"error":{"code":"rejected"}}""", + """{"id":"failed-id","usage":{"total_tokens":-1},"error":{"code":"rejected"}}""", + "not JSON") + Using.resource(new AIRecordingServer((n, _) => 400 -> bodies(n - 1))) { server => + val stage = prompt(server, "responses").setResponseIdCol("response_id") + val rows = bodies.map(_ => stage.transform(input.limit(1)).head()) + rows.foreach { row => + assert(row.getAs[Row]("error").getString(0).contains("http_error")) + assert(row.getAs[Row]("error").getStruct(1).getInt(1) == 400) + } + assert(rows.head.getAs[Row]("usage").getLong(2) == 3L) + assert(rows.head.getAs[String]("response_id") == "failed-id") + assert(rows(1).getAs[Row]("usage").getLong(2) == 3L && rows(1).getAs[String]("response_id") == null) + assert(rows(2).getAs[Row]("usage") == null && rows(2).getAs[String]("response_id") == "failed-id") + assert(rows.last.getAs[String]("raw") == "not JSON") + } + } + + test("invalid scalar settings fail before actions and never send requests") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val invalid = Seq( + prompt(server).setTemperature(0.5).setExtraBody(Map("temperature" -> 0.5)), + prompt(server).setMaxCompletionTokens(5).setExtraBody(Map("max_tokens" -> 5)), + prompt(server).setExtraBody(Map("model" -> "other")), + prompt(server).setExtraBody(Map("stream" -> true)), + prompt(server).setExtraBody(Map("background" -> true)), + prompt(server, "embeddings").setExtraBody(Map("encoding_format" -> "base64")), + prompt(server, "responses").setN(2)) + invalid.foreach(stage => intercept[IllegalArgumentException](stage.transform(input))) + assert(server.requests.isEmpty) + } + } + + test("invalid endpoint and typed column bindings fail before any HTTP or external input reads") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq("not a URL", "file:///tmp/model", "https://user:password@example.invalid", + "https://example.invalid#fragment", "https://example.invalid?version=ignored").foreach { url => + intercept[IllegalArgumentException](prompt(server).setUrl(url).transformSchema(input.schema)) + } + val data = input.withColumn("wrong_type", org.apache.spark.sql.functions.lit(1)) + val invalid = Seq( + prompt(server).setTemperatureCol("wrong_type"), + prompt(server).setStoreCol("wrong_type").setApiType("responses"), + prompt(server).setUserCol("absent"), + prompt(server).setModel("conflicts-with-explicit-deployment"), + prompt(server).setText("embedding-only"), + prompt(server).setStateCol("text"), + prompt(server, "embeddings").setSystemPrompt("unsupported"), + prompt(server, "embeddings").setDimensions(-1), + prompt(server, "responses").setApiVersion("unsupported-on-v1"), + prompt(server).setPromptTemplate("{text}")) + invalid.foreach(stage => intercept[IllegalArgumentException](stage.transformSchema(data.schema))) + assert(server.requests.isEmpty) + } + } + + test("native scalar credentials and row headers reach the selected protocol without entering the body") { + val cases = Seq( + ("systemone", NativeRequestFixtures.SystemResponse, NativeRequestFixtures.SystemQuestions, "api-key"), + ("decisions", NativeRequestFixtures.DecisionResponse, NativeRequestFixtures.DecisionQuestions, "authorization")) + cases.foreach { case (api, reply, questions, keyHeader) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = prompt(server, api).setQuestions(questions).setSubscriptionKey("synthetic-api-key") + if (api == "systemone") stage.setUrl(server.url) + stage.transform(input.limit(1)).collect() + val expected = if (api == "decisions") "Bearer synthetic-api-key" else "synthetic-api-key" + assert(server.headers.peek()(keyHeader) == expected) + assert(!server.requests.peek()._2.contains("synthetic-api-key")) + stage.clear(stage.subscriptionKey) + stage.setAADTokenCol("text").transform(input.limit(1)).collect() + assert(server.headers.asScala.last("authorization") == "Bearer hello") + } + } + } + + test("nested extension siblings merge but explicitly owned format subtree cannot be replaced") { + Using.resource(new AIRecordingServer((_, _) => 200 -> responseReply)) { server => + val stage = prompt(server, "responses").setResponseFormat("json_object") + .setExtraBody(Map("text" -> Map("verbosity" -> "low"))) + stage.transform(input.limit(1)).collect() + val text = AIJson.objectString(server.requests.peek()._2).fields("text").asJsObject + assert(text.fields.keySet == Set("format", "verbosity")) + stage.setExtraBody(Map("text" -> Map("format" -> Map("type" -> "json_object")))) + intercept[IllegalArgumentException](stage.transform(input)) + assert(server.requests.size() == 1) + } + } + + test("output roles replace inputs in place and reject resolver aliases") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val metadata = new MetadataBuilder().putString("description", "retained input metadata").build() + val data = spark.createDataFrame(Seq(Row("hello", 7, "keep")).asJava, + new StructType().add("text", StringType).add("answer", IntegerType) + .add("a.b", StringType, nullable = false, metadata)) + val stage = prompt(server) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + assert(result.schema.fieldNames.take(3).toSeq == Seq("text", "answer", "a.b")) + assert(result.head().getAs[String]("a.b") == "keep") + val old = spark.conf.get("spark.sql.caseSensitive") + try { + spark.conf.set("spark.sql.caseSensitive", "false") + intercept[IllegalArgumentException](stage.setUsageCol("ANSWER").transformSchema(data.schema)) + spark.conf.set("spark.sql.caseSensitive", "true") + assert(stage.transformSchema(data.schema).fieldNames.contains("ANSWER")) + } finally spark.conf.set("spark.sql.caseSensitive", old) + } + } + + test("missing input stays in output and bounded retry counts physical sends") { + Using.resource(new AIRecordingServer((attempt, _) => + if (attempt == 1) 429 -> """{"error":{"message":"retry"}}""" else 200 -> chatReply)) { server => + val data = spark.createDataFrame(Seq(Row(Option.empty[String].orNull), Row("valid")).asJava, + new StructType().add("text", StringType)) + val rows = prompt(server).transform(data.coalesce(1)).collect() + assert(rows.length == 2) + assert(rows.head.getAs[Row]("error").getString(0).contains("missing_input")) + assert(rows.last.getAs[String]("answer") == "answer") + assert(server.requests.size() == 2) + } + } + + test("scalar and column setters obey last setter wins and copy/save/load preserve behavior") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = prompt(server).setExtraBodyCol("does_not_exist").setExtraBody(Map("custom" -> true)) + val copied = stage.copy(ParamMap.empty) + val directory = Files.createTempDirectory("aiprompt-params-").resolve("stage").toString + try { + copied.write.session(spark).save(directory) + val loaded = AIPrompt.load(directory) + assert(loaded.getExtraBody == stage.getExtraBody) + assert(!loaded.isSet(loaded.store)) + assert(!loaded.isSet(loaded.postProcessing)) + assert(!loaded.isSet(loaded.handler)) + assert(loaded.transform(input.limit(1)).head().getAs[String]("answer") == "answer") + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(new java.io.File(directory).getParentFile) + } + } + } + + test("new API selectors do not change OpenAIPrompt's accepted API types") { + intercept[IllegalArgumentException](new OpenAIPrompt().setApiType("systemone")) + assert(new AISystemOne().getApiType == "systemone") + assert(new com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions().getApiType == "decisions") + intercept[IllegalArgumentException](new AISystemOne().setApiType("chat_completions")) + } + + test("JSON-valued body parameters preserve exact numbers and JVM arrays through save load") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val extension = Map[String, Any]( + "array" -> Array(1, 2), + "decimal" -> BigDecimal("0.12345678901234567890123456789"), + "integer" -> BigInt("9223372036854775808123456789"), + "nested" -> java.util.Arrays.asList("value", java.lang.Boolean.TRUE)) + val expected = AIJson.fromAny(extension) + val stage = prompt(server).setExtraBody(extension) + val directory = Files.createTempDirectory("aiprompt-json-params-").resolve("stage").toString + try { + stage.copy(ParamMap.empty).write.session(spark).save(directory) + val loaded = AIPrompt.load(directory) + assert(AIJson.fromAny(loaded.getExtraBody) == expected) + val output = stage.transformSchema(input.schema) + assert(loaded.transformSchema(input.schema) == output) + assert(loaded.transform(input.limit(1)).head().getAs[String]("answer") == "answer") + val body = AIJson.objectString(server.requests.peek()._2) + expected.asJsObject.fields.foreach { case (key, value) => assert(body.fields(key) == value) } + val questions = """[{"type":"score","minimum":0.12345678901234567890123456789}]""" + stage.setQuestions(questions) + val encoded = stage.questions.jsonEncode(stage.get(stage.questions).get) + val decoded = stage.questions.jsonDecode(encoded).left.get + assert(AIJson.fromAny(decoded) == AIJson.parse(questions)) + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(new java.io.File(directory).getParentFile) + } + } + } + + test("all supported chat scalar and column options reach their wire fields") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = prompt(server).setMaxTokens(32).setTemperature(0.7).setTopP(0.8).setUser("caller") + .setN(2).setStop("END").setPresencePenalty(0.1).setFrequencyPenalty(0.2).setSeed(7) + .setVerbosity("low").setReasoningEffort("low").setResponseFormat("json_object") + stage.transform(input.limit(1)).collect() + val scalar = AIJson.objectString(server.requests.peek()._2).fields + val expected = Map("max_completion_tokens" -> JsNumber(32), "temperature" -> JsNumber(0.7), + "top_p" -> JsNumber(0.8), "user" -> JsString("caller"), "n" -> JsNumber(2), "stop" -> JsString("END"), + "presence_penalty" -> JsNumber(0.1), "frequency_penalty" -> JsNumber(0.2), "seed" -> JsNumber(7), + "verbosity" -> JsString("low"), "reasoning_effort" -> JsString("low"), + "response_format" -> JsObject("type" -> JsString("json_object"))) + expected.foreach { case (key, value) => assert(scalar(key) == value, key) } + val data = spark.createDataFrame(Seq(Row("row input", 11, 0.3, 0.4, "row-user", 1, + "STOP", 0.5, 0.6, 42, "high", "high", """{"type":"json_object"}""")).asJava, + new StructType().add("text", StringType).add("tokens", IntegerType) + .add("temperature", DoubleType).add("top", DoubleType).add("user", StringType) + .add("n", IntegerType).add("stop", StringType).add("presence", DoubleType) + .add("frequency", DoubleType).add("seed", IntegerType).add("verbosity", StringType) + .add("reasoning", StringType).add("format", StringType)) + stage.setMaxTokensCol("tokens").setTemperatureCol("temperature").setTopPCol("top").setUserCol("user") + .setNCol("n").setStopCol("stop").setPresencePenaltyCol("presence").setFrequencyPenaltyCol("frequency") + .setSeedCol("seed").setVerbosityCol("verbosity").setReasoningEffortCol("reasoning") + .setResponseFormatCol("format").transform(data).collect() + val column = AIJson.objectString(server.requests.asScala.last._2).fields + assert(column("max_completion_tokens") == JsNumber(11)) + assert(column("temperature") == JsNumber(0.3) && column("top_p") == JsNumber(0.4)) + assert(column("user") == JsString("row-user") && column("n") == JsNumber(1)) + assert(column("stop") == JsString("STOP") && column("presence_penalty") == JsNumber(0.5)) + assert(column("frequency_penalty") == JsNumber(0.6) && column("seed") == JsNumber(42)) + assert(column("verbosity") == JsString("high") && column("reasoning_effort") == JsString("high")) + assert(column("response_format") == expected("response_format")) + } + } + + test("Responses column options map to nested wire paths and do not rewrite postprocessing") { + Using.resource(new AIRecordingServer((_, _) => 200 -> responseReply)) { server => + val data = spark.createDataFrame(Seq(Row("hello", true, "prior", """{"type":"json_object"}""")).asJava, + input.schema.add("store", BooleanType).add("prior", StringType).add("format", StringType)) + val stage = prompt(server, "responses").setStoreCol("store").setPreviousResponseIdCol("prior") + .setResponseFormatCol("format").setVerbosity("low").setReasoningEffort("medium").setMaxCompletionTokens(20) + val result = stage.transform(data).head() + assert(result.getAs[String]("answer") == "onetwothree") + val body = AIJson.objectString(server.requests.peek()._2).fields + assert(body("store") == JsTrue && body("previous_response_id") == JsString("prior")) + assert(body("max_output_tokens") == JsNumber(20)) + assert(body("text").asJsObject.fields == Map("format" -> JsObject("type" -> JsString("json_object")), + "verbosity" -> JsString("low"))) + assert(body("reasoning") == JsObject("effort" -> JsString("medium"))) + } + } + + test("postprocessing and retained messages have matching planned schemas and safe scratch columns") { + Using.resource(new AIRecordingServer((_, _) => 200 -> + chatReply.replace("answer", """{\"value\":\"ok\"}"""))) { server => + val stage = prompt(server).setPromptTemplate("Read {`a.b`}").setDropPrompt(false) + .setPostProcessingOptions(Map("jsonSchema" -> "value STRING")) + stage.clear(stage.inputCol) + val data = spark.createDataFrame(Seq(Row("hello")).asJava, new StructType().add("a.b", StringType)) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null) + assert(row.getAs[Row]("answer").getAs[String]("value") == "ok") + assert(row.getSeq[Row](row.fieldIndex(stage.getMessagesCol)).nonEmpty) + } + } + + test("unsupported message fields become explicit row errors instead of disappearing") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val message = new StructType().add("role", StringType).add("content", StringType).add("tool_call_id", StringType) + val data = spark.createDataFrame(Seq(Row(Seq(Row("tool", "answer", "call-1")))).asJava, + new StructType().add("messages", ArrayType(message))) + val stage = prompt(server).setMessagesCol("messages") + stage.clear(stage.inputCol) + assert(stage.transform(data).head().getAs[Row]("error").getString(0).contains("invalid_request")) + assert(server.requests.isEmpty) + } + } + + test("raw message arrays and header maps use Spark-produced collections without losing names") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val message = new StructType().add("role", StringType).add("content", StringType).add("name", StringType) + val data = spark.createDataFrame(Seq(Row(Seq(Row("user", "hello", "caller")), + Map("X-Test-Trace" -> "recording"))).asJava, + new StructType().add("messages", ArrayType(message)) + .add("headers", MapType(StringType, StringType))) + val stage = prompt(server).setMessagesCol("messages") + stage.setVectorParam(stage.customHeaders, "headers") + stage.clear(stage.inputCol) + val mapped = data.map { row => + Row(Seq(Row("user", "hello", "caller")), row.getMap[String, String](1).toMap) + }(Encoders.row(data.schema)) + assert(stage.transform(mapped).head().getAs[String]("answer") == "answer") + assert(server.headers.peek()("x-test-trace") == "recording") + val messages = AIJson.objectString(server.requests.peek()._2).fields("messages").asInstanceOf[JsArray] + assert(messages.elements.head.asJsObject.fields("name") == JsString("caller")) + } + } + + test("schema-less nested messages preserve inline media and map-backed content") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val part = new StructType().add("type", StringType).add("text", StringType) + .add("image_url", MapType(StringType, StringType)) + val message = new StructType().add("role", StringType).add("content", ArrayType(part)) + val schema = new StructType().add("messages", ArrayType(message)) + val image = "data:image/png;base64,AQID" + val data = spark.range(1).map { _ => + Row(Seq(Row("user", Seq(Row("image_url", Option.empty[String].orNull, Map("url" -> image)))))) + }(Encoders.row(schema)) + val stage = prompt(server).setMessagesCol("messages") + stage.clear(stage.inputCol) + assert(stage.transform(data).head().getAs[Row]("error") == null) + val messages = AIJson.objectString(server.requests.peek()._2).fields("messages").asInstanceOf[JsArray] + val content = messages.elements.head.asJsObject.fields("content").asInstanceOf[JsArray] + assert(content.elements.head.asJsObject.fields("image_url").asJsObject.fields("url") == JsString(image)) + } + } + + test("filename scratch columns never replace inputs or consume a public answer") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val data = spark.createDataFrame( + Seq(Row("data:text/plain;base64,aGVsbG8=", "original", "also original")).asJava, + new StructType().add("file.name", StringType).add("file.name_filename", StringType) + .add("file.name_filename_1", StringType)) + val stage = prompt(server).setPromptTemplate("Read {`file.name`}") + .setColumnTypes(Map("file.name" -> "path")).setOutputCol("file.name_filename_2") + stage.clear(stage.inputCol) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val row = result.head() + assert(row.getAs[String]("file.name_filename") == "original") + assert(row.getAs[String]("file.name_filename_1") == "also original") + assert(row.getAs[String]("file.name_filename_2") == "answer") + assert(row.getAs[Row]("error") == null) + assert(server.requests.peek()._2.contains("Read attachment.txt")) + assert(!result.columns.contains("file.name_filename_3")) + } + } + + test("postprocessing failure retains raw content with a distinct parse diagnostic") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = prompt(server).setPostProcessingOptions(Map("jsonSchema" -> "ARRAY")) + val result = stage.transform(input.limit(1)) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(row.isNullAt(row.fieldIndex("answer"))) + assert(row.getAs[Row]("error").getString(0).contains("parse_error")) + assert(row.getAs[String]("raw").contains("answer")) + } + } + + test("malformed provider answers retain valid usage and raw JSON on the public output") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply.replace("\"answer\"", "42"))) { server => + val row = prompt(server).transform(input.limit(1)).head() + assert(row.getAs[String]("answer") == None.orNull) + assert(row.getAs[Row]("error").getString(0).contains("parse_error")) + assert(row.getAs[Row]("usage").getLong(2) == 3L) + assert(row.getAs[String]("raw").contains("\"content\":42")) + assert(server.requests.size() == 1) + } + } + + test("template schema planning resolves literal names and rejects missing columns without HTTP") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = prompt(server).setPromptTemplate("Read {missing}") + stage.clear(stage.inputCol) + val error = intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + assert(error.getMessage.startsWith("Invalid prompt template") && error.getCause == null) + assert(server.requests.isEmpty) + } + } + + test("bounded 429 retries terminate and preserve the final raw response") { + Using.resource(new AIRecordingServer((_, _) => 429 -> """{"error":{"code":"rate_limit"}}""")) { server => + val row = prompt(server).transform(input.limit(1)).head() + assert(server.requests.size() == 4) + assert(row.getAs[Row]("error").getString(0).contains("http_error")) + assert(row.getAs[String]("raw").contains("rate_limit")) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestBase.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestBase.scala new file mode 100644 index 00000000000..f70919441a2 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestBase.scala @@ -0,0 +1,20 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase + +private[ai] abstract class AIPromptTestBase extends TestBase { + protected def path(): String = tmpDir.resolve(java.util.UUID.randomUUID().toString).toAbsolutePath.toString + + protected def withServer(test: AIPromptRecordingServer => Unit): Unit = { + val server = new AIPromptRecordingServer + try test(server) finally server.close() + } + + protected def prompt(server: AIPromptRecordingServer): AIPrompt = + new AIPrompt().setUrl(server.url + "/v1").setDeploymentName("local-model").setSubscriptionKey("synthetic-key") + .setPromptTemplate("{text}").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setRawResponseCol("raw") +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestFixtures.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestFixtures.scala new file mode 100644 index 00000000000..1948eaa3650 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestFixtures.scala @@ -0,0 +1,149 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.fuzzing.TransformerFuzzing +import com.microsoft.azure.synapse.ml.param.RPrinter +import com.sun.net.httpserver.{HttpExchange, HttpServer} +import org.apache.commons.io.IOUtils +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.types.{ArrayType, MapType, StringType, StructType} +import spray.json._ + +import java.net.InetSocketAddress +import java.nio.charset.StandardCharsets +import java.util.concurrent.{ConcurrentLinkedQueue, Executors} +import java.util.concurrent.atomic.AtomicInteger +import scala.collection.JavaConverters._ +import scala.util.Using + +private[ai] final class AIRecordingServer(reply: (Int, String) => (Int, String), + retryAfter: String = "0", + responseHeaders: Map[String, String] = Map.empty, + encodeResponse: String => Array[Byte] = + _.getBytes(StandardCharsets.UTF_8)) extends AutoCloseable { + private val server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + private val pool = Executors.newCachedThreadPool() + val requests = new ConcurrentLinkedQueue[(String, String)]() + val headers = new ConcurrentLinkedQueue[Map[String, String]]() + server.setExecutor(pool) + server.createContext("/", (exchange: HttpExchange) => { + val body = Using.resource(exchange.getRequestBody) { in => + new String(IOUtils.toByteArray(in), StandardCharsets.UTF_8) + } + requests.add(exchange.getRequestURI.toString -> body) + headers.add(exchange.getRequestHeaders.asScala.map { case (name, values) => + name.toLowerCase(java.util.Locale.ROOT) -> values.asScala.mkString(",") + }.toMap) + val (status, response) = reply(requests.size(), body) + val bytes = encodeResponse(response) + exchange.getResponseHeaders.add("Content-Type", "application/json") + responseHeaders.foreach { case (name, value) => exchange.getResponseHeaders.add(name, value) } + if (status == 429) exchange.getResponseHeaders.add("Retry-After", retryAfter) + exchange.sendResponseHeaders(status, bytes.length) + Using.resource(exchange.getResponseBody)(_.write(bytes)) + exchange.close() + }) + server.start() + def url: String = s"http://127.0.0.1:${server.getAddress.getPort}" + override def close(): Unit = { + server.stop(0) + pool.shutdownNow() + } +} + +private[ai] class AIPromptRecordingServer extends AutoCloseable { + val requests: AtomicInteger = new AtomicInteger() + private val server = new AIRecordingServer((_, body) => { + requests.incrementAndGet() + if (body.contains("ROW_ERROR")) { + 400 -> """{"error":{"code":"invalid_request","message":"synthetic row error"}}""" + } else { + 200 -> """{"id":"local-response","model":"local-model","choices":[{"index":0, + |"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}], + |"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}""".stripMargin + } + }) + val url: String = server.url + override def close(): Unit = server.close() +} + +private[ai] class TestableAIFabricAuthPrompt(override val uid: String) extends AIPrompt(uid) { + override private[ai] def fabricEndpoint: Option[String] = + Some("https://fabric.example.invalid/cognitive/openai") + + override private[ai] def fabricAuthHeader: String = "MwcToken synthetic" + + override private[ai] def assertFabricModelStatus(model: String): Unit = + require(model != "blocked-model", "Synthetic Fabric model is unavailable") +} + +private[ai] class TestableAIThrowingFabricAuthPrompt(override val uid: String) + extends TestableAIFabricAuthPrompt(uid) { + override private[ai] def fabricAuthHeader: String = + throw new IllegalStateException("The unused Fabric fallback must not be evaluated") +} + +private[ai] object AIChatContentFixtures { + def mixed(spark: org.apache.spark.sql.SparkSession): DataFrame = { + val content = ArrayType(MapType(StringType, StringType)) + val schema = new StructType().add("messages", + ArrayType(new StructType().add("role", StringType).add("content", content))) + val parts = Seq(Seq.empty[Map[String, String]], Seq(Map("type" -> "text", "text" -> "valid")), + Seq(Map("type" -> "audio", "data" -> "AQID")), Seq(Map("type" -> "unknown", "text" -> "not supported")), + Seq(Map("type" -> "text")), Seq(Map("type" -> "text", "text" -> "hello", "extra" -> "lost"))) + spark.createDataFrame(parts.map(content => Row(Seq(Row("user", content)))).asJava, schema) + } +} + +private[services] object AILiveTestEnvironment { + def get(name: String): Option[String] = sys.env.get(name).map(_.trim).filter(_.nonEmpty) +} + +private[services] object NativeRequestFixtures { + val SystemQuestions: String = + """{"repeat_contact":{"type":"noul","instructions":"Has the customer contacted support before?"}, + |"team":{"type":"choice","instructions":"Which team should handle this request?", + |"criteria":{"billing":"Charges, invoices and refunds","engineering":"Bugs, errors and outages"}}, + |"severity":{"type":"score","instructions":"How severe is this issue?", + |"criteria":["Cosmetic","Major","Critical"]}}""".stripMargin + val DecisionQuestions: String = + """[{"type":"predicate","name":"repeat_contact","instructions":"Has the customer contacted support before?"}, + |{"type":"choice","name":"team","instructions":"Which team should handle this request?", + |"choices":[{"value":"billing","description":"Charges, invoices and refunds"}, + |{"value":"engineering","description":"Bugs, errors and outages"}]}, + |{"type":"score","name":"severity","instructions":"How severe is this issue?", + |"levels":[{"label":"Cosmetic","description":"Appearance only"}, + |{"label":"Major","description":"A task fails but a workaround exists"}, + |{"label":"Critical","description":"A task is fully blocked"}]}]""".stripMargin + val SystemResponse: String = + """{"model":"microsoft-decision-1","answers":{"repeat_contact":true, + |"team":{"type":"choice","choice":"engineering","additional":[1,false,null]}, + |"severity":1.7},"usage":{"input_tokens":20,"output_tokens":0,"total_tokens":20},"future":true}""".stripMargin + val DecisionResponse: String = + """{"model":"gpt-6-luna","answers":[{"type":"predicate","name":"repeat_contact","probability":0.98}, + |{"type":"choice","name":"team","choice":"engineering","confidence":0.9, + |"probabilities":[{"value":"billing","probability":0.05},{"value":"engineering","probability":0.95}]}, + |{"type":"score","name":"severity","score":1.7,"confidence":0.8,"probabilities":[ + |{"value":0,"label":"Cosmetic","probability":0.1},{"value":1,"label":"Major","probability":0.1}, + |{"value":2,"label":"Critical","probability":0.8}]}], + |"usage":{"input_tokens":20,"output_tokens":0,"total_tokens":20}}""".stripMargin +} + +private[services] trait AIPromptFuzzing[S <: AIPrompt] extends TransformerFuzzing[S] { + override def rTestInstantiateModel(stage: S, num: Int): String = { + // Reuse the UID so native-incompatible defaults need not become explicit constructor arguments. + val lines = super.rTestInstantiateModel(stage, num).split("\n") + val uidLines = lines.indices.filter(index => lines(index).trim.startsWith(",uid = ")) + require(uidLines.length == 1, "Expected exactly one generated R constructor UID") + lines.updated(uidLines.head, s",uid = ${RPrinter(JsString(stage.uid))})").mkString("\n") + } + + override def pyTestInstantiateModel(stage: S, num: Int): String = { + // Keep UID-derived defaults implicit when the generator falls back to explicit Params. + val columns = JsObject("messagesCol" -> JsString(stage.getMessagesCol), + "responseIdCol" -> JsString(stage.getResponseIdCol)).compactPrint + super.pyTestInstantiateModel(stage, num) + s"model._setDefault(**$columns)\n" + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptUnicodeSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptUnicodeSuite.scala new file mode 100644 index 00000000000..fd4bfcd0e3c --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptUnicodeSuite.scala @@ -0,0 +1,240 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.Row +import spray.json.{JsArray, JsNumber, JsObject, JsString} + +import java.nio.file.Files +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptUnicodeSuite extends TestBase { + test("native metadata rejects malformed explicit and default string Params before saving") { + import spark.implicits._ + val data = Seq("hello").toDF("text") + val saved = Files.createTempDirectory(tmpDir, "unicode-metadata").resolve("model").toString + new AIPrompt().setSystemPrompt("keep").write.session(spark).save(saved) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + for (name <- Seq("systemPrompt", "promptTemplate"); binding <- Seq("explicit", "default", "shadowed"); + malformed <- Seq("private-\uD800", "private-\uDC00")) { + val stage = new TestableAIMutableJsonDefaultPrompt("unicode-metadata") + .setUrl(server.url + "/v1").setModel("model").setInputCol("text").setErrorCol("error") + val param = if (name == "systemPrompt") stage.systemPrompt else stage.promptTemplate + if (binding == "explicit") stage.set(param, malformed) + else { + stage.replaceJsonDefault(param, malformed) + if (binding == "shadowed") stage.set(param, "valid") + } + val writer = stage.write.session(spark).overwrite() + val error = intercept[IllegalArgumentException](writer.save(saved)) + assert(error.getMessage.contains("UTF-8") && !error.getMessage.contains("private-")) + assert(error.getCause == null) + assert(AIPrompt.read.session(spark).load(saved).getSystemPrompt == "keep") + val fresh = Files.createTempDirectory(tmpDir, "unicode-rejected").resolve("model") + intercept[IllegalArgumentException](writer.save(fresh.toString)) + assert(!Files.exists(fresh)) + if (name == "systemPrompt" && binding != "shadowed") { + val row = stage.transform(data).head() + assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + } + } + assert(server.requests.isEmpty) + } + } + + test("native metadata preserves valid Unicode and literal escape text through explicit and default Params") { + import spark.implicits._ + val data = Seq("hello").toDF("text") + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + for (value <- Seq("\\" + "ud800", "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD"); explicit <- Seq(false, true)) { + val stage = new TestableAIMutableJsonDefaultPrompt("valid-unicode-metadata") + .setUrl(server.url + "/v1").setModel("model").setOutputCol("answer").setErrorCol("error") + if (explicit) stage.setSystemPrompt(value).setPromptTemplate(value) + else { + stage.replaceJsonDefault(stage.systemPrompt, value).replaceJsonDefault(stage.promptTemplate, value) + stage.setInputCol("text") + } + val saved = Files.createTempDirectory(tmpDir, "valid-unicode-metadata").resolve("model").toString + stage.write.session(spark).save(saved) + val loaded = AIPrompt.read.session(spark).load(saved) + Seq(stage, loaded).foreach { candidate => + assert(candidate.getSystemPrompt == value && candidate.getPromptTemplate == value) + assert(candidate.isSet(candidate.systemPrompt) == explicit) + assert(candidate.isSet(candidate.promptTemplate) == explicit) + val row = candidate.transform(data).head() + assert(row.getAs[Row]("error") == null && row.getAs[String]("answer") == "ok") + } + val requests = server.requests.asScala.map(request => AIJson.objectString(request._2)).toVector.takeRight(2) + assert(requests.length == 2 && requests.forall { body => + body.fields("messages") match { + case JsArray(messages) => messages.head.asJsObject.fields("content") == JsString(value) + case _ => fail("Expected a request message array") + } + }) + } + assert(server.requests.size() == 8) + } + } + + test("prompt template Unicode is validated before Spark literal encoding") { + import spark.implicits._ + val data = Seq("row").toDF("text") + val unicode = "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD" + Seq(new AIPrompt(), new AIPrompt().setApiType("responses"), + new AISystemOne(), new OpenAIDecisions()).foreach { stage => + val reply = stage.getApiType match { + case "systemone" => NativeRequestFixtures.SystemResponse + case "decisions" => NativeRequestFixtures.DecisionResponse + case "responses" => + """{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}]}""" + case _ => """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val url = if (stage.getApiType == "systemone") server.url else server.url + "/v1" + stage.setUrl(url).setModel("model").setOutputCol("answer").setErrorCol("error") + if (stage.getApiType == "systemone") { + stage.setQuestions(NativeRequestFixtures.SystemQuestions).setInputMode("json") + } else if (stage.getApiType == "decisions") { + stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + } + Seq("\uD800", "\uDC00").foreach { malformed => + stage.setPromptTemplate("\"private-template-" + malformed + "\"") + Seq[() => Any](() => stage.transformSchema(data.schema), () => stage.transform(data).collect()) + .foreach { evaluate => + val error = intercept[IllegalArgumentException](evaluate()) + assert(error.getMessage.contains("template") && !error.getMessage.contains("private-template")) + assert(error.getCause == null) + } + } + assert(server.requests.isEmpty) + val result = stage.setPromptTemplate(JsString(unicode).compactPrint).transform(data).head() + assert(result.getAs[Row]("error") == null && !result.isNullAt(result.fieldIndex("answer"))) + assert(server.requests.size() == 1 && server.requests.peek()._2.contains(unicode)) + if (stage.getApiType == "systemone") { + assert(AIJson.objectString(server.requests.peek()._2).fields("state") == JsString(unicode)) + } + } + } + } + + test("JSON surrogate validation rejects unpaired strings and keys across native setters and persistence") { + val invalid = Seq("\uD800", "\uDC00", "\uDC00\uD800", "\uD800x", "x\uDC00") + invalid.foreach { text => + val documents = Seq(JsObject("private-field" -> JsString(text)), JsObject(text -> JsString("private-value"))) + documents.foreach { document => + val values = Seq[Any](document, AIJson.AnyFormat.read(document), + AIJson.AnyFormat.read(document).asInstanceOf[Map[String, Any]].asJava) + values.foreach { value => + val error = intercept[IllegalArgumentException](AIJson.fromAny(value)) + assert(error.getMessage.contains("surrogate") && !error.getMessage.contains("private-")) + } + Seq(new AIPrompt(), new AISystemOne(), new OpenAIDecisions()).foreach { stage => + Seq[() => AIPrompt]( + () => stage.setExtraBody(document), () => stage.setQuestions(document), + () => stage.setResponseFormat(AIJson.AnyFormat.read(document).asInstanceOf[Map[String, Any]]) + ).foreach { setter => + assert(intercept[IllegalArgumentException](setter()).getMessage.contains("surrogate")) + } + } + } + } + val high = "\\" + "ud800" + val low = "\\" + "udc00" + Seq(s"""{"private-field":"$high"}""", s"""{"private-field":"$low"}""", + s"""{"$high":1,"?":2}""", s"""{"$low":1,"?":2}""").foreach { json => + assert(intercept[IllegalArgumentException](AIJson.objectString(json)).getMessage.contains("surrogate")) + val stage = new AIPrompt() + assert(intercept[IllegalArgumentException](stage.extraBody.jsonDecode(s"""{"left":$json}""")) + .getMessage.contains("surrogate")) + } + val unicode = "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD" + val stage = new AIPrompt().setExtraBody(Map(unicode -> Seq(unicode))) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-unicode").resolve("model").toString + stage.write.session(spark).save(saved) + Seq(stage, stage.copy(ParamMap.empty), AIPrompt.load(saved)).foreach { candidate => + assert(candidate.getExtraBody == Map(unicode -> Seq(unicode))) + } + val escaped = """"caf""" + "\\" + "u00e9 " + "\\" + "u96ea " + "\\" + "ud83d" + "\\" + "ude00 " + + "\\" + "ufffd\"" + assert(AIJson.parse(escaped) == JsString(unicode)) + } + + test("JSON surrogate validation rejects malformed per-row bodies without colliding wire keys") { + import spark.implicits._ + val high = "\\" + "ud800" + val low = "\\" + "udc00" + val pair = "\\" + "ud83d" + "\\" + "ude00" + val invalid = Seq(s"""{"private-field":"$high"}""", s"""{"private-field":"$low"}""", + s"""{"$high":1,"?":2}""", s"""{"$low":1,"?":2}""") + val valid = s"""{"$pair":"$pair","?":2}""" + val data = (invalid :+ valid).zipWithIndex.map { case (body, id) => (id, "hello", body) } + .toDF("id", "text", "body") + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setExtraBodyCol("body").setOutputCol("answer").setErrorCol("error") + val rows = stage.transform(data).collect().sortBy(_.getAs[Int]("id")) + rows.take(invalid.size).foreach { row => + val diagnostic = row.getAs[Row]("error").getString(0) + assert(diagnostic.contains("invalid_request") && diagnostic.contains("surrogate")) + assert(!diagnostic.contains("private-field") && row.getAs[String]("answer") == null) + } + assert(rows.last.getAs[Row]("error") == null && server.requests.size() == 1) + val sent = AIJson.objectString(server.requests.peek()._2) + assert(sent.fields("\uD83D\uDE00") == JsString("\uD83D\uDE00") && sent.fields("?") == JsNumber(2)) + stage.clear(stage.extraBody) + Seq("\uD800", "\uDC00").foreach { malformed => + val row = stage.setSystemPrompt(malformed).transform(data.limit(1)).head() + assert(row.getAs[Row]("error").getString(0).contains("invalid_request")) + } + assert(server.requests.size() == 1) + } + } + + test("structured response JSON surrogate diagnostics cover decoded values and keys in both generative APIs") { + import spark.implicits._ + val high = "\\" + "ud800" + val low = "\\" + "udc00" + val pair = "\\" + "ud83d" + "\\" + "ude00" + val documents = Map( + "value-high" -> s"""{"value":"$high"}""", + "value-low" -> s"""{"value":"$low"}""", + "key-high" -> s"""{"$high":"private-value","?":"other"}""", + "key-low" -> s"""{"$low":"private-value","?":"other"}""", + "valid" -> s"""{'value':'$pair','$pair':'unchanged'}""") + Seq("chat_completions", "responses").foreach { api => + Using.resource(new AIRecordingServer((_, body) => { + val document = documents.find { case (name, _) => body.contains(name) }.get._2 + val text = JsString(document).compactPrint + 200 -> (if (api == "chat_completions") { + s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$text}}]}""" + } else { + s"""{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":$text}]}]}""" + }) + })) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> "MAP")) + val rows = stage.transform(documents.keys.toSeq.toDF("text")).collect() + rows.filter(_.getAs[String]("text") != "valid").foreach { row => + assert(row.getAs[Row]("error") != null) + val diagnostic = row.getAs[Row]("error").getString(0) + assert(diagnostic.contains("parse_error") && !diagnostic.contains("private-value")) + assert(row.getAs[String]("raw").contains("\\\\u")) + } + val valid = rows.find(_.getAs[String]("text") == "valid").get + assert(valid.getAs[Row]("error") == null) + assert(valid.getAs[Map[String, String]]("answer") == + Map("value" -> "\uD83D\uDE00", "\uD83D\uDE00" -> "unchanged")) + assert(server.requests.size() == documents.size) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodecSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodecSuite.scala new file mode 100644 index 00000000000..41517c5ff49 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIResponseCodecSuite.scala @@ -0,0 +1,619 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.openai.UsageUtils +import org.apache.spark.ml.linalg.Vector +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.StructType +import spray.json.{JsArray, JsNumber, JsObject, JsString, JsValue, enrichString} + +import java.nio.charset.StandardCharsets +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIResponseCodecSuite extends TestBase { + test("JSON postprocessing honors Spark read syntax without false row errors") { + import spark.implicits._ + val input = Seq("hello").toDF("text") + val cases = Seq( + "{'value':1.0}" -> 1.0, + """{"value":NaN}""" -> Double.NaN, + """{"value":Infinity}""" -> Double.PositiveInfinity, + """{"value":1.0,"value":2.0}""" -> 2.0, + ("""{"value":3.0,"ignored":""" + "[" * 150 + "0" + "]" * 150 + "}") -> 3.0) + Seq("chat_completions", "responses").foreach { api => + cases.foreach { case (text, expected) => + val quoted = JsString(text).compactPrint + val reply = if (api == "chat_completions") { + s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$quoted},"finish_reason":"stop"}]}""" + } else { + s"""{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":$quoted}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> "value DOUBLE")) + val row = stage.transform(input).head() + val actual = row.getAs[Row]("answer").getDouble(0) + assert(actual == expected || (actual.isNaN && expected.isNaN)) + assert(row.getAs[Row]("error") == null, s"$api reported an error for $text") + assert(row.getAs[String]("raw") == reply && server.requests.size() == 1) + } + } + } + } + + test("malformed answer projections retain independently validated usage and response identity") { + Seq("chat_completions", "responses", "embeddings", "systemone", "decisions").foreach { api => + val raw = """{"id":"response-with-invalid-answer","model":"returned", + |"choices":42,"output":42,"data":42,"answers":42,"usage":{"total_tokens":7}}""".stripMargin + val decoded = AIResponseCodec.decode(api, raw) + assert(decoded.value == None.orNull) + assert(decoded.diagnostic.exists(_._1 == "parse_error")) + assert(decoded.usage.getLong(2) == 7L) + assert(decoded.responseId == "response-with-invalid-answer") + } + } + + test("provider failures outrank malformed projections without weakening successful response checks") { + Seq("chat_completions", "responses").foreach { api => + val key = if (api == "chat_completions") "choices" else "output" + Seq("null", "42", "[null]", """[{"unexpected":true}]""").foreach { projection => + val observations = s""""id":"kept-id","usage":{"total_tokens":7},"$key":$projection""" + val errors = Seq(""""error":{"code":"failed","message":"private provider detail"}""") ++ + (if (api == "responses") Seq(""""status":"failed"""") else Seq.empty) + errors.foreach { failure => + val result = AIResponseCodec.decode(api, s"{$failure,$observations}") + assert(result.diagnostic.exists(_._1 == "provider_error")) + assert(result.value == None.orNull && result.usage.getLong(2) == 7L && result.responseId == "kept-id") + assert(!result.diagnostic.get._2.contains("private")) + } + val successful = AIResponseCodec.decode(api, s"{$observations}") + assert(successful.diagnostic.exists(_._1 == "parse_error") && successful.value == None.orNull) + } + val valid = if (api == "chat_completions") { + """{"index":0,"message":{"role":"assistant","content":"partial"}}""" + } else { + """{"type":"message","role":"assistant","content":[{"type":"output_text","text":"partial"}]}""" + } + val mixed = AIResponseCodec.decode(api, s"""{"$key":[$valid,null]}""") + assert(mixed.diagnostic.exists(_._1 == "parse_error") && mixed.value == None.orNull) + } + } + + private val missing: Null = None.orNull + // Answer examples from the public guide, completed with the reference's envelope/usage. + // https://developers.openai.com/api/docs/guides/decisions + // https://developers.openai.com/api/reference/python/resources/decisions/methods/create + private val decisionFixture = + """{"model":"gpt-6-luna","answers":[ + | {"type":"predicate","name":"visible_damage","probability":0.92}, + | {"type":"choice","name":"department","choice":"billing","confidence":0.93, + | "probabilities":[{"value":"billing","probability":0.95},{"value":"technical","probability":0.02}, + | {"value":"shipping","probability":0.01},{"value":"other","probability":0.02}]}, + | {"type":"score","name":"severity","score":1.1,"confidence":0.55, + | "probabilities":[{"value":0,"label":"Cosmetic","probability":0.1}, + | {"value":1,"label":"Workaround available","probability":0.7}, + | {"value":2,"label":"Fully blocked","probability":0.2}]} + |],"usage":{"input_tokens":20,"output_tokens":4,"total_tokens":24, + | "input_tokens_details":{"cached_tokens":3,"cache_write_tokens":2}, + | "output_tokens_details":{"reasoning_tokens":1}}}""".stripMargin + + private def decode(api: String, text: String): AIResponseCodec.Decoded = AIResponseCodec.decode(api, text) + + private def answer(text: String): String = s"""{"model":"model","answers":[$text]}""" + + private def diagnostic(api: String, text: String, code: String): AIResponseCodec.Decoded = { + val result = decode(api, text) + assert(result.diagnostic.map(_._1).contains(code)) + result + } + + private def row(result: AIResponseCodec.Decoded): Row = result.value.asInstanceOf[Row] + + test("Responses permits a bounded echo envelope without relaxing other JSON parsing checks") { + val replies = Seq( + "chat_completions" -> """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""", + "responses" -> """{"output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin, + "embeddings" -> """{"data":[{"index":0,"embedding":[1.0]}]}""", + "systemone" -> """{"model":"model","answers":{"q":true}}""", + "decisions" -> decisionFixture) + replies.foreach { case (api, reply) => + val limit = AIJson.MaxDepth + (if (api == "responses") 3 else 0) + def nested(depth: Int): String = + reply.dropRight(1) + ""","ignored":""" + "[" * (depth - 1) + "0" + "]" * (depth - 1) + "}" + assert(decode(api, nested(limit)).diagnostic.isEmpty, api) + diagnostic(api, nested(limit + 1), "parse_error") + } + Seq( + """{"output":[],"text":{"format":{"x":1,"x":2}}}""", + """{"output":[],"text":{"format":{"x":NaN}}}""", + ("""{"output":[],"text":{"format":{"x":""" + "1" * 1001 + "}}}"), + """{"output":[]} {}""" + ).foreach(diagnostic("responses", _, "parse_error")) + } + + test("Decisions public fixtures round-trip through the stable Spark schema and normalized usage") { + val result = decode("decisions", decisionFixture) + assert(result.diagnostic.isEmpty) + val schema = AIResponseCodec.outputType("decisions").asInstanceOf[StructType] + assert(schema("usage").dataType == UsageUtils.UsageStructType) + val restored = spark.createDataFrame(Collections.singletonList(row(result)), schema).head() + val answers = restored.getSeq[Row](1) + assert(answers.map(_.getString(0)) == Seq("predicate", "choice", "score")) + assert(answers.head.getDouble(2) == 0.92) + assert(answers(1).getStruct(3) == Row("string", "billing", missing)) + assert(answers(1).getSeq[Row](6).head.getStruct(0) == Row("string", "billing", missing)) + assert(answers(2).getSeq[Row](7)(1) == Row(1, "Workaround available", 0.7)) + assert(restored.getStruct(2) == result.usage) + assert(result.usage.getLong(0) == 20L) + assert(result.usage.getMap[String, Long](3) == Map("cached_tokens" -> 3L, "cache_write_tokens" -> 2L)) + assert(result.usage.getMap[String, Long](4) == Map("reasoning_tokens" -> 1L)) + } + + test("Decisions boolean choices are typed and mixed refusals retain ordered siblings") { + val result = diagnostic("decisions", answer( + """{"type":"refusal","name":"blocked"}, + |{"type":"choice","name":"typed","choice":true,"confidence":0.8, + |"probabilities":[{"value":true,"probability":0.8},{"value":"true","probability":0.2}]}, + |{"type":"predicate","name":null,"probability":0.1}""".stripMargin), "refusal") + val answers = row(result).getSeq[Row](1) + assert(answers.map(_.getString(0)) == Seq("refusal", "choice", "predicate")) + assert(answers(1).getStruct(3) == Row("boolean", missing, true)) + assert(answers(1).getSeq[Row](6)(1).getStruct(0) == Row("string", "true", missing)) + assert(answers(2).getDouble(2) == 0.1) + assert(result.usage == missing) + } + + test("Decisions validates required fields and finite bounded numeric probabilities") { + Seq( + """{"type":"predicate"}""", + """{"type":"predicate","probability":"0.3"}""", + """{"type":"predicate","probability":-0.1}""", + """{"type":"predicate","probability":1.0000000000000000001}""", + """{"type":"predicate","probability":1e1000}""", + """{"type":"predicate","probability":0.5,"name":true}""", + """{"type":"future_answer","probability":0.5}""", + """{"type":"choice","choice":true,"confidence":0.5, + |"probabilities":[{"value":"true","probability":0.5},{"value":false,"probability":0.5}]}""".stripMargin, + """{"type":"choice","choice":"a","confidence":2, + |"probabilities":[{"value":"a","probability":0.5},{"value":"b","probability":0.5}]}""".stripMargin, + """{"type":"choice","choice":1,"confidence":0.5,"probabilities":[]}""", + """{"type":"score","score":1,"confidence":0.5, + |"probabilities":[{"value":1,"label":"A","probability":0.5}, + |{"value":0,"label":"B","probability":0.5}]}""".stripMargin, + """{"type":"score","score":2,"confidence":0.5, + |"probabilities":[{"value":0,"label":"A","probability":1}]}""".stripMargin, + """{"type":"score","score":0,"confidence":0.5,"probabilities":[{"value":0,"probability":1}]}""" + ).foreach(a => diagnostic("decisions", answer(a), "parse_error")) + diagnostic("decisions", """{"answers":[{"type":"refusal"}]}""", "parse_error") + diagnostic("decisions", """{"model":"model","answers":[]}""", "parse_error") + diagnostic("decisions", """{"model":"model","answers":{}}""", "parse_error") + } + + test("Decisions score bounds are checked before Double rounding or underflow") { + def response(score: String): String = + NativeRequestFixtures.DecisionResponse.replace("\"score\":1.7", s""""score":$score""") + Seq("-1e-400", "2.0000000000000000001").foreach { score => + diagnostic("decisions", response(score), "parse_error") + } + Seq("0", "2").foreach { score => + val result = decode("decisions", response(score)) + assert(result.diagnostic.isEmpty) + assert(row(result).getSeq[Row](1)(2).getDouble(4) == score.toDouble) + } + } + + test("System One has lossless JSON answers and does not invent a provider answer schema") { + // The public guide establishes the envelope, not a complete answer/usage schema: + // https://learn.microsoft.com/en-us/azure/foundry/foundry-models/how-to/use-foundry-models-microsoft-decision + // These heterogeneous answer values exercise SynapseML's representation, not guessed provider projections. + val payload = + """{"model":"microsoft-decision-1","answers":{"n":0.75,"b":true,"s":"true","nothing":null, + |"a":[1,false],"o":{"type":"future","unknown":{"confidence":"not a known field"}}}, + |"usage":{"input_tokens":12,"output_tokens":3,"total_tokens":15, + |"input_tokens_details":{"cache_write_tokens":2},"future_usage":{"unit":"new"}}}""".stripMargin + val result = decode("systemone", payload) + assert(result.diagnostic.isEmpty) + val schema = AIResponseCodec.outputType("systemone").asInstanceOf[StructType] + val restored = spark.createDataFrame(Collections.singletonList(row(result)), schema).head() + val expected = payload.parseJson.asJsObject.fields("answers").asJsObject.fields + val answers = restored.getMap[String, String](1).map { case (k, v) => k -> v.parseJson } + assert(answers == expected) + assert(restored.getString(0) == "microsoft-decision-1") + assert(restored.getString(2).parseJson == payload.parseJson.asJsObject.fields("usage")) + assert(result.usage.getLong(0) == 12L) + assert(result.usage.getMap[String, Long](3) == Map("cache_write_tokens" -> 2L)) + } + + test("System One validates only the documented response envelope and known usage") { + Seq( + """{"model":false,"answers":{"q":1}}""", + """{"model":"m","answers":[]}""", + """{"model":"m","answers":{"q":1},"usage":12}""", + """{"model":"m","answers":{"q":1},"usage":{"input_tokens":1.5}}""" + ).foreach(diagnostic("systemone", _, "parse_error")) + val unobserved = decode("systemone", """{"model":"m","answers":{"q":true}}""") + assert(unobserved.diagnostic.isEmpty && unobserved.usage == missing && row(unobserved).isNullAt(2)) + } + + private def textPart(text: String): JsValue = JsObject("type" -> JsString("output_text"), "text" -> JsString(text)) + + private def message(parts: JsValue*): JsValue = + JsObject("type" -> JsString("message"), "role" -> JsString("assistant"), "content" -> JsArray(parts.toVector)) + + private def response(output: JsValue*): JsObject = JsObject( + "id" -> JsString("resp-example"), "status" -> JsString("completed"), "output" -> JsArray(output.toVector)) + + test("Responses concatenates every assistant text part in order without tool arguments") { + val raw = response(message(textPart("a"), textPart("b")), + JsObject("type" -> JsString("function_call"), "arguments" -> JsString("not prose")), + message(textPart("c"))).compactPrint + val result = decode("responses", raw) + assert(result.value == "abc" && result.responseId == "resp-example" && result.diagnostic.isEmpty) + assert(result.usage == missing) + } + + test("Responses partial text and usage survive provider failure, refusal and incomplete diagnostics") { + val base = response(message(textPart("partial"))) + val usage = decisionFixture.parseJson.asJsObject.fields("usage") + val cases = Seq( + ("incomplete", JsObject("status" -> JsString("incomplete"))), + ("provider_error", JsObject("status" -> JsString("failed"), + "error" -> JsObject("code" -> JsString("server_error"), "message" -> JsString("Sensitive provider text")))), + ("refusal", JsObject("output" -> JsArray(message(textPart("partial")), + message(JsObject("type" -> JsString("refusal"), "refusal" -> JsString("Sensitive refusal text")))))) + ) + cases.foreach { case (code, changes) => + val raw = JsObject(base.fields ++ changes.fields + ("usage" -> usage)).compactPrint + val result = diagnostic("responses", raw, code) + assert(result.value == "partial" && result.usage.getLong(0) == 20L && result.responseId == "resp-example") + assert(!result.diagnostic.get._2.contains("Sensitive")) + } + } + + test("Responses no-text output is not a refusal and malformed known text is a parse error") { + diagnostic("responses", response( + JsObject("type" -> JsString("function_call"), "arguments" -> JsString("{}"))).compactPrint, "no_text") + val malformed = message(JsObject("type" -> JsString("output_text"), "text" -> JsArray())) + diagnostic("responses", response(malformed).compactPrint, "parse_error") + diagnostic("responses", """{"output":[{"type":"message","role":"assistant"}]}""", "parse_error") + val emptyText = decode("responses", response(message(textPart(""))).compactPrint) + assert(emptyText.diagnostic.isEmpty && emptyText.value == "") + } + + private def chat(content: String, finish: String = "stop", refusal: String = "null"): String = + s"""{"id":"chat-example","choices":[{"index":0,"finish_reason":"$finish", + |"message":{"role":"assistant","content":$content,"refusal":$refusal}}]}""".stripMargin + + test("Chat empty refusal preserves successful text without masking other diagnostics") { + Seq("null", "\"\"").foreach { refusal => + val decoded = decode("chat_completions", chat("\"ok\"", refusal = refusal)) + assert(decoded.value == "ok" && decoded.diagnostic.isEmpty && decoded.responseId == "chat-example") + diagnostic("chat_completions", chat("null", refusal = refusal), "no_text") + diagnostic("chat_completions", chat("\"partial\"", "length", refusal), "incomplete") + diagnostic("chat_completions", chat("null", "content_filter", refusal), "refusal") + } + Seq("\"declined\"", "\" \"").foreach { refusal => + diagnostic("chat_completions", chat("\"partial\"", refusal = refusal), "refusal") + } + } + + Seq("chat_completions", "responses", "embeddings", "systemone", "decisions").foreach { api => + val fixture = api match { + case "chat_completions" => chat("\"ok\"").parseJson.asJsObject + case "responses" => response(message(textPart("ok"))) + case "embeddings" => """{"data":[{"index":0,"embedding":[0.5]}]}""".parseJson.asJsObject + case "systemone" => """{"model":"model","answers":{"q":true}}""".parseJson.asJsObject + case "decisions" => answer("""{"type":"predicate","name":"q","probability":0.5}""").parseJson.asJsObject + } + test(s"$api public HTTP decoding rejects malformed UTF-8 without changing valid Unicode") { + val cases = Seq("malformed-utf8", "http-error", "unicode", "replacement") + val payloads = cases.map { name => + val text = name match { + case "unicode" => "caf\u00e9 \u96ea \uD83D\uDE00" + case "replacement" => "\uFFFD" + case _ => "private-value~" + } + val payload = JsObject(fixture.fields ++ Map("model" -> JsString(text), "id" -> JsString("kept-id"), + "usage" -> """{"total_tokens":7}""".parseJson, "future" -> JsString(text))).compactPrint + .replace("\"ok\"", JsString(text).compactPrint) + name -> payload + }.toMap + Using.resource(new AIRecordingServer((_, body) => { + val name = cases.find(body.contains).get + (if (name == "http-error") 400 else 200) -> payloads(name) + }, encodeResponse = text => text.getBytes(StandardCharsets.UTF_8) + .map(byte => if (byte == '~'.toByte) 0xff.toByte else byte))) { server => + val endpoint = if (api == "systemone") server.url else server.url + "/v1" + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("id").setRawResponseCol("raw") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val data = spark.createDataFrame(cases.map(Row(_)).asJava, new StructType().add("text", "string")) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + assert(rows.length == cases.length && server.requests.size() == cases.length) + rows.foreach { row => + val name = row.getAs[String]("text") + if (name == "malformed-utf8" || name == "http-error") { + assert(Seq("answer", "raw", "usage", "id").forall(field => row.isNullAt(row.fieldIndex(field)))) + val error = row.getAs[Row]("error") + val code = if (name == "http-error") "http_error" else "parse_error" + assert(error.getString(0).parseJson.asJsObject.fields("code") == JsString(code)) + assert(!error.getString(0).contains("private-value") && !error.getString(0).contains("kept-id")) + if (name == "http-error") assert(error.getAs[Row]("status").getAs[Int]("statusCode") == 400) + else assert(error.getAs[Row]("status") == null) + } else { + assert(row.getAs[Row]("error") == null) + assert(row.getAs[String]("raw") == payloads(name)) + assert(row.getAs[Any]("answer") == AIResponseCodec.decode(api, payloads(name)).value) + assert(row.getAs[String]("id") == "kept-id" && row.getAs[Row]("usage").getLong(2) == 7L) + } + } + } + } + + test(s"$api JSON surrogate validation preserves escaped Unicode without repairing text") { + val escape = "\\" + "u" + val cases = Seq("d800", "dc00", "d800x", s"dc00${escape}d800", s"d83d${escape}de00") + .map(code => "\"" + escape + code + "\"") + Using.resource(new AIRecordingServer((_, body) => { + val index = cases.indices.find(i => body.contains(s"case-$i")).get + val payload = fixture.compactPrint.dropRight(1) + s""","future":${cases(index)}}""" + 200 -> payload.replace("\"ok\"", cases(index)) + })) { server => + val stage = new AIPrompt().setApiType(api) + .setUrl(if (api == "systemone") server.url else server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("id").setRawResponseCol("raw") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val data = spark.createDataFrame(cases.indices.map(i => Row(s"case-$i")).asJava, + new StructType().add("text", "string")) + val rows = stage.transform(data).collect().sortBy(_.getAs[String]("text")) + rows.dropRight(1).foreach { row => + assert(Seq("answer", "usage", "id").forall(field => row.isNullAt(row.fieldIndex(field)))) + assert(row.getAs[Row]("error").getString(0).parseJson.asJsObject.fields("code") == JsString("parse_error")) + assert(row.getAs[String]("raw").contains("\\ud")) + } + val valid = rows.last + assert(valid.getAs[Row]("error") == null) + assert(valid.getAs[Any]("answer") == AIResponseCodec.decode(api, valid.getAs[String]("raw")).value) + if (Set("chat_completions", "responses")(api)) assert(valid.getAs[String]("answer") == "\uD83D\uDE00") + assert(server.requests.size() == cases.size) + } + } + + test(s"$api optional response diagnostics preserve public answers and independent metadata") { + val fields = fixture.fields ++ Map("id" -> JsString("kept-id"), + "usage" -> """{"total_tokens":7}""".parseJson) + val malformedErrors = Seq(JsArray(), JsObject("code" -> JsNumber(42)), + JsObject("message" -> JsNumber(42))).map(error => fields.updated("error", error) -> "parse_error") + val protocolCases = api match { + case "chat_completions" => + def choice(finish: JsValue, refusal: JsValue): Map[String, JsValue] = + fields.updated("choices", JsArray(JsObject("index" -> JsNumber(0), "finish_reason" -> finish, + "message" -> JsObject("role" -> JsString("assistant"), "content" -> JsString("ok"), + "refusal" -> refusal)))) + Seq(choice(JsNumber(42), JsString("")) -> "parse_error", + choice(JsString("stop"), JsNumber(42)) -> "parse_error", + choice(JsNumber(42), spray.json.JsNull) -> "parse_error", + choice(JsString("length"), JsNumber(42)) -> "incomplete", + choice(JsNumber(42), JsString("declined")) -> "refusal", + choice(JsNumber(42), spray.json.JsNull).updated("error", + JsObject("code" -> JsString("failed"))) -> "provider_error") + case "responses" => + Seq(fields.updated("status", JsNumber(42)) -> "parse_error", + fields.updated("incomplete_details", JsNumber(42)) -> "incomplete", + fields.updated("incomplete_details", JsObject("reason" -> JsNumber(42))) -> "incomplete", + fields.updated("output", JsArray(JsObject(message(textPart("ok")).asJsObject.fields + .updated("status", JsNumber(42))))) -> "parse_error", + fields.updated("status", JsString("failed")).updated("error", JsArray()) -> "provider_error", + fields.updated("status", JsNumber(42)).updated("output", JsArray(message(textPart("ok"), + JsObject("type" -> JsString("refusal"), "refusal" -> JsString("declined"))))) -> "refusal") + case _ => Seq.empty + } + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + (malformedErrors ++ protocolCases).foreach { case (payloadFields, code) => + val payload = JsObject(payloadFields).compactPrint + Using.resource(new AIRecordingServer((_, _) => 200 -> payload)) { server => + val endpoint = if (api == "systemone") server.url else server.url + "/v1" + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("id").setRawResponseCol("raw") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val observed = output.head() + withClue(payload) { + api match { + case "embeddings" => assert(observed.getAs[Vector]("answer").toArray.sameElements(Array(0.5))) + case "systemone" => assert(observed.getAs[Row]("answer").getMap[String, String](1)("q") == "true") + case "decisions" => assert(observed.getAs[Row]("answer").getSeq[Row](1).head.getDouble(2) == 0.5) + case _ => assert(observed.getAs[String]("answer") == "ok") + } + assert(observed.getAs[String]("id") == "kept-id" && observed.getAs[Row]("usage").getLong(2) == 7L) + assert(observed.getAs[Row]("error").getString(0).contains(code)) + assert(observed.getAs[String]("raw") == payload && server.requests.size() == 1) + } + } + } + } + + test(s"$api nullable usage details preserve available counters without a public row error") { + val chatNames = api == "chat_completions" || api == "embeddings" + val input = if (chatNames) "prompt_tokens" else "input_tokens" + val output = if (chatNames) "completion_tokens" else "output_tokens" + val usage = JsObject(input -> JsNumber(7), output -> JsNumber(3), "total_tokens" -> JsNumber(10), + s"${input}_details" -> """{"cached_tokens":2,"audio_tokens":null}""".parseJson, + s"${output}_details" -> """{"reasoning_tokens":1,"audio_tokens":null}""".parseJson) + val payload = JsObject(fixture.fields ++ Map("usage" -> usage, "id" -> JsString("kept-id"))).compactPrint + Using.resource(new AIRecordingServer((_, _) => 200 -> payload)) { server => + val endpoint = if (api == "systemone") server.url else server.url + "/v1" + val stage = new AIPrompt().setApiType(api).setUrl(endpoint).setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("id").setRawResponseCol("raw") + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val observed = result.head() + val expected = Row(7L, 3L, 10L, Map("cached_tokens" -> 2L), Map("reasoning_tokens" -> 1L)) + assert(!observed.isNullAt(observed.fieldIndex("answer")) && observed.getAs[Row]("error") == null) + assert(observed.getAs[Row]("usage") == expected && observed.getAs[String]("id") == "kept-id") + if (api == "decisions") assert(observed.getAs[Row]("answer").getStruct(2) == expected) + if (api == "systemone") assert(AIJson.parse(observed.getAs[Row]("answer").getString(2)) == usage) + assert(observed.getAs[String]("raw") == payload && server.requests.size() == 1) + } + } + } + + Seq( + "chat_completions" -> """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""", + "responses" -> """{"output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"ok"}]}]}""".stripMargin, + "embeddings" -> """{"data":[{"index":0,"embedding":[0.5]}]}""" + ).foreach { case (api, fixture) => + test(s"$api malformed optional model metadata retains public answers and independent observations") { + val payload = JsObject(fixture.parseJson.asJsObject.fields ++ Map( + "model" -> JsNumber(42), "id" -> JsString("kept-id"), + "usage" -> """{"total_tokens":7}""".parseJson)).compactPrint + Using.resource(new AIRecordingServer((_, _) => 200 -> payload)) { server => + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setResponseIdCol("id") + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val observed = output.head() + if (api == "embeddings") assert(observed.getAs[Vector]("answer").toArray.sameElements(Array(0.5))) + else assert(observed.getAs[String]("answer") == "ok") + assert(observed.getAs[String]("id") == "kept-id" && observed.getAs[Row]("usage").getLong(2) == 7L) + assert(observed.getAs[Row]("error").getString(0).contains("parse_error")) + assert(server.requests.size() == 1) + } + } + } + + test("malformed optional model metadata does not replace partial-response provider diagnostics") { + val base = response(message(textPart("partial"))).fields.updated("model", JsNumber(42)) + Seq("incomplete" -> "incomplete", "failed" -> "provider_error").foreach { case (status, code) => + val decoded = diagnostic("responses", JsObject(base.updated("status", JsString(status))).compactPrint, code) + assert(decoded.value == "partial" && decoded.responseId == "resp-example") + } + Seq("systemone" -> """{"model":42,"answers":{"q":true}}""", + "decisions" -> """{"model":42,"answers":[{"type":"predicate","probability":0.5}]}""") + .foreach { case (api, payload) => assert(diagnostic(api, payload, "parse_error").value == missing) } + } + + test("Chat deliberately projects the first choice and preserves refusal and truncation") { + val base = chat("\"first\"").parseJson.asJsObject + val choices = base.fields("choices").asInstanceOf[JsArray].elements + val second = chat("\"second\"").parseJson.asJsObject.fields("choices").asInstanceOf[JsArray].elements + val raw = JsObject(base.fields.updated("choices", JsArray(choices ++ second))).compactPrint + assert(decode("chat_completions", raw).value == "first") + assert(diagnostic("chat_completions", chat("\"partial\"", "length"), "incomplete").value == "partial") + diagnostic("chat_completions", chat("null", refusal = "\"declined\""), "refusal") + diagnostic("chat_completions", chat("null", "content_filter"), "refusal") + diagnostic("chat_completions", chat("null", "tool_calls"), "no_text") + diagnostic("chat_completions", chat("42"), "parse_error") + diagnostic("chat_completions", """{"choices":[{"message":{"role":"assistant","content":"Missing index"}}]}""", + "parse_error") + } + + test("Embeddings accepts exactly one indexed float vector and rejects base64 or malformed dimensions") { + val raw = """{"object":"list","model":"embedding","data":[ + |{"object":"embedding","index":0,"embedding":[0.25,-0.5,1]}], + |"usage":{"prompt_tokens":3,"total_tokens":3}}""".stripMargin + val result = decode("embeddings", raw) + assert(result.diagnostic.isEmpty && result.value.asInstanceOf[Vector].toArray.sameElements(Array(0.25, -0.5, 1))) + assert(result.usage.getLong(0) == 3L && result.usage.isNullAt(1)) + Seq( + """{"data":[{"index":0,"embedding":"AACAPw=="}]}""", + """{"data":[{"index":1,"embedding":[1]}]}""", + """{"data":[{"index":0.1,"embedding":[1]}]}""", + """{"data":[{"embedding":[1]}]}""", + """{"data":[{"index":0,"embedding":[1e1000]}]}""", + """{"data":[{"index":0,"embedding":[]}]}""", + """{"data":[{"index":0,"embedding":[true]}]}""", + """{"object":42,"data":[{"index":0,"embedding":[1]}]}""", + """{"object":"wrong","data":[{"index":0,"embedding":[1]}]}""", + """{"data":[{"object":42,"index":0,"embedding":[1]}]}""", + """{"data":[{"object":"wrong","index":0,"embedding":[1]}]}""", + """{"data":[{"index":0,"embedding":[1]},{"index":1,"embedding":[2]}]}""", + """{"data":[]}""" + ).foreach(diagnostic("embeddings", _, "parse_error")) + } + + test("Usage counts reject negative fractional overflow and noninteger known details") { + val base = response(message(textPart("ok"))) + Seq( + """{"input_tokens":-1}""", + """{"input_tokens":0.5}""", + """{"input_tokens":null}""", + """{"input_tokens":9223372036854775808}""", + """{"input_tokens":"1"}""", + """{"input_tokens_details":{"cache_write_tokens":true}}""", + """{"output_tokens_details":{"reasoning_tokens":0.5}}""", + """{"input_tokens_details":[]}""" + ).foreach { usage => + val decoded = diagnostic("responses", + JsObject(base.fields.updated("usage", usage.parseJson)).compactPrint, "parse_error") + assert(decoded.value == "ok" && decoded.responseId == "resp-example" && decoded.usage == missing) + } + val chatUsage = """{"prompt_tokens":9,"completion_tokens":3,"total_tokens":12, + |"prompt_tokens_details":{"cached_tokens":2},"completion_tokens_details":{"reasoning_tokens":1}}""".stripMargin + val chatBase = chat("\"ok\"").parseJson.asJsObject + val result = decode("chat_completions", JsObject(chatBase.fields + ("usage" -> chatUsage.parseJson)).compactPrint) + assert(result.usage == Row(9L, 3L, 12L, Map("cached_tokens" -> 2L), Map("reasoning_tokens" -> 1L))) + } + + test("malformed optional metadata preserves valid answers and independent observations on every API") { + val fixtures = Seq( + "chat_completions" -> chat("\"ok\""), + "responses" -> response(message(textPart("ok"))).compactPrint, + "embeddings" -> """{"data":[{"index":0,"embedding":[0.5]}]}""", + "systemone" -> """{"model":"model","answers":{"q":true}}""", + "decisions" -> answer("""{"type":"predicate","name":"q","probability":0.5}""")) + fixtures.foreach { case (api, fixture) => + val base = fixture.parseJson.asJsObject.fields + val badId = diagnostic(api, JsObject(base ++ Map( + "id" -> JsArray(), "usage" -> """{"total_tokens":7}""".parseJson)).compactPrint, "parse_error") + assert(badId.value != missing && badId.usage.getLong(2) == 7L && badId.responseId == missing) + val badUsage = diagnostic(api, JsObject(base ++ Map( + "id" -> JsString("kept-id"), "usage" -> """{"total_tokens":-1}""".parseJson)).compactPrint, "parse_error") + assert(badUsage.value != missing && badUsage.responseId == "kept-id" && badUsage.usage == missing) + if (api == "systemone") assert(row(badUsage).getString(2) == """{"total_tokens":-1}""") + if (api == "decisions") assert(row(badUsage).isNullAt(2)) + } + val failed = JsObject(response(message(textPart("partial"))).fields ++ Map( + "status" -> JsString("failed"), "id" -> JsArray())) + assert(diagnostic("responses", failed.compactPrint, "provider_error").value == "partial") + } + + test("Provider errors and invalid JSON have distinct safe diagnostics on every API") { + Seq("chat_completions", "responses", "embeddings", "systemone", "decisions").foreach { api => + val failed = diagnostic(api, """{"error":{"code":"rejected","message":"Sensitive provider text"}}""", + "provider_error") + assert(failed.value == missing && failed.usage == missing) + Seq("not-json-private-value", "[]", """{"model":"one","model":"two"}""").foreach { raw => + val bad = diagnostic(api, raw, "parse_error") + assert(!bad.diagnostic.get._2.contains(raw)) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneLiveSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneLiveSuite.scala new file mode 100644 index 00000000000..e2197fb542a --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneLiveSuite.scala @@ -0,0 +1,44 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{StringType, StructType} + +import scala.collection.JavaConverters._ + +/** Opt-in inference only: this suite never provisions, changes, or deletes Azure resources. */ +class AISystemOneLiveSuite extends TestBase { + test("existing Decision-1 deployment supports all three question types through direct and facade APIs") { + val endpoint = AILiveTestEnvironment.get("SYNAPSEML_SYSTEMONE_ENDPOINT") + val deployment = AILiveTestEnvironment.get("SYNAPSEML_SYSTEMONE_DEPLOYMENT") + val token = AILiveTestEnvironment.get("SYNAPSEML_SYSTEMONE_TOKEN") + val key = AILiveTestEnvironment.get("SYNAPSEML_SYSTEMONE_KEY") + assume(endpoint.isDefined && deployment.isDefined && (token.isDefined || key.isDefined), + "Live System One endpoint, deployment and credentials are not configured") + val input = spark.createDataFrame(Seq( + Row("Checkout is down again for every customer. This is my third request.")).asJava, + new StructType().add("ticket", StringType)) + Seq(new AISystemOne(), new AIPrompt().setApiType("systemone")).foreach { stage => + stage.setUrl(endpoint.get).setDeploymentName(deployment.get).setStateCol("ticket") + .setQuestions(NativeRequestFixtures.SystemQuestions).setOutputCol("decision") + .setErrorCol("error").setUsageCol("usage") + token match { + case Some(value) => stage.setAADToken(value) + case None => key.foreach(stage.setSubscriptionKey) + } + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null, "Live System One request failed; inspect the protected error channel") + val decision = row.getAs[Row]("decision") + val answers = decision.getAs[Map[String, String]]("answers") + assert(answers.keySet == Set("repeat_contact", "team", "severity")) + answers.values.foreach(AIJson.parse) + assert(decision.getAs[String]("model").nonEmpty) + assert(row.getAs[Row]("usage") != null) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneSuite.scala new file mode 100644 index 00000000000..a1d1baa0f7d --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AISystemOneSuite.scala @@ -0,0 +1,186 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.fuzzing.TestObject +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.ml.util.MLReadable +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{StringType, StructType} +import spray.json.DefaultJsonProtocol._ +import spray.json._ + +import scala.collection.JavaConverters._ +import scala.util.Using + +class AISystemOneSuite extends AIPromptFuzzing[AISystemOne] { + override def reader: MLReadable[_] = AISystemOne + + override def testObjects(): Seq[TestObject[AISystemOne]] = Seq(new TestObject( + new AISystemOne("systemone-fuzzing").setUrl("https://example.invalid").setDeploymentName("synthetic-model") + .setSubscriptionKey("synthetic-key").setStateCol("state").setQuestions(NativeRequestFixtures.SystemQuestions) + .setOutputCol("decision").setErrorCol("error"), + spark.createDataFrame(Seq(Row(Option.empty[String].orNull)).asJava, inputSchema))) + + private val inputSchema = new StructType().add("state", StringType) + private val inputText = "Checkout is down again. This is my third request." + private def input = spark.createDataFrame(Seq(Row(inputText)).asJava, inputSchema) + + test("direct System One and facade use the native route and retain every answer JSON value") { + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + Seq(new AISystemOne(), new AIPrompt().setApiType("systemone")).foreach { stage => + stage.setUrl(server.url).setDeploymentName("Decision-1").setStateCol("state") + .setQuestions(NativeRequestFixtures.SystemQuestions).setOutputCol("decision").setErrorCol("error") + .setUsageCol("usage").setRawResponseCol("raw") + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null) + val decision = row.getAs[Row]("decision") + assert(decision.getAs[String]("model") == "microsoft-decision-1") + val answers = decision.getAs[Map[String, String]]("answers") + assert(answers("repeat_contact") == "true") + assert(answers("severity") == "1.7") + assert(AIJson.parse(answers("team")).asJsObject.fields("additional") == + JsArray(JsNumber(1), JsBoolean(false), JsNull)) + assert(row.getAs[String]("raw").contains("future")) + } + val expected = JsObject("model" -> JsString("Decision-1"), "state" -> JsString(inputText), + "questions" -> NativeRequestFixtures.SystemQuestions.parseJson) + assert(server.requests.size() == 2 && server.requests.asScala.forall { case (path, body) => + path == "/providers/microsoft/v1/systemone" && body.parseJson == expected + }) + } + } + + test("JSON state interpretation and question columns remain native JSON rather than escaped text") { + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + val data = spark.createDataFrame(Seq(Row("""{"message":"outage","count":3}""", + NativeRequestFixtures.SystemQuestions)).asJava, inputSchema.add("questions", StringType)) + val stage = new AISystemOne().setUrl(server.url).setDeploymentName("Decision-1") + .setStateCol("state").setStateMode("json").setQuestionsCol("questions") + stage.transform(data).collect() + assert(AIJson.objectString(server.requests.peek()._2).fields("state").isInstanceOf[JsObject]) + assert(AIJson.objectString(server.requests.peek()._2).fields("questions").isInstanceOf[JsObject]) + } + } + + test("direct Decisions and facade preserve question order and typed native answers") { + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + Seq(new OpenAIDecisions(), new AIPrompt().setApiType("decisions")).foreach { stage => + stage.setUrl(server.url + "/v1").setModel("gpt-6-luna").setInputCol("state") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setSafetyIdentifier("synthetic-test") + .setOutputCol("decision").setErrorCol("error") + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(row.getAs[Row]("error") == null) + val decision = row.getAs[Row]("decision") + val answers = decision.getSeq[Row](decision.fieldIndex("answers")) + assert(answers.map(_.getAs[String]("name")) == Seq("repeat_contact", "team", "severity")) + assert(answers(1).getAs[Row]("choice").getAs[String]("stringValue") == "engineering") + assert(answers(2).getSeq[Row](answers(2).fieldIndex("scoreProbabilities")).last.getAs[Int]("value") == 2) + } + val expected = JsObject("model" -> JsString("gpt-6-luna"), "input" -> JsString(inputText), + "questions" -> NativeRequestFixtures.DecisionQuestions.parseJson, + "safety_identifier" -> JsString("synthetic-test")) + assert(server.requests.size() == 2 && server.requests.asScala.forall { case (path, body) => + path == "/v1/decisions" && body.parseJson == expected + }) + } + } + + test("direct Decisions and facade preserve exact score bounds as terminal parse errors") { + val scores = Seq("0", "2", "-1e-400", "2.0000000000000000001") + val data = spark.createDataFrame(scores.map(Row(_)).asJava, inputSchema) + def response(score: String): String = + NativeRequestFixtures.DecisionResponse.replace("\"score\":1.7", s""""score":$score""") + Using.resource(new AIRecordingServer((_, body) => + 200 -> response(AIJson.objectString(body).fields("input").convertTo[String]))) { server => + Seq(new OpenAIDecisions(), new AIPrompt().setApiType("decisions")).foreach { stage => + stage.setUrl(server.url + "/v1").setModel("model").setInputCol("state") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setOutputCol("answer").setErrorCol("error") + .setRawResponseCol("raw").setUsageCol("usage") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + assert(rows.length == scores.size) + rows.foreach { row => + val score = row.getAs[String]("state") + assert(row.getAs[String]("raw") == response(score)) + assert(row.getAs[Row]("usage").getAs[Long]("total_tokens") == 20L) + if (score == "0" || score == "2") { + assert(row.getAs[Row]("error") == null) + assert(row.getAs[Row]("answer").getAs[Seq[Row]]("answers")(2).getAs[Double]("score") == score.toDouble) + } else { + assert(row.getAs[Row]("answer") == null) + assert(row.getAs[Row]("error").getString(0).parseJson.asJsObject.fields("code") == JsString("parse_error")) + } + } + } + assert(server.requests.size() == scores.size * 2) + } + } + + Seq( + ("systemone", () => new AISystemOne(), NativeRequestFixtures.SystemQuestions, + NativeRequestFixtures.SystemResponse, "/providers/microsoft/v1/systemone", "state"), + ("decisions", () => new OpenAIDecisions(), NativeRequestFixtures.DecisionQuestions, + NativeRequestFixtures.DecisionResponse, "/v1/decisions", "input") + ).foreach { case (api, create, questions, reply, route, inputKey) => + val expected = JsObject("model" -> JsString("synthetic-model"), inputKey -> JsString(inputText), + "questions" -> questions.parseJson) + val serviceError = """{"error":{"code":"synthetic_error","message":"Synthetic native service failure"}}""" + + test(s"$api protocol mock checks native requests authentication and transient retries") { + Seq(create(), new AIPrompt().setApiType(api)).foreach { stage => + Using.resource(new AIRecordingServer((count, body) => { + if (body.parseJson != expected) 400 -> serviceError + else if (count == 1) 503 -> serviceError + else 200 -> reply + })) { server => + stage.setUrl(server.url).setModel("synthetic-model").setInputCol("state").setQuestions(questions) + .setSubscriptionKey("synthetic-key").setOutputCol("answer").setErrorCol("error").setConcurrency(1) + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + assert(result.head().getAs[Row]("error") == null) + assert(server.requests.size() == 2 && server.requests.asScala.forall { case (path, body) => + path == route && body.parseJson == expected + }) + val auth = if (api == "systemone") Some("synthetic-key") -> None + else None -> Some("Bearer synthetic-key") + assert(server.headers.asScala.forall { headers => + (headers.get("api-key"), headers.get("authorization")) == auth + }) + } + } + } + + test(s"$api protocol mock preserves terminal HTTP and malformed response diagnostics without retries") { + Seq((400, serviceError, "http_error"), (200, "{}", "parse_error")).foreach { + case (status, response, code) => + Seq(create(), new AIPrompt().setApiType(api)).foreach { stage => + Using.resource(new AIRecordingServer((_, _) => status -> response)) { server => + stage.setUrl(server.url).setModel("synthetic-model").setInputCol("state").setQuestions(questions) + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + val row = stage.transform(input).head() + assert(row.getAs[Row]("error").getString(0).contains(code) && row.isNullAt(row.fieldIndex("answer"))) + assert(row.getAs[String]("raw") == response && server.requests.size() == 1) + assert(server.requests.peek()._1 == route && server.requests.peek()._2.parseJson == expected) + } + } + } + } + } + + test("native malformed responses are parse errors and do not discard raw JSON") { + Using.resource(new AIRecordingServer((_, _) => 200 -> """{"model":"Decision-1","answers":[]}""")) { server => + val stage = new AISystemOne().setUrl(server.url).setDeploymentName("Decision-1").setStateCol("state") + .setQuestions(NativeRequestFixtures.SystemQuestions).setErrorCol("error").setRawResponseCol("raw") + val row = stage.transform(input).head() + assert(row.getAs[Row]("error").getString(0).contains("parse_error")) + assert(row.getAs[String]("raw").contains("answers")) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodecSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodecSuite.scala new file mode 100644 index 00000000000..3f3b6ae9c75 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodecSuite.scala @@ -0,0 +1,232 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import org.apache.spark.ml.linalg.{DenseVector, Matrices, SQLDataTypes, SparseVector, Vector, Vectors} +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import spray.json._ + +import java.util.Arrays +import scala.util.Using + +class AIValueCodecSuite extends AIPromptTestBase { + test("schema reconciliation widens non-nullable expressions without changing values") { + val metadata = new MetadataBuilder().putString("source", "schema-reconciliation").build() + val input = spark.range(1).select(lit("hello").as("text", metadata), + struct(lit(1).as("value")).as("record"), array(lit(2)).as("items")) + assert(input.schema.fields.forall(!_.nullable)) + val expected = StructType(input.schema.fields.map(_.copy(nullable = true))) + val output = AIValueCodec.reconcileSchema(input, expected) + assert(output.schema == expected) + assert(output.collect().toSeq == input.collect().toSeq) + } + + test("negative decimal scales remain valid at every nesting position") { + val setting = "spark.sql.legacy.allowNegativeScaleOfDecimal" + val previous = spark.conf.getOption(setting) + try { + spark.conf.set(setting, "true") + val amount = lit(120).cast(DecimalType(8, -1)) + val values = Seq(amount, struct(amount.as("amount")), array(amount), + map(lit("key"), amount), map(amount, lit("value"))) + withServer { server => + values.zipWithIndex.foreach { case (retained, index) => + val data = spark.range(1).select(lit("hello").as("text"), retained.as("retained")) + val stage = prompt(server) + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val row = output.head() + assert(row.getAs[String]("answer") == "answer") + assert(row.getAs[Any]("retained") == data.head().getAs[Any]("retained")) + assert(server.requests.get() == index + 1) + } + } + } finally previous.fold(spark.conf.unset(setting))(spark.conf.set(setting, _)) + } + + Seq("retained", "input", "extraBody", "responseFormat").foreach { binding => + test(s"empty input names remain valid for $binding bindings") { + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + val value = binding match { + case "extraBody" => """{"metadata":{"source":"empty-name"}}""" + case "responseFormat" => """{"type":"json_object"}""" + case _ => "empty-name" + } + val input = spark.range(1).select(lit("hello").as("text"), lit(value).as("")) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + binding match { + case "input" => stage.setInputCol("") + case "extraBody" => stage.setExtraBodyCol("") + case "responseFormat" => stage.setResponseFormatCol("") + case _ => () + } + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val row = output.head() + assert(row.getAs[String]("") == value) + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + val body = AIJson.objectString(server.requests.peek()._2) + val message = body.fields("messages").asInstanceOf[JsArray].elements.last.asJsObject + assert(message.fields("content") == JsString(if (binding == "input") value else "hello")) + if (binding == "extraBody") assert(body.fields("metadata") == JsObject("source" -> JsString("empty-name"))) + if (binding == "responseFormat") assert(body.fields("response_format") == AIJson.parse(value)) + assert(server.requests.size() == 1) + } + } + } + + test("input names reject nulls and ambiguity and output roles require nonempty names") { + val stage = new AIPrompt().setUrl("http://localhost/v1").setModel("model").setInputCol("text") + val valid = new StructType().add("text", StringType) + val invalid = Seq(valid.add(Option.empty[String].orNull, StringType), valid.add("", StringType).add("", StringType)) + invalid.foreach(schema => intercept[IllegalArgumentException](stage.transformSchema(schema))) + intercept[IllegalArgumentException](stage.setOutputCol("").transformSchema(valid)) + } + + test("matrix map keys reject before Row conversion instead of losing SQL-distinct entries") { + val matrices = spark.createDataFrame(Arrays.asList( + Row(Matrices.dense(1, 1, Array(1.0)), Matrices.sparse(1, 1, Array(0, 1), Array(0), Array(1.0)))), + new StructType().add("dense", SQLDataTypes.MatrixType).add("sparse", SQLDataTypes.MatrixType)) + val weights = map_from_arrays(array(col("dense"), col("sparse")), array(lit("dense"), lit("sparse"))) + val input = matrices.select(lit("hello").as("text"), weights.as("weights"), + array(struct(weights.as("weights"))).as("nested")) + assert(input.select(size(col("weights")), size(col("nested").getItem(0).getField("weights"))).head() == + Row(2, 2)) + withServer { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + Seq(input.drop("nested"), input.drop("weights")).foreach { data => + Seq[() => Any](() => stage.transformSchema(data.schema), () => stage.transform(data)).foreach { call => + assert(intercept[IllegalArgumentException](call()).getMessage.contains("UDT map keys")) + } + } + assert(server.requests.get() == 0) + val ordinary = matrices.select(lit("hello").as("text"), col("dense"), col("sparse")) + val output = stage.transform(ordinary) + assert(output.schema == stage.transformSchema(ordinary.schema)) + assert(output.head().getAs[org.apache.spark.ml.linalg.Matrix]("dense") == Matrices.dense(1, 1, Array(1.0))) + assert(server.requests.get() == 1) + } + } + + test("nullable vector-map containers preserve non-nullable children and null parents") { + val vectors = spark.createDataFrame(Arrays.asList(Row(Vectors.dense(1.0))), + new StructType().add("vector", SQLDataTypes.VectorType, nullable = false)) + val weights = map_from_arrays(array(col("vector")), array(lit("weight"))) + val payload = when(col("id") === 0, struct(lit(1).as("k"), weights.as("weights"))) + val input = vectors.crossJoin(spark.range(2)).select(col("id"), lit("hello").as("text"), + payload.as("payload"), array(payload).as("items"), map(lit("item"), payload).as("mapping")) + assert(input.schema("payload").nullable) + assert(input.schema("payload").dataType.asInstanceOf[StructType].fields.forall(!_.nullable)) + withServer { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect().sortBy(_.getAs[Long]("id")) + val present = rows.head + Seq(present.getAs[Row]("payload"), present.getAs[Seq[Row]]("items").head, + present.getAs[Map[String, Row]]("mapping")("item")).foreach { value => + assert(value.getInt(0) == 1 && value.getMap[Vector, String](1) == Map(Vectors.dense(1.0) -> "weight")) + } + val absent = rows.last + val missing = Option.empty[Row].orNull + assert(absent.getAs[Row]("payload") == null && absent.getAs[Seq[Row]]("items") == Seq(missing)) + assert(absent.getAs[Map[String, Row]]("mapping") == Map("item" -> missing)) + assert(rows.forall(row => row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + assert(server.requests.get() == 2) + } + } + + Seq(false, true).foreach { withVectors => + test(s"case-distinct nested fields are retained by ordinal with vectors=$withVectors") { + val caseSensitive = spark.conf.get("spark.sql.caseSensitive") + try { + spark.conf.set("spark.sql.caseSensitive", "false") + val vectors = spark.createDataFrame(Arrays.asList(Row(Vectors.dense(1.0))), + new StructType().add("vector", SQLDataTypes.VectorType)) + val fields = Seq(lit(1).as("a"), lit(2).as("A")) ++ (if (withVectors) + Seq(map_from_arrays(array(col("vector")), array(lit("weight"))).as("weights")) else Seq.empty) + val payload = struct(fields: _*) + val metadata = new MetadataBuilder().putString("source", "case-distinct").build() + val input = vectors.select(lit("hello").as("text"), payload.as("payload", metadata), + array(payload).as("items"), map(lit("item"), payload).as("mapping")) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"{\"value\":\"ok\"}"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(false, true).foreach { parsed => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error") + if (parsed) stage.setPostProcessingOptions(Map("jsonSchema" -> "value STRING")) + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + val row = output.head() + Seq(row.getAs[Row]("payload"), row.getAs[Seq[Row]]("items").head, + row.getAs[Map[String, Row]]("mapping")("item")).foreach { value => + assert(value.getInt(0) == 1 && value.getInt(1) == 2) + if (withVectors) assert(value.getMap[Vector, String](2) == Map(Vectors.dense(1.0) -> "weight")) + } + assert(row.getAs[Row]("error") == null) + if (parsed) assert(row.getAs[Row]("answer").getString(0) == "ok") + else assert(row.getAs[String]("answer") == """{"value":"ok"}""") + } + val replacing = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("payload") + val output = replacing.transform(input) + assert(output.schema == replacing.transformSchema(input.schema)) + assert(output.head().getAs[String]("payload") == """{"value":"ok"}""") + assert(server.requests.size() == 3) + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + } + + test("vector map keys retain SQL identity through inference and output replacement") { + withServer { server => + val stage = prompt(server) + val schema = new StructType().add("dense", SQLDataTypes.VectorType).add("sparse", SQLDataTypes.VectorType) + val vectors = spark.createDataFrame(Arrays.asList( + Row(Vectors.dense(1.0), Vectors.sparse(1, Array(0), Array(1.0)))), schema) + val weights = map_from_arrays(array(col("dense"), col("sparse")), array(lit("dense"), lit("sparse"))) + val vectorValues = map_from_arrays(array(col("dense"), col("sparse")), array(col("sparse"), col("dense"))) + val nested = array(struct(weights.as("weights"), array(col("dense"), col("sparse")).as("vectors"))) + val optional = struct(when(lit(false), weights).as("weights"), when(lit(false), nested).as("nested")) + val metadata = new MetadataBuilder().putString("source", "vector-key-map").build() + val input = vectors.select(lit("test").as("text"), weights.as("weights", metadata), nested.as("nested"), + vectorValues.as("vectorValues"), optional.as("optional"), when(lit(false), optional).as("missing")) + assert(input.select(size(col("weights")), size(col("nested").getItem(0).getField("weights"))).head() == + Row(2, 2)) + val output = stage.transform(input).cache() + try { + assert(output.schema == stage.transformSchema(input.schema)) + Seq(col("weights"), col("nested").getItem(0).getField("weights")).foreach { column => + val entries = output.select(explode(column)).collect() + assert(entries.length == 2) + val keys = entries.map(row => row.getString(1) -> row.getAs[Vector](0)).toMap + assert(keys("dense").isInstanceOf[DenseVector] && keys("sparse").isInstanceOf[SparseVector]) + assert(keys.values.forall(_.toArray.toSeq == Seq(1.0))) + } + val nestedVectors = output.select(col("nested").getItem(0).getField("vectors")).head().getSeq[Vector](0) + assert(nestedVectors.head.isInstanceOf[DenseVector] && nestedVectors(1).isInstanceOf[SparseVector]) + val pairs = output.select(explode(col("vectorValues"))).collect().map { row => + row.getAs[Vector](0).getClass -> row.getAs[Vector](1).getClass + }.toSet + assert(pairs == Set(classOf[DenseVector] -> classOf[SparseVector], + classOf[SparseVector] -> classOf[DenseVector])) + assert(output.select(col("optional.weights").isNull, col("optional.nested").isNull, + col("missing").isNull).head() == Row(true, true, true)) + } finally output.unpersist() + val replacing = stage.setOutputCol("weights") + val replaced = replacing.transform(input).cache() + try { + assert(replaced.schema == replacing.transformSchema(input.schema)) + assert(replaced.select(col("weights"), size(col("nested").getItem(0).getField("weights"))).head() == + Row("answer", 2)) + } finally replaced.unpersist() + assert(server.requests.get() == 2) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestionsSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestionsSuite.scala new file mode 100644 index 00000000000..98c71b8e684 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/NativeQuestionsSuite.scala @@ -0,0 +1,182 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.ai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.sql.Row +import spray.json.{JsArray, JsNumber, JsObject, JsString, JsValue, enrichString} + +import java.util.Base64 +import scala.util.Try +import scala.util.Using + +class NativeQuestionsSuite extends TestBase { + // Public request examples, retrieved 2026-10-13: + // https://learn.microsoft.com/en-us/azure/foundry/foundry-models/how-to/use-foundry-models-microsoft-decision + private val systemOne = + """{ + |"team":{"type":"choice","instructions":"Which team should handle this request?", + | "criteria":{"billing":"Charges, invoices, and refunds","engineering":"Bugs, errors, and outages", + | "support":"How-to and account questions"}}, + |"severity":{"type":"score","instructions":"How severe is this issue?", + | "criteria":["Cosmetic","Minor","Major","Critical"]}, + |"repeat_contact":{"type":"noul","instructions":"Has the customer contacted support before?"} + |}""".stripMargin.parseJson + + // https://developers.openai.com/api/docs/guides/decisions + // https://developers.openai.com/api/reference/python/resources/decisions/methods/create + private val decisions = + """[ + |{"type":"predicate","name":"visible_damage","instructions":"Does the product have visible damage?"}, + |{"type":"choice","name":"department","instructions":"Which department should handle this complaint?", + | "choices":[{"value":"billing","description":"Payments, invoices, and refunds."}, + | {"value":"technical","description":"Problems using the product."}]}, + |{"type":"score","name":"severity","instructions":"How severe is this issue?", + | "levels":[{"label":"Cosmetic","description":"Appearance only; no lost functionality."}, + | {"label":"Workaround available","description":"A task fails, but another way works."}, + | {"label":"Fully blocked","description":"A task fails with no workaround."}]} + |]""".stripMargin.parseJson + + private def validate(api: String, text: String): Unit = NativeQuestions.validate(api, text.parseJson) + + private def invalid(api: String, text: String): Unit = { + intercept[IllegalArgumentException](validate(api, text)) + } + + test("System One accepts the public combined question and optional noul criteria examples") { + NativeQuestions.validate("systemone", systemOne) + validate("systemone", + """{"refund":{"type":"noul","instructions":"Does this reply promise the customer a refund?", + |"criteria":{"true":"Commits to returning money","false":"Makes no commitment about money"}}}""".stripMargin) + } + + test("System One enforces native named questions instead of a guessed common API shape") { + Seq( + "[]", "{}", + """{"q":{"type":"predicate","instructions":"Test"}}""", + """{"q":{"type":"noul"}}""", + """{"q":{"type":"noul","instructions":true}}""", + """{"q":{"type":"noul","instructions":"Test","criteria":{"yes":"Yes","no":"No"}}}""", + """{"q":{"type":"choice","instructions":"Test","choices":["one","two"]}}""", + """{"q":{"type":"choice","instructions":"Test","criteria":{"one":1,"two":"two"}}}""", + """{"q":{"type":"choice","instructions":"Test","criteria":{"one":"Only one"}}}""", + """{"q":{"type":"score","instructions":"Test","criteria":[{"label":"first"}]}}""", + """{"q":{"type":"score","instructions":"Test","criteria":["first",true]}}""" + ).foreach(invalid("systemone", _)) + } + + test("Decisions accepts the public ordered questions and typed distinct choice values") { + NativeQuestions.validate("decisions", decisions) + validate("decisions", + """[{"type":"choice","instructions":"Pick one","choices":[{"value":true},{"value":"true"}]}, + |{"type":"predicate","name":null,"instructions":"Test"}]""".stripMargin) + } + + test("Decisions rejects malformed types, duplicate names and ambiguous choice values") { + Seq( + "{}", "[]", + """[{"type":"noul","instructions":"Test"}]""", + """[{"type":"predicate","instructions":"Test","name":true}]""", + """[{"type":"predicate","name":"same","instructions":"A"}, + |{"type":"predicate","name":"same","instructions":"B"}]""".stripMargin, + """[{"type":"choice","instructions":"Pick","choices":[{"value":1},{"value":2}]}]""", + """[{"type":"choice","instructions":"Pick","choices":[{"value":true},{"value":true}]}]""", + """[{"type":"choice","instructions":"Pick","choices":[{"value":"a","description":42},{"value":"b"}]}]""", + """[{"type":"score","instructions":"Rate","levels":["Low","High"]}]""", + """[{"type":"score","instructions":"Rate","levels":[{"label":"first"},{"description":"No label"}]}]""" + ).foreach(invalid("decisions", _)) + val choices = (0 until 256).map(i => JsObject("value" -> JsString(i.toString))) + val tooMany = JsArray(JsObject("type" -> JsString("choice"), "instructions" -> JsString("Pick"), + "choices" -> JsArray(choices.toVector))) + intercept[IllegalArgumentException](NativeQuestions.validate("decisions", tooMany)) + } + + private def input(parts: JsValue*): JsValue = JsArray(JsObject( + "role" -> JsString("user"), "content" -> JsArray(parts.toVector))) + + private def image(url: String): JsValue = JsObject("type" -> JsString("input_image"), "image_url" -> JsString(url)) + + test("direct and facade Decisions inputs distinguish missing, blank and unchanged nonblank text") { + import spark.implicits._ + Seq("text", "json").foreach { mode => + Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions()).foreach { stage => + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val nonblank = " evidence\t" + val values = Seq(Option.empty[String].orNull, "", " \t\n", nonblank) + val data = values.zipWithIndex.map { case (text, id) => + (id, if (mode == "json" && text != null && text.nonEmpty) JsString(text).compactPrint else text) + }.toDF("id", "text") + stage.setUrl(server.url + "/v1").setModel("model").setInputCol("text").setInputMode(mode) + .setQuestions(NativeRequestFixtures.DecisionQuestions).setErrorCol("error") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect().sortBy(_.getAs[Int]("id")) + rows.take(2).foreach { row => + val error = AIJson.objectString(row.getAs[Row]("error").getString(0)) + assert(error.fields("code") == JsString("missing_input")) + assert(error.fields("message") == JsString("input is null or empty")) + } + assert(rows(2).getAs[Row]("error").getString(0).contains("invalid_request")) + assert(rows(3).getAs[Row]("error") == null && server.requests.size() == 1) + assert(AIJson.objectString(server.requests.peek()._2).fields("input") == JsString(nonblank)) + } + } + } + } + + test("Decisions accepts text and user-message content with inline image data") { + NativeQuestions.validateDecisionsInput(JsString("Evidence")) + NativeQuestions.validateDecisionsInput("""[{"role":"user","type":"message","content":"Evidence"}]""".parseJson) + NativeQuestions.validateDecisionsInput(input( + JsObject("type" -> JsString("input_text"), "text" -> JsString("Inspect the product.")), + image("data:image/png;base64,AQID"))) + } + + test("Decisions inline image syntax matches the basic base64 decoder without exposing content") { + val samples = Seq("", "A", "AQ", "AQI", "AQID", "AQ==", "AQI=", "AR==", "AQJ=", + "AQ=", "AQI==", "AQID=", "====", "=AQI", "AQ I", "AQ\nI", "AQ-I", "AQ_I", "private@image") + samples.foreach { encoded => + val expected = Try(Base64.getDecoder.decode(encoded)).toOption.exists(_.nonEmpty) + val actual = Try(NativeQuestions.validateDecisionsInput(input(image(s"data:image/png;base64,$encoded")))) + assert(actual.isSuccess == expected, s"Base64 syntax acceptance differed for ${encoded.length} characters") + actual.failed.foreach { error => + assert(error.isInstanceOf[IllegalArgumentException] && error.getCause == null) + assert(!error.getMessage.contains("private@image")) + } + } + } + + test("Decisions rejects hosted images, invalid base64, unsupported roles and modalities") { + val unsupported = Seq( + input(image("https://example.com/image.png")), + input(image("data:image/png;base64,not-base64")), + input(image("data:image/png;base64,")), + input(JsObject("type" -> JsString("input_image"), "file_id" -> JsString("file-example"))), + input(JsObject(image("data:image/png;base64,AQID").asJsObject.fields + + ("file_id" -> JsString("file-example")))), + input(JsObject(image("data:image/png;base64,AQID").asJsObject.fields + ("detail" -> JsNumber(42)))), + input(JsObject("type" -> JsString("input_audio"))), + """[{"role":"system","content":"Evidence"}]""".parseJson, + """[{"role":"user","content":null}]""".parseJson, + """[{"role":"user","content":[]}]""".parseJson, + """[{"role":"user","type":"function_call","content":"Evidence"}]""".parseJson, + """[{"role":"user","content":[{"type":"input_text","text":42}]}]""".parseJson + ) + unsupported.foreach(v => intercept[IllegalArgumentException](NativeQuestions.validateDecisionsInput(v))) + val inline = image("data:image/png;base64,AQID") + NativeQuestions.validateDecisionsInput(input(Seq.fill(128)(inline): _*)) + intercept[IllegalArgumentException] { + NativeQuestions.validateDecisionsInput(input(Seq.fill(129)(inline): _*)) + } + } + + test("Native validation errors do not disclose supplied values") { + val error = intercept[IllegalArgumentException] { + validate("decisions", """[{"type":"private-sensitive-value","instructions":"Secret"}]""") + } + assert(!error.getMessage.contains("private-sensitive-value")) + assert(!error.getMessage.contains("Secret")) + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsLiveSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsLiveSuite.scala new file mode 100644 index 00000000000..9fa2ff52f95 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsLiveSuite.scala @@ -0,0 +1,56 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.openai + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.services.ai.{AILiveTestEnvironment, AIPrompt, NativeRequestFixtures} +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{StringType, StructType} +import spray.json._ + +import scala.collection.JavaConverters._ + +/** Opt-in inference against an existing deployment; never provisions or changes resources. */ +class OpenAIDecisionsLiveSuite extends TestBase { + test("live Decisions supports typed questions through direct and facade APIs") { + val endpoint = AILiveTestEnvironment.get("SYNAPSEML_DECISIONS_ENDPOINT") + val model = AILiveTestEnvironment.get("SYNAPSEML_DECISIONS_MODEL") + val token = AILiveTestEnvironment.get("SYNAPSEML_DECISIONS_TOKEN") + val key = AILiveTestEnvironment.get("SYNAPSEML_DECISIONS_KEY") + assume(endpoint.isDefined && model.isDefined && (token.isDefined || key.isDefined), + "Live Decisions endpoint, model and credentials are not configured") + val input = spark.createDataFrame(Seq( + Row("Checkout is down again for every customer. This is my third request.")).asJava, + new StructType().add("ticket", StringType)) + Seq(new OpenAIDecisions(), new AIPrompt().setApiType("decisions")).foreach { stage => + stage.setUrl(endpoint.get).setModel(model.get).setInputCol("ticket") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setOutputCol("decision") + .setErrorCol("error").setUsageCol("usage").setRawResponseCol("raw") + token match { + case Some(value) => stage.setAADToken(value) + case None => key.foreach(stage.setSubscriptionKey) + } + val result = stage.transform(input) + assert(result.schema == stage.transformSchema(input.schema)) + val row = result.head() + assert(row.isNullAt(row.fieldIndex("error")), + "Live Decisions request failed; inspect the protected error channel") + val decision = row.getAs[Row]("decision") + val answers = decision.getAs[Seq[Row]]("answers") + assert(answers.map(_.getAs[String]("name")) == Seq("repeat_contact", "team", "severity")) + assert(answers.map(_.getAs[String]("type")) == Seq("predicate", "choice", "score")) + val probability = answers.head.getAs[Double]("probability") + assert(probability >= 0 && probability <= 1) + val choice = answers(1).getAs[Row]("choice") + assert(choice.getAs[String]("type") == "string") + assert(Set("billing", "engineering")(choice.getAs[String]("stringValue"))) + val score = answers(2).getAs[Double]("score") + assert(score >= 0 && score <= 2) + val raw = row.getAs[String]("raw").parseJson.asJsObject + assert(raw.fields("answers").asInstanceOf[JsArray].elements.size == answers.size) + assert(decision.getAs[String]("model").nonEmpty) + assert(!row.isNullAt(row.fieldIndex("usage"))) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsSuite.scala new file mode 100644 index 00000000000..834b0bb460a --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIDecisionsSuite.scala @@ -0,0 +1,22 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.services.openai + +import com.microsoft.azure.synapse.ml.core.test.fuzzing.TestObject +import com.microsoft.azure.synapse.ml.services.ai.{AIPromptFuzzing, NativeRequestFixtures} +import org.apache.spark.ml.util.MLReadable +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{StringType, StructType} + +import scala.collection.JavaConverters._ + +class OpenAIDecisionsSuite extends AIPromptFuzzing[OpenAIDecisions] { + override def reader: MLReadable[_] = OpenAIDecisions + + override def testObjects(): Seq[TestObject[OpenAIDecisions]] = Seq(new TestObject( + new OpenAIDecisions("decisions-fuzzing").setUrl("https://example.invalid/v1").setModel("synthetic-model") + .setSubscriptionKey("synthetic-key").setInputCol("input").setQuestions(NativeRequestFixtures.DecisionQuestions) + .setOutputCol("decision").setErrorCol("error"), + spark.createDataFrame(Seq(Row(Option.empty[String].orNull)).asJava, new StructType().add("input", StringType)))) +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptMultimodalRequestSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptMultimodalRequestSuite.scala index 439df77a5f0..9028138052c 100644 --- a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptMultimodalRequestSuite.scala +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptMultimodalRequestSuite.scala @@ -10,6 +10,7 @@ import org.apache.commons.io.{FileUtils, IOUtils} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.FileSystem import org.apache.spark.TaskKilledException +import org.apache.spark.injections.TaskContextUtils import org.apache.spark.ml.param.ParamMap import org.apache.spark.sql.Row import org.apache.spark.sql.functions.{col, lit, struct, to_json} @@ -407,7 +408,6 @@ class OpenAIPromptMultimodalRequestSuite extends TestBase { assert(!error.contains(missingUrl)) } } - test("attachment file reads propagate cancellation without converting ordinary IO failures") { withAttachmentFailureFileSystem { fs => def read(name: String): Unit = { @@ -440,7 +440,6 @@ class OpenAIPromptMultimodalRequestSuite extends TestBase { } finally Thread.interrupted() } } - Seq(false, true).foreach { includeCauses => test(s"attachment reads detect interruption in cyclic exception graphs (causes=$includeCauses)") { val (root, last) = AttachmentReadFailures.cyclicExceptionGraph(includeCauses) @@ -544,8 +543,9 @@ class OpenAIPromptMultimodalRequestSuite extends TestBase { }) } else { assert(!gate.completionInterrupted.get()) - assert(causes.exists { - case error: TaskKilledException => error.reason == "Attachment read canceled" + val reason = TaskContextUtils.getKillReason(gate.context.get()) + assert(reason.nonEmpty && causes.exists { + case error: TaskKilledException => reason.contains(error.reason) case _ => false }) } diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParserSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParserSuite.scala index 5ce4fb8eb2c..6f7f8b975c7 100644 --- a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParserSuite.scala +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptParserSuite.scala @@ -65,6 +65,53 @@ class OpenAIPromptParserSuite extends TestBase { assert(parser.outputSchema == StringType) } + test("JSON postprocessing diagnostics share explicit read options and nesting limits") { + val data = Seq("{'value':1.0}", """{"value":NaN}""", """{"value":2.0,"ignored":[[0]]}""").toDF("text") + Seq(true, false).foreach { enabled => + val parser = new JsonParser("value DOUBLE", Map( + "allowSINGLEQuotes" -> enabled.toString, "allowNonNumericNumbers" -> enabled.toString, + "maxNestingDepth" -> "3")) + val rows = data.select(parser.parse(col("text")).as("answer"), + parser.hasParseError(col("text")).as("invalid")).collect() + assert(rows.take(2).forall(_.getAs[Boolean]("invalid") == !enabled)) + assert(!rows.last.getAs[Boolean]("invalid")) + assert(rows.last.getAs[Row]("answer").getDouble(0) == 2.0) + } + } + + test("JSON postprocessing preserves trailing line comments when comments are enabled") { + val data = Seq( + """{"value":1} // trailing comment""", + """{"value":"wrong"} // trailing comment""", + "{\"value\":1} // trailing comment\n{\"value\":2}").toDF("text") + Seq(true, false).foreach { enabled => + val parser = new JsonParser("value INT", Map("allowComments" -> enabled.toString)) + val rows = data.select(parser.parse(col("text")).as("answer"), + parser.hasParseError(col("text")).as("invalid")).collect() + assert(rows.head.getAs[Row]("answer").getInt(0) == 1) + assert(rows.head.getAs[Boolean]("invalid") == !enabled) + assert(rows.tail.forall(_.getAs[Boolean]("invalid"))) + } + } + + test("JSON postprocessing diagnoses escaped surrogates in nested and ignored tokens without changing syntax") { + val high = "\\" + "ud800" + val low = "\\" + "udc00" + val pair = "\\" + "ud83d" + "\\" + "ude00" + val data = Seq( + s"""{"value":"$high"}""", s"""{"value":"$low"}""", + s"""{"value":"ordinary","ignored":[{"$high":"other"}]}""", + s"""{"value":"ordinary","ignored":[["$low"]]}""", + s"""{value:'$pair',ignored:['$pair']} // permitted""").toDF("text") + val parser = new JsonParser("value STRING", Map("allowUnquotedFieldNames" -> "true", + "allowSingleQuotes" -> "true", "allowComments" -> "true")) + val rows = data.select(parser.parse(col("text")).as("answer"), + parser.hasParseError(col("text")).as("invalid")).collect() + assert(rows.take(4).forall(_.getAs[Boolean]("invalid"))) + assert(!rows.last.getAs[Boolean]("invalid")) + assert(rows.last.getAs[Row]("answer").getString(0) == "\uD83D\uDE00") + } + test("stringMessageWrapper changes text type for responses API") { val prompt = new OpenAIPrompt() assert(prompt.stringMessageWrapper("hello") == Map("type" -> "text", "text" -> "hello")) diff --git a/core/src/main/python/synapse/ml/core/schema/Utils.py b/core/src/main/python/synapse/ml/core/schema/Utils.py index 30a3d850406..0ed25a078b4 100644 --- a/core/src/main/python/synapse/ml/core/schema/Utils.py +++ b/core/src/main/python/synapse/ml/core/schema/Utils.py @@ -144,7 +144,7 @@ def _set_params_via_setters(self, kwargs, skip_none=False): getattr(self, setter)(value) return self - def _transfer_params_from_java(self): + def _transfer_params_from_java(self, include_defaults=False): """ Transforms the embedded com.microsoft.azure.synapse.ml.core.serialize.params from the companion Java object. """ @@ -160,6 +160,15 @@ def _transfer_params_from_java(self): java_param.getClass(), ) is_service_param = self._is_service_param(java_param) + if include_defaults and not is_service_param and not is_complex_param: + default = self._java_obj.getDefault(java_param) + if default.isDefined(): + value = _java2py(sc, default.get()) + if isinstance(value, JavaObject): + value = json.loads(java_param.jsonEncode(default.get())) + self._setDefault(**{param.name: value}) + else: + self._defaultParamMap.pop(param, None) if self._java_obj.isSet(java_param): if is_complex_param: value = self._java_obj.getOrDefault(java_param) diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/Wrappable.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/Wrappable.scala index 321f64f2f30..c15e2b639b2 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/Wrappable.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/Wrappable.scala @@ -17,7 +17,6 @@ import java.nio.charset.StandardCharsets import java.nio.file.Files import java.util.WeakHashMap - trait BaseWrappable extends Params { protected val thisStage: Params = this @@ -29,7 +28,6 @@ trait BaseWrappable extends Params { protected lazy val classNameHelper: String = thisStage.getClass.getName.split(".".toCharArray).last - protected def companionModelClassName: String = { val symbol = scala.reflect.runtime.currentMirror.classSymbol(thisStage.getClass) val superClassSymbol = symbol.baseClasses @@ -75,6 +73,8 @@ trait PythonWrappable extends BaseWrappable { private def getPythonTypeInfo(p: Param[_]): PythonTypeInfo = PythonStubCodegen.pythonTypeInfo(thisStage, p) + protected def pyParamInputType(p: Param[_]): Option[String] = None + def pyAdditionalMethods: String = { "" } @@ -123,7 +123,7 @@ trait PythonWrappable extends BaseWrappable { // TODO add default values protected lazy val pyClassDoc: String = { val argLines = thisStage.params.map { p => - s"""${p.name} (${getPythonTypeInfo(p).pyType}): ${p.doc}""" + s"""${p.name} (${pyParamInputType(p).getOrElse(getPythonTypeInfo(p).pyType)}): ${p.doc}""" }.mkString("\n") s"""|"\"" |Args: @@ -341,7 +341,7 @@ trait PythonWrappable extends BaseWrappable { } private def pyStubParamArgs(p: Param[_]): Seq[String] = { - val pyiType = getPythonTypeInfo(p).pyiType + val pyiType = pyParamInputType(p).getOrElse(getPythonTypeInfo(p).pyiType) (p, safeGetDefault(p)) match { case (_: ServiceParam[_], _) => Seq( @@ -390,7 +390,7 @@ trait PythonWrappable extends BaseWrappable { private def pyStubParamSetter(p: Param[_]): String = { val capName = p.name.capitalize - val pyiType = getPythonTypeInfo(p).pyiType + val pyiType = pyParamInputType(p).getOrElse(getPythonTypeInfo(p).pyiType) p match { case _: ServiceParam[_] => s"""|def set$capName(self: _T, value: $pyiType) -> _T: ... @@ -619,7 +619,7 @@ trait PythonWrappable extends BaseWrappable { val paramGetters = thisStage.params.map(pyStubParamGetter).mkString("\n") s"""|$copyrightLines | - |from typing import Any, Dict, List, Optional, TypeVar + |from typing import Any, Dict, List, Optional, TypeVar, Union | |from pyspark.ml.evaluation import JavaEvaluator |from pyspark.ml.param import Param, ParamMap diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala new file mode 100644 index 00000000000..94232a84456 --- /dev/null +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala @@ -0,0 +1,254 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.io.http + +import org.apache.http.client.methods.HttpRequestBase +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.injections.TaskContextUtils +import org.apache.spark.internal.Logging + +import java.util.concurrent.{Callable, CancellationException, ConcurrentHashMap, ExecutionException, Executors, Future, + FutureTask, RejectedExecutionException, ScheduledFuture, ScheduledThreadPoolExecutor, TimeUnit, TimeoutException} +import java.util.concurrent.atomic.AtomicBoolean +import scala.annotation.tailrec +import scala.collection.JavaConverters._ +import scala.collection.mutable +import scala.concurrent.duration.Duration +import scala.util.control.NonFatal + +private[http] object CancellableRequests extends Logging { + private val PollMillis = 100L + private val PollNanos = TimeUnit.MILLISECONDS.toNanos(PollMillis) + private val DrainSeconds = 5L + private val DrainNanos = TimeUnit.SECONDS.toNanos(DrainSeconds) + private val Active = new ThreadLocal[Scope]() + private lazy val CancellationMonitor = { + val monitor = new ScheduledThreadPoolExecutor(1, (runnable: Runnable) => { + val thread = Executors.defaultThreadFactory().newThread(runnable) + thread.setName("synapseml-http-cancellation-monitor") + thread.setDaemon(true) + thread.setContextClassLoader(CancellableRequests.getClass.getClassLoader) + thread + }) + monitor.setRemoveOnCancelPolicy(true) + monitor + } + + private def preserveInterruption[T](action: => T): T = try action catch { + case error: InterruptedException => + Thread.currentThread().interrupt() + throw error + } + + def check(): Unit = { + if (Thread.currentThread().isInterrupted) { + throw new InterruptedException("HTTP request canceled") + } + if (Option(Active.get()).exists(_.isCancelled)) { + throw new CancellationException("HTTP request canceled") + } + } + + def withRequest[T](request: HttpRequestBase)(action: => T): T = { + val scope = Option(Active.get()) + scope.foreach(_.requests.add(request)) + try preserveInterruption { + check() + action + } finally scope.foreach(_.requests.remove(request)) + } + + private class Scope(concurrency: Int) extends AutoCloseable { + val cancelled = new AtomicBoolean(false) + val requests = ConcurrentHashMap.newKeySet[HttpRequestBase]() + val context: Option[TaskContext] = Option(TaskContext.get()) + private val pool = Executors.newFixedThreadPool(concurrency, (runnable: Runnable) => { + val thread = Executors.defaultThreadFactory().newThread(runnable) + thread.setName("synapseml-bounded-http-" + thread.getId) + thread.setDaemon(true) + thread + }) + private var closed = false + private val cleanupLock = new Object() + private var cleanupAttempted = false + private var abortWork: Option[FutureTask[Unit]] = None + private var cancellationWatch: Option[ScheduledFuture[_]] = None + + def isCancelled: Boolean = cancelled.get() || context.exists(_.isInterrupted()) + + def observeCancellation(): Unit = synchronized { + if (!closed) { + cancellationWatch = context.map { task => + CancellationMonitor.scheduleWithFixedDelay(new Runnable { + override def run(): Unit = if (task.isInterrupted()) cancel() + }, PollMillis, PollMillis, TimeUnit.MILLISECONDS) + } + } + } + + def checkTask(): Unit = { + check() + context.foreach { task => + if (task.isInterrupted()) { + throw new TaskKilledException(TaskContextUtils.getKillReason(task).getOrElse("HTTP request task canceled")) + } + } + if (cancelled.get()) { + throw new CancellationException("HTTP request canceled") + } + } + + def submit[T](action: => T): Future[T] = { + checkTask() + try pool.submit(new Callable[T] { + override def call(): T = { + Active.set(Scope.this) + try { + check() + action + } finally Active.remove() + } + }) catch { + case error: RejectedExecutionException => + checkTask() + throw error + } + } + + @tailrec + final def await[T](future: Future[T], started: Long, timeout: Duration): T = { + checkTask() + val remaining = if (timeout.isFinite) timeout.toNanos - (System.nanoTime() - started) else Long.MaxValue + val waitNanos = math.max(0L, math.min(remaining, PollNanos)) + val result = try Some(future.get(waitNanos, TimeUnit.NANOSECONDS)) catch { + case _: TimeoutException if remaining <= 0 => + checkTask() + throw new TimeoutException("HTTP concurrent request timed out") + case _: TimeoutException => None + case failure: ExecutionException => + checkTask() + throw failure.getCause + } + result match { + case Some(value) => value + case None => await(future, started, timeout) + } + } + + private def cancel(): Option[FutureTask[Unit]] = synchronized { + if (!closed) { + closed = true + cancelled.set(true) + cancellationWatch.foreach(_.cancel(false)) + pool.shutdownNow() + val activeRequests = requests.asScala.toVector + if (activeRequests.nonEmpty) { + val work = new FutureTask[Unit](new Callable[Unit] { + override def call(): Unit = activeRequests.foreach { request => + try request.abort() catch { + case NonFatal(error) => + logWarning(s"HTTP request abort failed (${error.getClass.getSimpleName})") + } + } + }) + abortWork = Some(work) + val thread = new Thread(work) + thread.setName("synapseml-http-abort-" + thread.getId) + thread.setDaemon(true) + thread.setContextClassLoader(CancellableRequests.getClass.getClassLoader) + thread.start() + } + } + abortWork + } + + override def close(): Unit = cleanupLock.synchronized { + if (!cleanupAttempted) { + cleanupAttempted = true + val work = cancel() + var interrupted = Thread.interrupted() + val started = System.nanoTime() + def remaining: Long = math.max(0L, DrainNanos - (System.nanoTime() - started)) + var failure: Option[Exception] = None + def attempt(action: => Unit): Unit = try action catch { + case error: Exception => + interrupted = interrupted || error.isInstanceOf[InterruptedException] + failure match { + case Some(original) => original.addSuppressed(error) + case None => failure = Some(error) + } + } + try { + attempt { + work.foreach { abort => + try abort.get(remaining, TimeUnit.NANOSECONDS) catch { + case error: TimeoutException => + abort.cancel(true) + throw new IllegalStateException( + "HTTP request abort did not stop; request quiescence is not confirmed", error) + case error: Exception => + abort.cancel(true) + throw error + } + } + } + attempt { + require(pool.awaitTermination(remaining, TimeUnit.NANOSECONDS), + "HTTP request workers did not stop; request quiescence is not confirmed") + } + failure.foreach(error => throw error) + } finally { + if (interrupted) Thread.currentThread().interrupt() + } + } + } + } + + def map[A, B](input: Iterator[A], concurrency: Int, timeout: Duration) + (send: A => B): Iterator[B] with AutoCloseable = { + val scope = new Scope(concurrency) + scope.context.foreach(_.addTaskCompletionListener[Unit](_ => scope.close())) + scope.observeCancellation() + + def guarded[T](action: => T): T = try preserveInterruption(action) catch { + case error: Exception => + try scope.close() catch { + case cleanup: Exception => error.addSuppressed(cleanup) + } + throw error + } + + new Iterator[B] with AutoCloseable { + private val pending = mutable.Queue.empty[Future[B]] + private var finished = false + + override def close(): Unit = { + finished = true + pending.clear() + scope.close() + } + + @tailrec + private def fill(): Unit = { + scope.checkTask() + if (pending.size < concurrency && input.hasNext) { + val request = input.next() + pending.enqueue(scope.submit(send(request))) + fill() + } + } + + override def hasNext: Boolean = if (finished) false else guarded { + fill() + if (pending.isEmpty) close() + !finished + } + + override def next(): B = guarded { + if (!hasNext) throw new NoSuchElementException("No more HTTP responses") + scope.await(pending.dequeue(), System.nanoTime(), timeout) + } + } + } +} diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala index 5e148ca8d55..e3acab7807d 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala @@ -6,12 +6,14 @@ package com.microsoft.azure.synapse.ml.io.http import com.microsoft.azure.synapse.ml.fabric.FabricClient import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging import org.apache.commons.io.IOUtils -import org.apache.http.HttpEntity +import org.apache.http.{Header, HttpEntity, HttpStatus} import org.apache.http.client.config.RequestConfig +import org.apache.http.client.utils.DateUtils import org.apache.http.client.methods.{CloseableHttpResponse, HttpPost, HttpRequestBase} -import org.apache.http.entity.{AbstractHttpEntity, BasicHttpEntity, ByteArrayEntity} +import org.apache.http.entity.{AbstractHttpEntity, BasicHttpEntity, ByteArrayEntity, StringEntity} import org.apache.http.impl.client.{CloseableHttpClient, HttpClientBuilder} import org.apache.http.impl.conn.PoolingHttpClientConnectionManager +import org.apache.http.message.BasicHttpResponse import org.apache.spark.injections.UDFUtils import org.apache.spark.internal.{Logging => SparkLogging} import org.apache.spark.sql.expressions.UserDefinedFunction @@ -19,6 +21,8 @@ import org.apache.spark.sql.types.StringType import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InputStream, SequenceInputStream} import java.nio.charset.StandardCharsets +import java.util.concurrent.{CancellationException, TimeUnit} +import scala.annotation.tailrec import scala.concurrent.duration.Duration import scala.concurrent.{ExecutionContext, blocking} import scala.util.{Random, Try} @@ -56,23 +60,52 @@ private[ml] trait HTTPClient extends BaseClient .setConnectionManager(connectionManager) .setDefaultRequestConfig(requestConfig).build() + private val cancellableClient: Client = HttpClientBuilder.create() + .setConnectionManager(connectionManager) + .setConnectionManagerShared(true) + .setDefaultRequestConfig(RequestConfig.copy(requestConfig).setNormalizeUri(false).build()) + .disableAutomaticRetries() + .disableCookieManagement() + .disableRedirectHandling().build() + override def close(): Unit = { - internalClient.close() + try cancellableClient.close() finally internalClient.close() } - protected def sendRequestWithContext(request: RequestWithContext): ResponseWithContext = { + protected def sendRequestWithContext(request: RequestWithContext): ResponseWithContext = + sendRequestWithClient(request, internalClient) + + private def sendRequestWithClient(request: RequestWithContext, client: Client): ResponseWithContext = { request.request.map(req => - ResponseWithContext(Some(handle(internalClient, req)), request.context) + ResponseWithContext(Some(handle(client, req)), request.context) ).getOrElse(ResponseWithContext(None, request.context)) } + // Outside a Spark task, callers must close the iterator if they do not consume it fully. + private[ml] def sendCancellableRequestsWithContext(requests: Iterator[RequestWithContext], + concurrency: Int, + timeout: Duration): Iterator[ResponseWithContext] + with AutoCloseable = + CancellableRequests.map(requests, concurrency, timeout)(request => + sendRequestWithClient(request, cancellableClient)) + } object HandlingUtils extends SparkLogging { + private class BufferedResponseEntity(val bytes: Array[Byte]) extends ByteArrayEntity(bytes) + private[ml] def convertAndClose(response: CloseableHttpResponse): HTTPResponseData = { var failure = Option.empty[Throwable] try { - new HTTPResponseData(response) + response.getEntity match { + case entity: BufferedResponseEntity => + HTTPResponseData(response.getAllHeaders.map(new HeaderData(_)), + Some(EntityData(entity.bytes, Option(entity.getContentEncoding).map(new HeaderData(_)), + Some(entity.getContentLength), Option(entity.getContentType).map(new HeaderData(_)), + entity.isChunked, entity.isRepeatable, entity.isStreaming)), + new StatusLineData(response.getStatusLine), response.getLocale.toString) + case _ => new HTTPResponseData(response) + } } catch { case error: Throwable => failure = Some(error) @@ -90,21 +123,44 @@ object HandlingUtils extends SparkLogging { type HandlerFunc = (CloseableHttpClient, HTTPRequestData) => HTTPResponseData + private def retryPause(milliseconds: Long, boundedPolicy: Boolean): Unit = { + if (!boundedPolicy) Thread.sleep(milliseconds) else { + require(milliseconds >= 0, "Retry delay cannot be negative") + val duration = TimeUnit.MILLISECONDS.toNanos(milliseconds) + val started = System.nanoTime() + @tailrec + def waitRemaining(): Unit = { + CancellableRequests.check() + val remaining = duration - (System.nanoTime() - started) + if (remaining > 0) { + TimeUnit.NANOSECONDS.sleep(remaining) + waitRemaining() + } + } + waitRemaining() + } + } + private def keepTrying(client: CloseableHttpClient, request: HttpRequestBase, retriesLeft: Array[Int], e: Throwable, extraCodesToRetry: Set[Int] = Set(), - backoff429Ms: Long = 0): CloseableHttpResponse = { + backoff429Ms: Long = 0, + boundedPolicy: Boolean = false, + quiet: Boolean = false, + refreshAuth: Option[String => String] = None): CloseableHttpResponse = { if (retriesLeft.isEmpty) { throw e } else { - Thread.sleep(retriesLeft.head.toLong) - sendWithRetries(client, request, retriesLeft.tail, extraCodesToRetry, backoff429Ms) + retryPause(retriesLeft.head.toLong, boundedPolicy) + sendWithRetryPolicy(client, request, retriesLeft.tail, extraCodesToRetry, backoff429Ms, + boundedPolicy, quiet, refreshAuth) } } private val MaxBackoffMs: Long = 60000L // 1 minute cap for 429 backoff + private val MaxRetryAfterLength = 64 private val MaxResponseInspectionBytes = 1024 * 1024L private def copyResponseEntityMetadata(source: HttpEntity, target: AbstractHttpEntity): Unit = { @@ -113,6 +169,53 @@ object HandlingUtils extends SparkLogging { target.setChunked(source.isChunked) } + private def preserveResponseReadFailure(buffered: BasicHttpResponse, redirect: Boolean): Unit = { + val bodyRead = buffered.getEntity != null + val message = if (bodyRead) { + "Could not read the HTTP response metadata; preserving its buffered body and status." + } else if (redirect) { + "Could not read the HTTP redirect response body; preserving its status without retrying." + } else { + "Could not read the HTTP error response body; preserving its status and available retry headers." + } + logWarning(message) + buffered.setLocale(java.util.Locale.ROOT) + if (!bodyRead) buffered.setEntity(new StringEntity(message, StandardCharsets.UTF_8)) + } + + private def bufferAndClose(response: CloseableHttpResponse): CloseableHttpResponse = { + try { + val buffered = new BasicHttpResponse(response.getStatusLine) with CloseableHttpResponse { + override def close(): Unit = () + } + val code = buffered.getStatusLine.getStatusCode + val redirect = code >= HttpStatus.SC_MULTIPLE_CHOICES && code < HttpStatus.SC_BAD_REQUEST + var retryAfter = Option.empty[Header] + try { + retryAfter = Option(response.getFirstHeader("Retry-After")) + .filter(h => code >= HttpStatus.SC_BAD_REQUEST && + retryAfterMillis(h.getValue, System.currentTimeMillis()).isDefined) + buffered.setHeaders(response.getAllHeaders) + buffered.setLocale(response.getLocale) + Option(response.getEntity).foreach { entity => + val bytes = new BufferedResponseEntity(IOUtils.toByteArray(entity.getContent)) + buffered.setEntity(bytes) + copyResponseEntityMetadata(entity, bytes) + } + } catch { + case _: java.io.IOException if redirect || code >= HttpStatus.SC_BAD_REQUEST => + CancellableRequests.check() + retryAfter.foreach(buffered.setHeader) + preserveResponseReadFailure(buffered, redirect) + } + buffered + } finally { + try response.close() catch { + case _: java.io.IOException => logWarning("Could not close the HTTP response.") + } + } + } + private def replayResponseEntity(response: CloseableHttpResponse, source: HttpEntity, bytes: Array[Byte], @@ -134,7 +237,11 @@ object HandlingUtils extends SparkLogging { } private[ml] def responseBodyForInspection(response: CloseableHttpResponse): Option[String] = { - Option(response.getEntity).flatMap { entity => + Option(response.getEntity).flatMap { + case entity: BufferedResponseEntity => + Some(new String(entity.bytes, 0, + math.min(entity.bytes.length.toLong, MaxResponseInspectionBytes).toInt, StandardCharsets.UTF_8)) + case entity => val output = new ByteArrayOutputStream() var input = Option.empty[InputStream] var keepInputOpen = false @@ -207,87 +314,156 @@ object HandlingUtils extends SparkLogging { retriesLeft: Array[Int], extraCodesToRetry: Set[Int] = Set(), backoff429Ms: Long = 0 - ): CloseableHttpResponse = { + ): CloseableHttpResponse = + sendWithRetryPolicy(client, request, retriesLeft, extraCodesToRetry, backoff429Ms, + boundedPolicy = false, quiet = false) + + private[ml] def retryAfterMillis(value: String, nowMillis: Long): Option[Long] = + Option(value).filter(_.length <= MaxRetryAfterLength).flatMap { hint => + Try(BigDecimal(hint) * 1000).toOption.filter(_ >= 0) + .map { millis => + if (millis == 0) 0L + else millis.max(BigDecimal(1)).min(BigDecimal(Long.MaxValue)) + .setScale(0, BigDecimal.RoundingMode.CEILING).toLong + } + .orElse(Try(Option(DateUtils.parseDate(hint))).toOption.flatten + .map(date => math.max(0L, date.getTime - nowMillis))) + } + + private[ml] def sendWithRetryPolicy(client: CloseableHttpClient, + request: HttpRequestBase, + retriesLeft: Array[Int], + extraCodesToRetry: Set[Int], + backoff429Ms: Long, + boundedPolicy: Boolean, + quiet: Boolean, + refreshAuth: Option[String => String] = None): CloseableHttpResponse = { + require(boundedPolicy || refreshAuth.isEmpty, "Authentication replay requires the bounded retry policy") + if (boundedPolicy) CancellableRequests.check() + var executingRequest = true try { - val response = client.execute(request) + val received = client.execute(request) + // Read failures belong to this attempt's retry budget, not to successful response decoding. + val response = if (boundedPolicy) bufferAndClose(received) else received + if (boundedPolicy) CancellableRequests.check() + executingRequest = false val code = response.getStatusLine.getStatusCode + lazy val boundedRetryAfterMs = if (boundedPolicy) { + Option(response.getFirstHeader("Retry-After")) + .flatMap(h => retryAfterMillis(h.getValue, System.currentTimeMillis())) + } else None //scalastyle:off magic.number - val dontRetry = code match { - case 200 => true - case 201 => true - case 202 => true - case 429 => - // Inspect body to distinguish capacity errors from transient rate limits. - if (capacityLimitExceeded(response)) { - // Fabric capacity-exceeded 429s are NOT transient rate limits — - // retrying will not help and causes hangs - logWarning(s"Capacity limit exceeded (non-retryable 429) on ${request.getURI}") - true - } else { - Option(response.getFirstHeader("Retry-After")) - .foreach { h => - logInfo(s"waiting ${h.getValue} on ${ - request match { - case p: HttpPost => p.getURI + " " + - Try(IOUtils.toString(p.getEntity.getContent, "UTF-8")).getOrElse("") - case _ => request.getURI - } - }") - } - false + if (code == 401 && refreshAuth.isDefined && retriesLeft.nonEmpty) { + if (boundedRetryAfterMs.exists(_ > MaxBackoffMs)) { + response + } else { + response.close() + boundedRetryAfterMs.foreach(delay => retryPause(delay, boundedPolicy)) + val rejected = Option(request.getFirstHeader("Authorization")).map(_.getValue).getOrElse("") + val refreshed = try Some(refreshAuth.get(rejected)) catch { + case canceled: CancellationException => throw canceled + case NonFatal(_) if boundedPolicy => + CancellableRequests.check() + logWarning("Fabric authentication refresh failed; returning the rejected HTTP response") + None } - case code => - logWarning(s"got error $code: ${response.getStatusLine.getReasonPhrase} on ${ - request match { - case p: HttpPost => p.getURI + " " + - Try(IOUtils.toString(p.getEntity.getContent, "UTF-8")).getOrElse("") - case _ => request.getURI + refreshed match { + case Some(auth) => + request.setHeader("Authorization", auth) + sendWithRetryPolicy(client, request, retriesLeft.tail, extraCodesToRetry, backoff429Ms, + boundedPolicy, quiet, None) + case None => response + } + } + } else { + val dontRetry = code match { + case code if boundedPolicy && code >= 200 && code < 400 => true + case 200 => true + case 201 => true + case 202 => true + case 429 => + // Inspect body to distinguish capacity errors from transient rate limits. + if (capacityLimitExceeded(response)) { + // Fabric capacity-exceeded 429s are NOT transient rate limits — + // retrying will not help and causes hangs + if (!quiet) logWarning(s"Capacity limit exceeded (non-retryable 429) on ${request.getURI}") + true + } else { + Option(response.getFirstHeader("Retry-After")).filter(_ => !quiet) + .foreach { h => + logInfo(s"waiting ${h.getValue} on ${ + request match { + case p: HttpPost => p.getURI + " " + + Try(IOUtils.toString(p.getEntity.getContent, "UTF-8")).getOrElse("") + case _ => request.getURI + } + }") + } + false } - }") + case code => + if (!quiet) logWarning(s"got error $code: ${response.getStatusLine.getReasonPhrase} on ${ + request match { + case p: HttpPost => p.getURI + " " + + Try(IOUtils.toString(p.getEntity.getContent, "UTF-8")).getOrElse("") + case _ => request.getURI + } + }") - if (extraCodesToRetry(code)) { - false + if (extraCodesToRetry(code)) { + false + } else { + code.toString.startsWith("4") // Retry only when code isn't a 4XX + } + } + //scalastyle:on magic.number + if (dontRetry || retriesLeft.isEmpty) { + response + } else { + // Bounded requests honor seconds and HTTP dates, terminating rather than retrying a long hint early. + // Legacy 429 handling caps numeric hints and falls back to exponential backoff for HTTP dates. + // Retry-After: 0 means "retry immediately" (RFC 7231). + val retryAfterMs = if (boundedPolicy) { + boundedRetryAfterMs + } else if (code == 429) { + Option(response.getFirstHeader("Retry-After")) + .flatMap(h => Try(h.getValue.toLong * 1000).toOption) + .filter(_ >= 0) + .map(math.min(_, MaxBackoffMs)) + } else None + if (boundedPolicy && retryAfterMs.exists(_ > MaxBackoffMs)) { + response } else { - code.toString.startsWith("4") // Retry only when code isn't a 4XX - } - } - //scalastyle:on magic.number - if (dontRetry || retriesLeft.isEmpty) { - response - } else { - // Retry-After may be delta-seconds (e.g. "120") or an HTTP-date. - // Parse as Long; if it's an HTTP-date or otherwise non-numeric, fall back to exponential backoff. - // Cap server-provided values to MaxBackoffMs; Retry-After: 0 means "retry immediately" (RFC 7231). - val retryAfterMs = if (code == 429) { - Option(response.getFirstHeader("Retry-After")) - .flatMap(h => Try(h.getValue.toLong * 1000).toOption) - .filter(_ >= 0) - .map(math.min(_, MaxBackoffMs)) - } else None - response.close() - if (code == 429) { // Do not count rate limiting in number of failures - val baseBackoff = retryAfterMs.getOrElse { - // No usable Retry-After header; use exponential backoff - val current = math.max(backoff429Ms, retriesLeft.head.toLong) - math.min(current * 2, MaxBackoffMs) + response.close() + if (code == 429) { // Legacy callers do not count rate limiting in number of failures. + val baseBackoff = retryAfterMs.getOrElse { + // No usable Retry-After header; use exponential backoff + val current = math.max(backoff429Ms, retriesLeft.head.toLong) + math.min(current * 2, MaxBackoffMs) + } + // Add jitter (up to 10% of base) to prevent thundering herd + val jitter = Random.nextInt(math.max((baseBackoff / 10).toInt, 1)) + val sleepMs = math.min(baseBackoff + jitter, MaxBackoffMs) + val source = if (retryAfterMs.isDefined) "Retry-After header" else "exponential backoff" + if (!quiet) logInfo(s"429 rate-limited on ${request.getURI}: " + + s"sleeping ${sleepMs}ms ($source, jitter=${jitter}ms)") + retryPause(sleepMs, boundedPolicy) + sendWithRetryPolicy(client, request, if (boundedPolicy) retriesLeft.tail else retriesLeft, + extraCodesToRetry, baseBackoff, boundedPolicy, quiet, refreshAuth) + } else { + retryPause(if (boundedPolicy) retryAfterMs.getOrElse(retriesLeft.head.toLong) + else retriesLeft.head.toLong, boundedPolicy) + sendWithRetryPolicy(client, request, retriesLeft.tail, extraCodesToRetry, 0, + boundedPolicy, quiet, refreshAuth) + } } - // Add jitter (up to 10% of base) to prevent thundering herd - val jitter = Random.nextInt(math.max((baseBackoff / 10).toInt, 1)) - val sleepMs = math.min(baseBackoff + jitter, MaxBackoffMs) - val source = if (retryAfterMs.isDefined) "Retry-After header" else "exponential backoff" - logInfo(s"429 rate-limited on ${request.getURI}: " + - s"sleeping ${sleepMs}ms ($source, jitter=${jitter}ms)") - Thread.sleep(sleepMs) - sendWithRetries(client, request, retriesLeft, extraCodesToRetry, baseBackoff) - } else { - Thread.sleep(retriesLeft.head.toLong) - sendWithRetries(client, request, retriesLeft.tail, extraCodesToRetry) } } } catch { - case e: java.io.IOException => - logError("Encountering a connection error", e) - keepTrying(client, request, retriesLeft, e, extraCodesToRetry, backoff429Ms) + case e: java.io.IOException if executingRequest || !boundedPolicy => + if (quiet) logWarning("HTTP connection failed") else logError("Encountering a connection error", e) + keepTrying(client, request, retriesLeft, e, extraCodesToRetry, backoff429Ms, + boundedPolicy, quiet, refreshAuth) } } //scalastyle:on method.length @@ -441,6 +617,37 @@ object HandlingUtils extends SparkLogging { def advancedUDF(retryTimes: Int*): UserDefinedFunction = UDFUtils.oldUdf(advanced(retryTimes: _*) _, StringType) + /** Internal retry policy for the cancellable client; interruption and programming errors propagate. */ + private[ml] def bounded(retryTimes: Int*)(client: CloseableHttpClient, + request: HTTPRequestData): HTTPResponseData = + boundedWithFabricAuth(retryTimes, + uri => FabricClient.isOpenAIEndpoint(uri), + rejected => FabricClient.refreshCognitiveMWCTokenAuthHeader(rejected))(client, request) + + private[ml] def boundedWithFabricAuth(retryTimes: Seq[Int], + isFabricEndpoint: String => Boolean, + refreshFabricAuth: String => String)( + client: CloseableHttpClient, request: HTTPRequestData): HTTPResponseData = { + CancellableRequests.check() + val httpRequest = request.toHTTPCore + val refresh = if (request.usesFabricAuth && isFabricEndpoint(request.requestLine.uri)) { + Some(refreshFabricAuth) + } else None + try { + CancellableRequests.withRequest(httpRequest) { + convertAndClose(sendWithRetryPolicy(client, httpRequest, retryTimes.toArray, Set.empty, 0, + boundedPolicy = true, quiet = true, refreshAuth = refresh)) + } + } catch { + case _: java.io.IOException => + CancellableRequests.check() + logWarning("HTTP connection failed after retry exhaustion") + Option.empty[HTTPResponseData].orNull + } finally { + httpRequest.releaseConnection() + } + } + def basic(client: CloseableHttpClient, request: HTTPRequestData): HTTPResponseData = { val req = request.toHTTPCore val data = convertAndClose(client.execute(req)) diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala index 4fba4c03852..f2eed67b093 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala @@ -19,7 +19,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.{DataFrame, Dataset, Row} import scala.concurrent.ExecutionContext -import scala.concurrent.duration.Duration +import scala.concurrent.duration.{Duration, SECONDS} trait HasHandler extends Params { val handler: UDFParam = new UDFParam( @@ -104,13 +104,14 @@ class HTTPTransformer(val uid: String) def this() = this(Identifiable.randomUID("HTTPTransformer")) + private def concurrentWait: Duration = get(concurrentTimeout) + .map(ct => Duration.fromNanos((ct * SECONDS.toNanos(1)).toLong)).getOrElse(Duration.Inf) + val clientHolder = SharedVariable { getConcurrency match { case 1 => new SingleThreadedHTTPClient(getHandler, (getTimeout * 1000).toInt) case n if n > 1 => - val dur = get(concurrentTimeout) - .map(ct => Duration.fromNanos((ct * math.pow(10, 9)).toLong)) //scalastyle:ignore magic.number - .getOrElse(Duration.Inf) + val dur = concurrentWait val ec = ExecutionContext.global new AsyncHTTPClient(getHandler, n, dur, (getTimeout * 1000).toInt)(ec) } @@ -119,21 +120,31 @@ class HTTPTransformer(val uid: String) /** @param dataset - The input dataset, to be transformed * @return The DataFrame that results from column selection */ - override def transform(dataset: Dataset[_]): DataFrame = { + override def transform(dataset: Dataset[_]): DataFrame = transformRequests(dataset, cancellable = false) + + private[ml] def transformCancellable(dataset: Dataset[_]): DataFrame = + transformRequests(dataset, cancellable = true) + + private def transformRequests(dataset: Dataset[_], cancellable: Boolean): DataFrame = { logTransform[DataFrame]({ val df = dataset.toDF() val enc = ExpressionEncoder(transformSchema(df.schema)) val colIndex = df.schema.fieldNames.indexOf(getInputCol) val fromRow = HTTPRequestData.makeFromRowConverter val toRow = HTTPResponseData.makeToRowConverter + val concurrency = getConcurrency + val wait = if (cancellable) concurrentWait else Duration.Inf df.mapPartitions { it => if (!it.hasNext) { Iterator() } else { val c = clientHolder.get - val responsesWithContext = c.sendRequestsWithContext(it.map { row => + val requests = it.map { row => c.RequestWithContext(Option(row.getStruct(colIndex)).map(fromRow), Some(row)) - }) + } + val responsesWithContext = if (cancellable) { + c.sendCancellableRequestsWithContext(requests, concurrency, wait) + } else c.sendRequestsWithContext(requests) responsesWithContext.map { rwc => Row.fromSeq(rwc.context.get.asInstanceOf[Row].toSeq :+ rwc.response.flatMap(Option.apply).map(toRow).orNull) } diff --git a/core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala b/core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala new file mode 100644 index 00000000000..8ff56f0fc77 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala @@ -0,0 +1,10 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.injections + +import org.apache.spark.TaskContext + +object TaskContextUtils { + def getKillReason(context: TaskContext): Option[String] = context.getKillReason() +} diff --git a/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala b/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala index 94fd9141eab..21a9e15f2f7 100644 --- a/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala +++ b/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala @@ -111,6 +111,14 @@ private[ml] object ComplexParamsWriter { spark: SparkSession, extraMetadata: Option[JObject] = None, paramMap: Option[JValue] = None): String = { + getMetadataToSave(instance, spark, extraMetadata, paramMap, None) + } + + def getMetadataToSave(instance: Params, + spark: SparkSession, + extraMetadata: Option[JObject], + paramMap: Option[JValue], + defaultParamMap: Option[JValue]): String = { val uid = instance.uid val cls = instance.getClass.getName val params = instance.extractParamMap().toSeq @@ -120,9 +128,9 @@ private[ml] object ComplexParamsWriter { val jsonParams = paramMap.getOrElse(render(params.map { case ParamPair(p, v) => p.name -> parse(p.jsonEncode(v)) }.toList)) - val jsonDefaultParams = render(defaultParams.map { case ParamPair(p, v) => + val jsonDefaultParams = defaultParamMap.getOrElse(render(defaultParams.map { case ParamPair(p, v) => p.name -> parse(p.jsonEncode(v)) - }.toList) + }.toList)) val basicMetadata = ("class" -> cls) ~ ("timestamp" -> System.currentTimeMillis()) ~ ("sparkVersion" -> spark.version) ~ diff --git a/core/src/main/scala/org/apache/spark/ml/ExplicitComplexParamsWritable.scala b/core/src/main/scala/org/apache/spark/ml/ExplicitComplexParamsWritable.scala new file mode 100644 index 00000000000..3c9f87132a9 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/ExplicitComplexParamsWritable.scala @@ -0,0 +1,116 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import com.fasterxml.jackson.core.exc.StreamConstraintsException +import com.fasterxml.jackson.databind.ObjectMapper +import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.param.{ParamPair, Params} +import org.apache.spark.ml.util.{DefaultParamsReader, MLReader, MLWriter} +import org.apache.spark.util.Utils +import org.json4s.{DefaultFormats, JValue} +import org.json4s.JsonDSL._ +import org.json4s.jackson.JsonMethods +import org.json4s.jackson.JsonMethods.{compact, parse, render} + +import java.nio.charset.StandardCharsets + +private[ml] object ExplicitParamsJson extends JsonMethods { + private lazy val Mapper: ObjectMapper = { + val result = super.mapper.copy() + // Native parameter strings have no smaller parser limit; do not mutate json4s' shared factory. + val factory = result.getFactory + factory.setStreamReadConstraints(factory.streamReadConstraints().rebuild().maxStringLength(Int.MaxValue).build()) + result + } + + override def mapper: ObjectMapper = Mapper +} + +/** Opt-in persistence for stages whose behavior distinguishes explicit Params from defaults. + * + * ComplexParam defaults are not persisted and must be reproducible by the UID constructor. + * This writer replaces earlier write overrides; validateForWrite must preserve their required checks. + */ +trait ExplicitComplexParamsWritable extends ComplexParamsWritable { + self: Params => + + /** Validate constructor-reproducible ComplexParam defaults and checks from replaced write overrides. */ + protected def validateForWrite(): Unit + + override def write: MLWriter = { + validateForWrite() + new MLWriter { + private var prepared: Option[(String, Map[ParamPair[_], Path])] = None + + override def save(path: String): Unit = { + validateForWrite() + def encode[T](pair: ParamPair[T]): (String, JValue) = { + val encoded = parse(pair.param.jsonEncode(pair.value), useBigDecimalForDouble = true) + pair.param.jsonDecode(compact(render(encoded))) + pair.param.name -> encoded + } + + def json(pairs: Seq[ParamPair[_]]): JValue = + render(pairs.filterNot(_.param.isInstanceOf[ComplexParam[_]]).map(pair => encode(pair)).toList) + + val locations = ComplexParamsWriter.getComplexParamLocations(self, path) + .filter { case (pair, _) => self.isSet(pair.param) } + val metadata = ComplexParamsWriter.getComplexMetadata(locations) + val metadataJson = try { + val encoded = ComplexParamsWriter.getMetadataToSave(self, sparkSession, metadata, + Some(json(self.paramMap.toSeq)), Some(json(self.defaultParamMap.toSeq))) + require(StandardCharsets.UTF_8.newEncoder().canEncode(encoded), + "Stage metadata cannot be encoded losslessly as UTF-8. Correct unpaired surrogates in string parameters.") + // Spark's Pipeline loaders parse metadata before invoking this stage's reader. + DefaultParamsReader.parseMetadata(encoded) + encoded + } catch { + case _: StreamConstraintsException => + throw new IllegalArgumentException( + "Stage metadata exceeds this runtime's Spark JSON parsing limits. " + + "Reduce saved parameter sizes or use input columns for large values.") + } + prepared = Some(metadataJson -> locations) + try super.save(path) finally prepared = None + } + + override protected def saveImpl(path: String): Unit = { + val (metadataJson, locations) = prepared.getOrElse { + throw new IllegalStateException("Stage metadata must be prepared before saving.") + } + val metadataPath = new Path(path, "metadata").toString + sparkSession.createDataFrame(Seq(Tuple1(metadataJson))).toDF("value").write.text(metadataPath) + ComplexParamsWriter.saveComplexParams(path, locations, shouldOverwrite, sparkSession) + } + } + } +} + +/** Restores opt-in exact JSON numbers while retaining legacy parameter decoding for other stages. */ +trait ExplicitComplexParamsReadable[T] extends ComplexParamsReadable[T] { + protected def expectedClass: Class[T] + + override def read: MLReader[T] = new MLReader[T] { + override def load(path: String): T = { + val metadataPath = new Path(path, "metadata").toString + val text = sparkSession.read.text(metadataPath).first().getString(0) + implicit val formats: DefaultFormats.type = DefaultFormats + val className = (ExplicitParamsJson.parse(text) \ "class").extract[String] + val cls = Utils.classForName(className, initialize = false) + require(expectedClass.isAssignableFrom(cls), "Saved stage class is incompatible with the requested reader.") + val metadata = if (classOf[ExplicitComplexParamsWritable].isAssignableFrom(cls)) { + val json = ExplicitParamsJson.parse(text, useBigDecimalForDouble = true) + DefaultParamsReader.Metadata(className, (json \ "uid").extract[String], + (json \ "timestamp").extract[Long], (json \ "sparkVersion").extract[String], + json \ "paramMap", json \ "defaultParamMap", json, text) + } else DefaultParamsReader.parseMetadata(text) + val instance = cls.getConstructor(classOf[String]).newInstance(metadata.uid).asInstanceOf[Params] + metadata.getAndSetParams(instance) + ComplexParamsReader.getAndSetComplexParams(instance, metadata, path, sparkSession) + instance.asInstanceOf[T] + } + } +} diff --git a/core/src/main/scala/org/apache/spark/ml/NamespaceInjections.scala b/core/src/main/scala/org/apache/spark/ml/NamespaceInjections.scala index 62e70df49b6..98a4e80e322 100644 --- a/core/src/main/scala/org/apache/spark/ml/NamespaceInjections.scala +++ b/core/src/main/scala/org/apache/spark/ml/NamespaceInjections.scala @@ -4,6 +4,7 @@ package org.apache.spark.ml import org.apache.spark.ml.image.ImageSchema +import org.apache.spark.ml.param.{Param, Params} import org.apache.spark.ml.param.shared.HasParallelism import org.apache.spark.ml.util.Identifiable import org.apache.spark.sql.Row @@ -27,6 +28,10 @@ object ImageInjections { } object ParamInjections { + def clearDefault[T](params: Params, param: Param[T]): Unit = { + params.defaultParamMap.remove(param) + } + trait HasParallelismInjected extends HasParallelism { def getExecutionContextProxy: ExecutionContext = { diff --git a/core/src/main/scala/org/apache/spark/sql/internal/SQLConfInjections.scala b/core/src/main/scala/org/apache/spark/sql/internal/SQLConfInjections.scala new file mode 100644 index 00000000000..1758e738a04 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/sql/internal/SQLConfInjections.scala @@ -0,0 +1,12 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.sql.internal + +import com.fasterxml.jackson.core.JsonFactory +import org.apache.spark.sql.catalyst.json.JSONOptions + +object SQLConfInjections { + def jsonFactory(options: Map[String, String]): JsonFactory = + new JSONOptions(options, SQLConf.get.sessionLocalTimeZone).buildJsonFactory() +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala index 65131a47e21..ba374ebdc40 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala @@ -19,6 +19,7 @@ import org.apache.spark.ml.{ ComplexParamsReadable, ComplexParamsTestHelper, ComplexParamsWritable, + ExplicitComplexParamsWritable, ObjectSerializer, PipelineStage, Serializer @@ -120,6 +121,11 @@ class MixedParamTest(override val uid: String) extends TestEstimatorBase(uid) object MixedParamTest extends ComplexParamsReadable[MixedParamTest] +class TestableExplicitSessionParamTest(override val uid: String) extends TestableSessionParamTest(uid) + with ExplicitComplexParamsWritable { + override protected def validateForWrite(): Unit = () +} + class ValidateComplexParamSerializer extends TestBase { val saveFile = new File(tmpDir.toFile, "m1.model").toString val saveFile2 = new File(tmpDir.toFile, "m2.model").toString @@ -141,16 +147,19 @@ class ValidateComplexParamSerializer extends TestBase { } } - test("complex writers use their assigned Spark session") { - val stage = new TestableSessionParamTest("assigned-session") - val session = spark.newSession() - session.conf.set(TestableSessionParamTest.ConfigKey, stage.uid) - withActiveSession(Some(spark)) { - val bytes = Array[Byte](1, 2, 3) - val path = tmpDir.resolve("assigned-session.model").toString - stage.set(stage.byteArray, bytes).write.session(session).save(path) - val restored = TestableSessionParamTest.read.session(session).load(path) - assert(restored.getOrDefault(restored.byteArray).sameElements(bytes)) + Seq(false, true).foreach { explicit => + test(s"complex writers use their assigned Spark session with explicit=$explicit") { + val stage = if (explicit) new TestableExplicitSessionParamTest("explicit-session") + else new TestableSessionParamTest("legacy-session") + val session = spark.newSession() + session.conf.set(TestableSessionParamTest.ConfigKey, stage.uid) + withActiveSession(Some(spark)) { + val bytes = Array[Byte](1, 2, 3) + val path = tmpDir.resolve(s"session-$explicit.model").toString + stage.set(stage.byteArray, bytes).write.session(session).save(path) + val restored = TestableSessionParamTest.read.session(session).load(path) + assert(restored.getOrDefault(restored.byteArray).sameElements(bytes)) + } } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala new file mode 100644 index 00000000000..cd10d2bf3a2 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala @@ -0,0 +1,797 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.io.split1 + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HTTPSchema, HTTPTransformer, HandlingUtils, + SingleThreadedHTTPClient} +import com.sun.net.httpserver.{HttpExchange, HttpServer} +import org.apache.commons.io.IOUtils +import org.apache.http.{Header, HttpEntity, HttpHost, HttpRequest, ProtocolVersion, StatusLine} +import org.apache.http.client.methods.{CloseableHttpResponse, HttpPost} +import org.apache.http.conn.ClientConnectionManager +import org.apache.http.entity.{BasicHttpEntity, ByteArrayEntity, StringEntity} +import org.apache.http.impl.client.{CloseableHttpClient, HttpClients} +import org.apache.http.message.BasicHttpResponse +import org.apache.http.params.{BasicHttpParams, HttpParams} +import org.apache.http.protocol.HttpContext +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.StructType +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} + +import java.io.{IOException, InputStream, OutputStream} +import java.net.{InetSocketAddress, SocketException} +import java.nio.charset.StandardCharsets +import java.util.concurrent.{CancellationException, ConcurrentLinkedQueue, CountDownLatch, Executors, + Semaphore, TimeUnit, TimeoutException} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference} +import scala.collection.JavaConverters._ +import scala.concurrent.duration._ +import scala.util.Using + +object VerifyBoundedRequests { + private[split1] def client(next: () => CloseableHttpResponse): CloseableHttpClient = new CloseableHttpClient { + override protected def doExecute(host: HttpHost, request: HttpRequest, + context: HttpContext): CloseableHttpResponse = next() + override def getParams: HttpParams = new BasicHttpParams() + override def getConnectionManager: ClientConnectionManager = throw new UnsupportedOperationException() + override def close(): Unit = () + } + + private[split1] def response(code: Int, entity: HttpEntity, onClose: () => Unit): CloseableHttpResponse = { + val result = new BasicHttpResponse(new ProtocolVersion("HTTP", 1, 1), code, "Synthetic") + with CloseableHttpResponse { + override def close(): Unit = onClose() + } + result.setEntity(entity) + result.setLocale(java.util.Locale.ROOT) + result + } +} + +class VerifyBoundedRequests extends TestBase { + import VerifyBoundedRequests.{client, response} + + private def truncatedEntity(): HttpEntity = { + val entity = new BasicHttpEntity() + entity.setContent(new InputStream { + private var first = true + override def read(): Int = { + if (first) { + first = false + '{'.toInt + } else throw new SocketException("Synthetic body reset") + } + }) + entity.setContentLength(100) + entity + } + + private def withServer(reply: HttpExchange => Unit)(body: String => Unit): Unit = { + val server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + val pool = Executors.newFixedThreadPool(2) + server.setExecutor(pool) + server.createContext("/test", (exchange: HttpExchange) => reply(exchange)) + server.start() + try body(s"http://127.0.0.1:${server.getAddress.getPort}/test") + finally { + server.stop(0) + pool.shutdownNow() + } + } + + private def cancelRemaining[A](results: Iterator[A]): Unit = { + val interrupted = Thread.interrupted() + Thread.currentThread().interrupt() + try { + val stopped = try !results.hasNext catch { + case error: InterruptedException => + assert(error.getSuppressed.isEmpty, "HTTP worker cleanup failed") + true + } + assert(stopped, "HTTP iterator did not stop after cancellation") + } finally { + Thread.interrupted() + if (interrupted) Thread.currentThread().interrupt() + } + } + + test("already buffered HTTP responses survive the minimum concurrent timeout without waiting") { + val worker = new AtomicReference[Thread]() + val sent = new CountDownLatch(1) + val reply = HTTPSchema.stringToResponse("buffered", 200, "OK") + val handler: HandlingUtils.HandlerFunc = (_, _) => { + worker.set(Thread.currentThread()) + sent.countDown() + reply + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val requests = Iterator(production.RequestWithContext(Some(request), None)) + val results = production.sendCancellableRequestsWithContext(requests, 1, 1.nanosecond) + try { + assert(results.hasNext && sent.await(5, TimeUnit.SECONDS)) + // This handler never blocks; an idle pool worker has completed the queued Future, not just its response. + eventually(timeout(Span(5, Seconds))) { + assert(worker.get().getState == Thread.State.WAITING) + } + val result = results.next().response.get + assert(result.statusLine.statusCode == 200) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == "buffered") + assert(!results.hasNext) + worker.get().join(5000) + assert(!worker.get().isAlive) + } finally cancelRemaining(results) + } + } + + test("the minimum concurrent timeout still cancels pending HTTP work and drains its worker") { + val worker = new AtomicReference[Thread]() + val started = new CountDownLatch(1) + val blocked = new CountDownLatch(1) + val stopped = new CountDownLatch(1) + val handler: HandlingUtils.HandlerFunc = (_, _) => { + worker.set(Thread.currentThread()) + started.countDown() + try { + blocked.await() + HTTPSchema.stringToResponse("unexpected", 200, "OK") + } finally stopped.countDown() + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val requests = Iterator(production.RequestWithContext(Some(request), None)) + val results = production.sendCancellableRequestsWithContext(requests, 1, 1.nanosecond) + try { + assert(results.hasNext && started.await(5, TimeUnit.SECONDS)) + val error = intercept[TimeoutException](results.next()) + assert(error.getSuppressed.isEmpty && stopped.await(5, TimeUnit.SECONDS)) + worker.get().join(5000) + assert(!worker.get().isAlive) + } finally { + cancelRemaining(results) + blocked.countDown() + } + } + } + + Seq(false, true).foreach { interruptWorker => + val interruptedSource = if (interruptWorker) "retry workers" else "waiting callers" + test(s"cancellable HTTP preserves interruption from $interruptedSource") { + val worker = new AtomicReference[Thread]() + val failure = new AtomicReference[Throwable]() + val callerInterrupted = new AtomicBoolean() + val workerInterrupted = new AtomicBoolean() + val sends = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + throw new IOException("Synthetic connection failure") + })) { failingClient => + val handler: HandlingUtils.HandlerFunc = (_, request) => { + worker.set(Thread.currentThread()) + try HandlingUtils.bounded(60000)(failingClient, request) + finally workerInterrupted.set(Thread.currentThread().isInterrupted) + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val requests = Iterator(production.RequestWithContext(Some(request), None)) + val results = production.sendCancellableRequestsWithContext(requests, 1, 30.seconds) + val consumer = new Thread(() => { + try results.next() catch { + case error: Throwable => failure.set(error) + } finally callerInterrupted.set(Thread.currentThread().isInterrupted) + }) + try { + assert(results.hasNext) + eventually(timeout(Span(5, Seconds))) { + assert(sends.get() == 1 && worker.get() != null) + assert(worker.get().getStackTrace.exists(_.getMethodName == "sleep")) + } + consumer.start() + eventually(timeout(Span(5, Seconds))) { + assert(consumer.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.FutureTask" && frame.getMethodName == "get")) + assert(consumer.getState == Thread.State.TIMED_WAITING) + } + if (interruptWorker) worker.get().interrupt() else consumer.interrupt() + consumer.join(5000) + assert(!consumer.isAlive) + assert(failure.get().isInstanceOf[InterruptedException]) + assert(failure.get().getSuppressed.isEmpty, "HTTP worker cleanup failed") + assert(callerInterrupted.get(), "The consumer lost its cancellation flag") + assert(workerInterrupted.get(), "The bounded retry worker lost its cancellation flag") + worker.get().join(5000) + assert(!worker.get().isAlive && sends.get() == 1) + } finally { + consumer.interrupt() + consumer.join(5000) + cancelRemaining(results) + } + } + } + } + } + + test("cancellable HTTP preserves interruption while draining a worker after timeout") { + val release = new Semaphore(0) + val worker = new AtomicReference[Thread]() + val failure = new AtomicReference[Throwable]() + val interrupted = new AtomicBoolean() + val handler: HandlingUtils.HandlerFunc = (_, _) => { + worker.set(Thread.currentThread()) + // Model an uncooperative request so cancellation must wait for explicit quiescence. + release.acquireUninterruptibly() + HTTPSchema.stringToResponse("released", 200, "OK") + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val requests = Iterator(production.RequestWithContext(Some(request), None)) + val results = production.sendCancellableRequestsWithContext(requests, 1, 100.millis) + val consumer = new Thread(() => { + try results.next() catch { + case error: Throwable => failure.set(error) + } finally interrupted.set(Thread.currentThread().isInterrupted) + }) + try { + assert(results.hasNext) + eventually(timeout(Span(5, Seconds))) { + assert(worker.get() != null && worker.get().getState == Thread.State.WAITING) + } + consumer.start() + eventually(timeout(Span(5, Seconds))) { + assert(consumer.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.ThreadPoolExecutor" && + frame.getMethodName == "awaitTermination")) + } + consumer.interrupt() + consumer.join(5000) + assert(!consumer.isAlive && failure.get().isInstanceOf[TimeoutException]) + assert(failure.get().getSuppressed.exists(_.isInstanceOf[InterruptedException])) + assert(interrupted.get(), "Worker draining lost the caller's cancellation flag") + } finally { + release.release() + consumer.interrupt() + consumer.join(5000) + cancelRemaining(results) + } + worker.get().join(5000) + assert(!worker.get().isAlive) + } + } + + test("large finite concurrent timeouts work in legacy and cancellable HTTP transforms") { + val sends = new AtomicInteger() + withServer { exchange => + sends.incrementAndGet() + exchange.sendResponseHeaders(200, -1) + exchange.close() + } { address => + val request = HTTPRequestData.makeToRowConverter(new HTTPRequestData(new HttpPost(address))) + val data = spark.createDataFrame(Seq(Row(request)).asJava, + new StructType().add("request", HTTPRequestData.schema)) + Seq(1, 2).foreach { concurrency => + val stage = new HTTPTransformer().setInputCol("request").setOutputCol("response") + .setConcurrency(concurrency).setConcurrentTimeout(1e20) + Seq(stage.transform _, stage.transformCancellable _).foreach { transform => + val response = transform(data).head().getAs[Row]("response") + assert(response.getAs[Row]("statusLine").getAs[Int]("statusCode") == 200) + } + } + assert(sends.get() == 4) + } + } + + test("bounded policy respects Retry-After seconds and HTTP dates without retrying long hints early") { + val now = 1700000000000L + val date = org.apache.http.client.utils.DateUtils.formatDate(new java.util.Date(now + 2000)) + assert(HandlingUtils.retryAfterMillis("2", now).contains(2000L)) + assert(HandlingUtils.retryAfterMillis("1.0005", now).contains(1001L)) + assert(HandlingUtils.retryAfterMillis("0.0005", now).contains(1L)) + assert(HandlingUtils.retryAfterMillis("1e-1000000", now).contains(1L)) + assert(HandlingUtils.retryAfterMillis("0", now).contains(0L)) + assert(HandlingUtils.retryAfterMillis(date, now).contains(2000L)) + assert(HandlingUtils.retryAfterMillis("-1", now).isEmpty) + assert(HandlingUtils.retryAfterMillis("not a date", now).isEmpty) + Seq("120", "120.0005", "60.0005", "9223372036854776", + "9999999999999999999999999999999999999999").foreach { hint => + Seq(429, 503).foreach { code => + val sends = new AtomicInteger() + withServer { exchange => + sends.incrementAndGet() + exchange.getResponseHeaders.add("Retry-After", hint) + exchange.sendResponseHeaders(code, -1) + exchange.close() + } { address => + Using.resource(HttpClients.custom().disableAutomaticRetries().build()) { client => + val result = HandlingUtils.bounded(0, 0)(client, new HTTPRequestData(new HttpPost(address))) + assert(result.statusLine.statusCode == code) + assert(sends.get() == 1, s"Long Retry-After $hint must not trigger another send for HTTP $code") + } + } + } + } + } + + test("bounded fractional Retry-After waits until at least the next whole millisecond") { + Seq(429, 503).foreach { code => + val sends = new ConcurrentLinkedQueue[Long]() + withServer { exchange => + sends.add(System.nanoTime()) + val first = sends.size() == 1 + exchange.getResponseHeaders.add("Connection", "close") + if (first) exchange.getResponseHeaders.add("Retry-After", "1.0005") + exchange.sendResponseHeaders(if (first) code else 200, -1) + exchange.close() + } { address => + Using.resource(HttpClients.custom().disableAutomaticRetries().build()) { client => + Using.resource(HandlingUtils.sendWithRetryPolicy(client, new HttpPost(address), Array(0), + Set.empty, 0, boundedPolicy = true, quiet = false)) { result => + assert(result.getStatusLine.getStatusCode == 200) + } + val times = sends.asScala.toVector + assert(times.size == 2) + assert(times(1) - times.head >= TimeUnit.MILLISECONDS.toNanos(1001)) + } + } + } + } + + test("bounded connection failures consume exactly the configured retry budget") { + val sends = new AtomicInteger() + withServer { exchange => + sends.incrementAndGet() + exchange.close() + } { address => + Using.resource(HttpClients.custom().disableAutomaticRetries().build()) { client => + val result = HandlingUtils.bounded(0, 0, 0)(client, new HTTPRequestData(new HttpPost(address))) + assert(result == null) + assert(sends.get() == 4, "Recursive IOException propagation must not restart exhausted retry budgets") + } + } + } + + test("bounded rejected response read failures preserve status and never resend terminal errors") { + Seq(400, 401, 403, 404, 413, 429, 503).foreach { code => + Seq("body", "firstHeader", "headers").foreach { fault => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + val result = new BasicHttpResponse(new ProtocolVersion("HTTP", 1, 1), code, "Synthetic") + with CloseableHttpResponse { + override def getFirstHeader(name: String): Header = { + if (fault == "firstHeader") throw new IOException("Synthetic retry header failure") + super.getFirstHeader(name) + } + override def getAllHeaders: Array[Header] = { + if (fault == "headers") throw new IOException("Synthetic header failure") + super.getAllHeaders + } + override def close(): Unit = { closes.incrementAndGet(); () } + } + result.setEntity(truncatedEntity()) + result.setLocale(java.util.Locale.ROOT) + result + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + val expected = if (code == 429 || code == 503) 3 else 1 + assert(sends.get() == expected && closes.get() == expected, s"Unexpected retries for HTTP $code") + assert(result != null && result.statusLine.statusCode == code) + val diagnostic = new String(result.entity.get.content, StandardCharsets.UTF_8) + assert(diagnostic.contains("Could not read") && !diagnostic.contains("Synthetic body reset")) + } + } + } + } + + test("bounded truncated responses retain short Retry-After hints and share the retry budget") { + Seq(429, 503).foreach { code => + Seq(0, 1, 2).foreach { retries => + val sends = new ConcurrentLinkedQueue[Long]() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.add(System.nanoTime()) + val failing = sends.size() <= 2 + val result = response(if (failing) code else 200, + if (failing) truncatedEntity() else new StringEntity("complete"), + () => { closes.incrementAndGet(); () }) + if (failing) result.setHeader("Retry-After", "0.05") + result + })) { http => + val result = HandlingUtils.bounded(Array.fill(retries)(0): _*)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result != null && result.statusLine.statusCode == (if (retries == 2) 200 else code)) + val times = sends.asScala.toVector + assert(times.size == retries + 1 && closes.get() == retries + 1) + times.sliding(2).filter(_.size == 2).foreach { pair => + assert(pair(1) - pair.head >= TimeUnit.MILLISECONDS.toNanos(50)) + } + } + } + } + } + + test("bounded truncated responses never retry Retry-After hints above the cap") { + val future = org.apache.http.client.utils.DateUtils.formatDate( + new java.util.Date(System.currentTimeMillis() + 120000)) + Seq(429, 503).foreach { code => + Seq("120", "60.0005", future).foreach { hint => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + val result = response(code, truncatedEntity(), () => { closes.incrementAndGet(); () }) + result.setHeader("Retry-After", hint) + result + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(sends.get() == 1 && closes.get() == 1) + assert(result != null && result.statusLine.statusCode == code) + val diagnostic = new String(result.entity.get.content, StandardCharsets.UTF_8) + assert(diagnostic.contains("Could not read") && !diagnostic.contains("Synthetic body reset")) + } + } + } + } + + test("production cancellable clients own the complete write-failure retry budget") { + Seq(0, 2).foreach { retries => + val writes = new AtomicInteger() + withServer { exchange => + exchange.sendResponseHeaders(200, -1) + exchange.close() + } { address => + val handler: HandlingUtils.HandlerFunc = (http, data) => { + val request = new HttpPost(data.requestLine.uri) + request.setEntity(new ByteArrayEntity(Array[Byte](1)) { + override def writeTo(output: OutputStream): Unit = { + writes.incrementAndGet() + throw new SocketException("Synthetic request write failure") + } + }) + try HandlingUtils.convertAndClose(HandlingUtils.sendWithRetryPolicy( + http, request, Array.fill(retries)(0), Set.empty, 0, boundedPolicy = true, quiet = true)) + finally request.releaseConnection() + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val data = new HTTPRequestData(new HttpPost(address)) + def requests: Iterator[production.RequestWithContext] = + Iterator(production.RequestWithContext(Some(data), None)) + intercept[IOException] { + production.sendCancellableRequestsWithContext(requests, 1, 5.seconds).toVector + } + assert(writes.get() == retries + 1, "Apache retries must not multiply the bounded send budget") + writes.set(0) + intercept[IOException](production.sendRequestsWithContext(requests).toVector) + assert(writes.get() == (retries + 1) * 4, "Legacy Apache retry behavior must remain unchanged") + } + } + } + } + + test("bounded metadata failures preserve buffered bodies and capacity retry classification") { + Seq("encoding", "type", "chunked").foreach { fault => + Seq(429 -> "CapacityLimitExceeded", 429 -> "transient", 503 -> "unavailable").foreach { case (code, body) => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + val entity = new StringEntity(body) { + override def getContentEncoding: Header = + if (fault == "encoding") throw new IOException("Synthetic metadata failure") + else super.getContentEncoding + override def getContentType: Header = + if (fault == "type") throw new IOException("Synthetic metadata failure") + else super.getContentType + override def isChunked: Boolean = + if (fault == "chunked") throw new IOException("Synthetic metadata failure") + else super.isChunked + } + val result = response(code, entity, () => { closes.incrementAndGet(); () }) + result.setHeader("Retry-After", "0") + result + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result != null && result.statusLine.statusCode == code) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == body) + val expectedSends = if (body == "CapacityLimitExceeded") 1 else 3 + assert(sends.get() == expectedSends && closes.get() == expectedSends) + } + } + } + } + + test("production cancellable clients do not follow cross-origin redirects or forward credentials") { + Seq(0, 2).foreach { retries => + val sends = new AtomicInteger() + val targetHeaders = new ConcurrentLinkedQueue[Map[String, String]]() + withServer { exchange => + targetHeaders.add(exchange.getRequestHeaders.asScala.map { case (name, values) => + name.toLowerCase(java.util.Locale.ROOT) -> values.asScala.mkString(",") + }.toMap) + exchange.sendResponseHeaders(200, -1) + exchange.close() + } { destination => + withServer { exchange => + sends.incrementAndGet() + Using.resource(exchange.getRequestBody)(IOUtils.toByteArray) + exchange.getResponseHeaders.add("Location", destination) + exchange.sendResponseHeaders(303, -1) + exchange.close() + } { address => + val handler: HandlingUtils.HandlerFunc = (http, request) => + HandlingUtils.bounded(Array.fill(retries)(0): _*)(http, request) + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val post = new HttpPost(address) + post.setHeader("api-key", "synthetic-redirect-key") + post.setHeader("Authorization", "Bearer synthetic-redirect-token") + post.setHeader("X-Custom-Auth", "synthetic-custom-auth") + val data = new HTTPRequestData(post) + def requests: Iterator[production.RequestWithContext] = + Iterator(production.RequestWithContext(Some(data), None)) + Seq(1, 2).foreach { concurrency => + val before = sends.get() + val result = production.sendCancellableRequestsWithContext(requests, concurrency, 5.seconds) + .toVector.head.response.get + assert(result != null, s"Expected redirect response: retries=$retries, " + + s"concurrency=$concurrency, sends=${sends.get() - before}") + assert(targetHeaders.isEmpty, "Redirect destinations must receive no requests or credentials") + assert(result.statusLine.statusCode == 303) + assert(sends.get() - before == 1, "Redirect responses must not repeat the original request") + } + val legacy = production.sendRequestsWithContext(requests).toVector.head.response.get + assert(legacy.statusLine.statusCode == 200 && targetHeaders.size() == 1, + "Legacy automatic redirect behavior must remain unchanged") + } + } + } + } + } + + test("bounded capacity metadata failures remain terminal without a usable Retry-After") { + Seq(None, Some("unusable")).foreach { hint => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + val result = response(429, new StringEntity("CapacityLimitExceeded") { + override def getContentType: Header = throw new IOException("Synthetic metadata failure") + }, () => { closes.incrementAndGet(); () }) + hint.foreach(result.setHeader("Retry-After", _)) + result + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result != null && result.statusLine.statusCode == 429) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == "CapacityLimitExceeded") + assert(sends.get() == 1 && closes.get() == 1) + } + } + } + + test("bounded policy returns every redirect status without repeating accepted requests") { + Seq(300, 301, 302, 303, 304, 305, 306, 307, 308, 399).foreach { code => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(code, new StringEntity("see other"), () => { closes.incrementAndGet(); () }) + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result.statusLine.statusCode == code) + assert(sends.get() == 1 && closes.get() == 1, s"HTTP $code must not replay the original POST") + } + } + } + + test("bounded redirect read failures preserve terminal status without replaying") { + Seq(300, 303, 399).foreach { code => + Seq("body", "headers", "metadata").foreach { fault => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + val result = new BasicHttpResponse(new ProtocolVersion("HTTP", 1, 1), code, "Synthetic") + with CloseableHttpResponse { + override def getAllHeaders: Array[Header] = { + if (fault == "headers") throw new IOException("Synthetic header failure") + super.getAllHeaders + } + override def close(): Unit = { closes.incrementAndGet(); () } + } + result.setEntity(if (fault == "body") truncatedEntity() else new StringEntity("see other") { + override def getContentType: Header = throw new IOException("Synthetic metadata failure") + }) + result.setLocale(java.util.Locale.ROOT) + result + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, + new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(sends.get() == 1 && closes.get() == 1, s"HTTP $code with $fault failure must not replay") + assert(result != null && result.statusLine.statusCode == code) + val diagnostic = new String(result.entity.get.content, StandardCharsets.UTF_8) + if (fault == "metadata") assert(diagnostic == "see other") + else assert(diagnostic.contains("Could not read") && !diagnostic.contains("Synthetic")) + } + } + } + } + + test("bounded policy accepts every successful HTTP status without repeating completed requests") { + Seq(203, 204, 206, 299).foreach { code => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(code, new StringEntity("completed"), () => { closes.incrementAndGet(); () }) + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result.statusLine.statusCode == code) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == "completed") + assert(sends.get() == 1 && closes.get() == 1) + } + } + } + + test("bounded policy propagates cancellation before performing HTTP") { + Using.resource(HttpClients.createDefault()) { client => + try { + Thread.currentThread().interrupt() + intercept[InterruptedException] { + HandlingUtils.bounded(0)(client, new HTTPRequestData(new HttpPost("http://127.0.0.1:1/cancelled"))) + } + } finally Thread.interrupted() + } + } + + test("bounded Fabric refresh consumes the same retry budget and never refreshes twice") { + val sends = new AtomicInteger() + val refreshes = new AtomicInteger() + val auth = new ConcurrentLinkedQueue[String]() + withServer { exchange => + auth.add(exchange.getRequestHeaders.getFirst("Authorization")) + sends.incrementAndGet() + exchange.sendResponseHeaders(401, -1) + exchange.close() + } { address => + Using.resource(HttpClients.custom().disableAutomaticRetries().build()) { client => + val request = new HttpPost(address) + request.setHeader("Authorization", "MwcToken stale") + try { + val response = HandlingUtils.sendWithRetryPolicy(client, request, Array(0, 0, 0), Set.empty, 0, + boundedPolicy = true, quiet = true, refreshAuth = Some { rejected => + assert(rejected == "MwcToken stale") + refreshes.incrementAndGet() + "MwcToken fresh" + }) + try assert(response.getStatusLine.getStatusCode == 401) finally response.close() + assert(sends.get() == 2 && refreshes.get() == 1) + assert(auth.asScala.toSeq == Seq("MwcToken stale", "MwcToken fresh")) + } finally request.releaseConnection() + } + } + } + + test("bounded body failures share the status retry budget and preserve the eventual response") { + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + val attempt = sends.incrementAndGet() + val entity = if (attempt == 2) truncatedEntity() else new StringEntity("complete") + response(if (attempt == 1) 503 else 200, entity, () => { closes.incrementAndGet(); () }) + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result.statusLine.statusCode == 200) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == "complete") + assert(sends.get() == 3 && closes.get() == 3) + } + + } + + test("bounded Fabric refresh failures retain the terminal rejection without replaying") { + Seq(new NoSuchMethodException("Synthetic missing token method"), + new IllegalStateException("Synthetic token failure"), new IOException("Synthetic refresh failure")) + .foreach { failure => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(401, new StringEntity("rejected"), () => { closes.incrementAndGet(); () }) + })) { http => + val request = new HttpPost("http://example.invalid/test") + try { + val result = HandlingUtils.convertAndClose(HandlingUtils.sendWithRetryPolicy( + http, request, Array(0, 0), Set.empty, 0, boundedPolicy = true, quiet = true, + refreshAuth = Some(_ => throw failure))) + assert(result.statusLine.statusCode == 401) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == "rejected") + assert(sends.get() == 1 && closes.get() == 1) + } finally request.releaseConnection() + } + } + } + + test("bounded Fabric refresh preserves interruption and cancellation") { + Seq(new InterruptedException("Synthetic interruption"), new CancellationException("Synthetic cancellation")) + .foreach { failure => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(401, new StringEntity("rejected"), () => { closes.incrementAndGet(); () }) + })) { http => + val request = new HttpPost("http://example.invalid/test") + try { + val observed = intercept[Exception] { + HandlingUtils.sendWithRetryPolicy(http, request, Array(0, 0), Set.empty, 0, + boundedPolicy = true, quiet = true, refreshAuth = Some(_ => throw failure)) + } + assert(observed eq failure) + assert(sends.get() == 1 && closes.get() == 1) + } finally request.releaseConnection() + } + } + } + + test("bounded exhausted body resets return transport failure rather than an empty successful response") { + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(200, truncatedEntity(), () => { closes.incrementAndGet(); () }) + })) { http => + val result = HandlingUtils.bounded(0, 0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result == null && sends.get() == 3 && closes.get() == 3) + } + } + + test("bounded response metadata IO failures close the attempt and consume one retry") { + Seq("status", "headers").foreach { fault => + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + val attempt = sends.incrementAndGet() + val result = new BasicHttpResponse(new ProtocolVersion("HTTP", 1, 1), 200, "Synthetic") + with CloseableHttpResponse { + override def getStatusLine: StatusLine = { + if (attempt == 1 && fault == "status") throw new IOException("Synthetic status failure") + super.getStatusLine + } + override def getAllHeaders: Array[Header] = { + if (attempt == 1 && fault == "headers") throw new IOException("Synthetic header failure") + super.getAllHeaders + } + override def close(): Unit = { closes.incrementAndGet(); () } + } + result.setEntity(new StringEntity("complete")) + result.setLocale(java.util.Locale.ROOT) + result + })) { http => + val result = HandlingUtils.bounded(0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result != null && new String(result.entity.get.content, StandardCharsets.UTF_8) == "complete") + assert(sends.get() == 2 && closes.get() == 2) + } + } + } + + test("bounded close IO failure does not discard a fully materialized answer or repeat the request") { + val sends = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(200, new StringEntity("complete"), () => { + closes.incrementAndGet() + throw new IOException("Synthetic close failure") + }) + })) { http => + val result = HandlingUtils.bounded(0)(http, new HTTPRequestData(new HttpPost("http://example.invalid/test"))) + assert(result != null && new String(result.entity.get.content, StandardCharsets.UTF_8) == "complete") + assert(sends.get() == 1 && closes.get() == 1) + } + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala new file mode 100644 index 00000000000..b21cdb2aefb --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala @@ -0,0 +1,682 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.io.split1 + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HTTPSchema, HandlingUtils, SingleThreadedHTTPClient} +import org.apache.http.client.methods.HttpPost +import org.apache.http.entity.StringEntity +import org.apache.spark.{TaskContextTestSupport, TaskKilledException} +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} + +import java.net.URLClassLoader +import java.util.concurrent.{CancellationException, ConcurrentHashMap, CountDownLatch, ExecutionException, Semaphore, + TimeUnit, TimeoutException} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference} +import scala.collection.JavaConverters._ +import scala.concurrent.duration._ +import scala.util.Using + +class VerifyCancellableRequests extends TestBase { + test("cancellation infrastructure does not retain the task classloader while request workers inherit it") { + val previousLoader = Thread.currentThread().getContextClassLoader + val libraryLoader = classOf[SingleThreadedHTTPClient].getClassLoader + Using.resource(new URLClassLoader(Array.empty, previousLoader)) { taskLoader => + val monitorLoader = new AtomicReference[ClassLoader]() + val workerLoader = new AtomicReference[ClassLoader]() + val abortLoader = new AtomicReference[ClassLoader]() + val monitorChecked = new CountDownLatch(1) + val workerStarted = new CountDownLatch(1) + val release = new Semaphore(0) + val context = new TaskContextTestSupport.ControlledContext(() => { + monitorLoader.set(Thread.currentThread().getContextClassLoader) + monitorChecked.countDown() + }) + val post = new HttpPost("http://example.invalid/test") { + override def abort(): Unit = { + abortLoader.set(Thread.currentThread().getContextClassLoader) + release.release() + super.abort() + } + } + Thread.currentThread().setContextClassLoader(taskLoader) + try { + Using.resource(VerifyBoundedRequests.client(() => { + workerLoader.set(Thread.currentThread().getContextClassLoader) + workerStarted.countDown() + release.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val handler: HandlingUtils.HandlerFunc = (_, data) => HandlingUtils.bounded()(http, data) + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val results = TaskContextTestSupport.withContext(context) { + production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(request), None)), 1, 30.seconds) + } + try { + assert(results.hasNext && workerStarted.await(5, TimeUnit.SECONDS)) + assert(monitorChecked.await(5, TimeUnit.SECONDS)) + results.close() + assert(monitorLoader.get() eq libraryLoader) + assert(abortLoader.get() eq libraryLoader) + assert(workerLoader.get() eq taskLoader) + } finally { + release.release() + try results.close() finally context.complete() + } + } + } + } finally Thread.currentThread().setContextClassLoader(previousLoader) + } + } + + test("task cancellation during worker submission preserves the recorded kill reason") { + val cancellationObserved = new AtomicBoolean() + val poolShutdown = new CountDownLatch(1) + val sent = new AtomicInteger() + val context = new TaskContextTestSupport.ControlledContext(() => ()) { + override def isInterrupted(): Boolean = { + val interrupted = super.isInterrupted() + if (interrupted && Thread.currentThread().getName == "synapseml-http-cancellation-monitor") { + cancellationObserved.set(true) + } + interrupted + } + } + val observer = new TaskContextTestSupport.ControlledContext(() => { + if (cancellationObserved.get()) poolShutdown.countDown() + }) + val cancelOnWorkerCreation = new InheritableThreadLocal[Boolean]() { + override def childValue(value: Boolean): Boolean = { + context.interrupt() + // The next monitor watch runs after this scope's shutdown, while submission is still in progress. + assert(poolShutdown.await(5, TimeUnit.SECONDS)) + false + } + } + val handler: HandlingUtils.HandlerFunc = (_, _) => { + sent.incrementAndGet() + HTTPSchema.stringToResponse("unused", 200, "OK") + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val results = TaskContextTestSupport.withContext(context) { + production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(request), None)), 1, 30.seconds) + } + val observation = TaskContextTestSupport.withContext(observer) { + production.sendCancellableRequestsWithContext(Iterator.empty, 1, 30.seconds) + } + try { + cancelOnWorkerCreation.set(true) + val failure = intercept[TaskKilledException](results.hasNext) + assert(failure.reason == "Synthetic task cancellation") + assert(sent.get() == 0 && !Thread.currentThread().isInterrupted) + } finally { + cancelOnWorkerCreation.remove() + try results.close() finally { + try observation.close() finally { + try context.complete() finally observer.complete() + } + } + } + } + } + + test("expired HTTP waits preserve a task kill that races with the initial cancellation check") { + val consumer = Thread.currentThread() + val blockMonitor = new AtomicBoolean() + val cancelAtAwait = new AtomicBoolean() + val monitorEntered = new CountDownLatch(1) + val monitorRelease = new CountDownLatch(1) + val started = new CountDownLatch(1) + val release = new CountDownLatch(1) + val worker = new AtomicReference[Thread]() + val context = new TaskContextTestSupport.ControlledContext(() => { + if (blockMonitor.get()) { + monitorEntered.countDown() + assert(monitorRelease.await(10, TimeUnit.SECONDS)) + } + }) { + override def isInterrupted(): Boolean = { + val interrupted = super.isInterrupted() + if ((Thread.currentThread() eq consumer) && cancelAtAwait.get() && + consumer.getStackTrace.exists(frame => + frame.getClassName.endsWith("CancellableRequests$Scope") && frame.getMethodName == "await")) { + // Publish the kill after the initial read, returning the value observed before it. + if (cancelAtAwait.compareAndSet(true, false)) interrupt() + } + interrupted + } + } + val handler: HandlingUtils.HandlerFunc = (_, _) => { + worker.set(Thread.currentThread()) + started.countDown() + release.await() + HTTPSchema.stringToResponse("released", 200, "OK") + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val results = TaskContextTestSupport.withContext(context) { + production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(request), None)), 1, 1.nanosecond) + } + try { + assert(results.hasNext && started.await(5, TimeUnit.SECONDS)) + blockMonitor.set(true) + assert(monitorEntered.await(5, TimeUnit.SECONDS)) + cancelAtAwait.set(true) + val failure = intercept[TaskKilledException](results.next()) + assert(failure.reason == "Synthetic task cancellation") + assert(!cancelAtAwait.get() && !Thread.currentThread().isInterrupted) + } finally { + monitorRelease.countDown() + release.countDown() + try results.close() finally { + try context.complete() finally { + Option(worker.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + } + } + } + + test("driver response iterators can close before submitting requests") { + val sent = new AtomicInteger() + val handler: HandlingUtils.HandlerFunc = (_, _) => { + sent.incrementAndGet() + HTTPSchema.stringToResponse("unused", 200, "OK") + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val data = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val results = production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(data), None)), 1, 30.seconds) + try { + results.close() + results.close() + assert(!results.hasNext && sent.get() == 0) + intercept[NoSuchElementException](results.next()) + } finally results.toVector + } + } + + test("driver response iterators close abandoned work and drain every owned worker") { + val started = new CountDownLatch(1) + val release = new Semaphore(0) + val workers = ConcurrentHashMap.newKeySet[Thread]() + val abortWorker = new AtomicReference[Thread]() + val aborts = new AtomicInteger() + val post = new HttpPost("http://example.invalid/blocked") { + override def abort(): Unit = { + abortWorker.set(Thread.currentThread()) + aborts.incrementAndGet() + release.release() + super.abort() + } + } + Using.resource(VerifyBoundedRequests.client(() => { + started.countDown() + release.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val handler: HandlingUtils.HandlerFunc = (_, data) => { + workers.add(Thread.currentThread()) + if (data.requestLine.uri.endsWith("/first")) HTTPSchema.stringToResponse("first", 200, "OK") + else HandlingUtils.bounded()(http, data) + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val first = new HTTPRequestData(new HttpPost("http://example.invalid/first")) + val blocked = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val results = production.sendCancellableRequestsWithContext( + Iterator(first, blocked).map(data => production.RequestWithContext(Some(data), None)), 2, 30.seconds) + try { + assert(results.next().response.nonEmpty && started.await(5, TimeUnit.SECONDS)) + results.close() + results.close() + assert(!results.hasNext && aborts.get() == 1) + workers.asScala.foreach(_.join(5000)) + assert(workers.size() == 2 && workers.asScala.forall(!_.isAlive)) + abortWorker.get().join(5000) + assert(!abortWorker.get().isAlive) + } finally { + release.release() + results.toVector + workers.asScala.foreach { worker => + worker.join(5000) + assert(!worker.isAlive) + } + } + } + } + } + + Seq(false, true).foreach { cooperative => + test(s"a logged abort failure only fails task cleanup while HTTP workers remain active: $cooperative") { + val context = new TaskContextTestSupport.ControlledContext(() => ()) + val started = new CountDownLatch(1) + val release = new Semaphore(0) + val worker = new AtomicReference[Thread]() + val abortWorker = new AtomicReference[Thread]() + val aborts = new AtomicInteger() + val post = new HttpPost("http://example.invalid/test") { + override def abort(): Unit = { + abortWorker.set(Thread.currentThread()) + aborts.incrementAndGet() + if (cooperative) release.release() + throw new IllegalStateException("Synthetic abort failure") + } + } + Using.resource(VerifyBoundedRequests.client(() => { + worker.set(Thread.currentThread()) + started.countDown() + release.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val handler: HandlingUtils.HandlerFunc = (_, data) => HandlingUtils.bounded()(http, data) + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val data = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val results = TaskContextTestSupport.withContext(context) { + production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(data), None)), 1, 30.seconds) + } + try { + assert(results.hasNext && started.await(5, TimeUnit.SECONDS)) + val before = System.nanoTime() + if (cooperative) { + context.complete() + worker.get().join(5000) + assert(!worker.get().isAlive) + } else { + val error = intercept[Exception](context.complete()) + assert(error.getMessage.contains("HTTP request workers did not stop")) + assert(worker.get().isAlive) + } + assert((System.nanoTime() - before).nanos < 8.seconds) + assert(aborts.get() == 1) + intercept[CancellationException](results.hasNext) + } finally { + release.release() + Seq(worker, abortWorker).foreach { reference => + Option(reference.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + } + } + } + } + + test("an unresponsive abort fails bounded cleanup without starting another abort pass") { + val started = new CountDownLatch(1) + val release = new Semaphore(0) + val abortRelease = new CountDownLatch(1) + val aborts = new AtomicInteger() + val worker = new AtomicReference[Thread]() + val abortWorker = new AtomicReference[Thread]() + val post = new HttpPost("http://example.invalid/test") { + override def abort(): Unit = { + abortWorker.set(Thread.currentThread()) + aborts.incrementAndGet() + assert(abortRelease.await(10, TimeUnit.SECONDS)) + super.abort() + } + } + Using.resource(VerifyBoundedRequests.client(() => { + worker.set(Thread.currentThread()) + started.countDown() + release.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val handler: HandlingUtils.HandlerFunc = (_, data) => HandlingUtils.bounded()(http, data) + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val data = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val results = production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(data), None)), 1, 1.nanosecond) + try { + assert(results.hasNext && started.await(5, TimeUnit.SECONDS)) + val before = System.nanoTime() + val error = intercept[TimeoutException](results.next()) + assert(error.getSuppressed.exists(_.getMessage.contains("HTTP request abort did not stop"))) + assert((System.nanoTime() - before).nanos < 8.seconds) + intercept[CancellationException](results.hasNext) + assert(aborts.get() == 1) + eventually(timeout(Span(2, Seconds))) { + assert(!abortWorker.get().isAlive, "Timed-out abort did not receive an interrupt") + } + } finally { + abortRelease.countDown() + release.release() + Seq(worker, abortWorker).foreach { reference => + Option(reference.get()).filterNot(_ == Thread.currentThread()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + } + } + } + + Seq(false, true).foreach { interruptCaller => + val failureSource = if (interruptCaller) "abort wait interruption" else "exceptional abort completion" + test(s"cleanup still drains HTTP workers after $failureSource") { + val started = new CountDownLatch(1) + val release = new Semaphore(0) + val abortStarted = new CountDownLatch(1) + val abortRelease = new CountDownLatch(1) + val worker = new AtomicReference[Thread]() + val abortWorker = new AtomicReference[Thread]() + val failure = new AtomicReference[Throwable]() + val callerInterrupted = new AtomicBoolean() + val post = new HttpPost("http://example.invalid/test") { + override def abort(): Unit = { + abortWorker.set(Thread.currentThread()) + abortStarted.countDown() + if (interruptCaller) { + assert(abortRelease.await(10, TimeUnit.SECONDS)) + super.abort() + } else { + throw new InterruptedException("Synthetic abort failure") + } + } + } + Using.resource(VerifyBoundedRequests.client(() => { + worker.set(Thread.currentThread()) + started.countDown() + release.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val handler: HandlingUtils.HandlerFunc = (_, data) => HandlingUtils.bounded()(http, data) + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val data = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val results = production.sendCancellableRequestsWithContext( + Iterator(production.RequestWithContext(Some(data), None)), 1, 1.nanosecond) + val consumer = new Thread(() => { + try results.next() catch { + case error: Throwable => failure.set(error) + } finally callerInterrupted.set(Thread.currentThread().isInterrupted) + }) + try { + assert(results.hasNext && started.await(5, TimeUnit.SECONDS)) + consumer.start() + assert(abortStarted.await(5, TimeUnit.SECONDS)) + if (interruptCaller) { + eventually(timeout(Span(2, Seconds))) { + assert(consumer.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.FutureTask" && frame.getMethodName == "get")) + } + consumer.interrupt() + } + eventually(timeout(Span(2, Seconds))) { + assert(consumer.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.ThreadPoolExecutor" && + frame.getMethodName == "awaitTermination"), "Abort-wait failure skipped the worker drain") + assert(!abortWorker.get().isAlive, "Failed abort wait left the owned abort thread running") + } + } finally { + abortRelease.countDown() + release.release() + consumer.join(10000) + Seq(worker, abortWorker).foreach { reference => + Option(reference.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + assert(!consumer.isAlive && failure.get().isInstanceOf[TimeoutException]) + val cleanup = failure.get().getSuppressed + assert(cleanup.length == 1 && cleanup.head.getSuppressed.isEmpty) + if (interruptCaller) { + assert(cleanup.head.isInstanceOf[InterruptedException] && callerInterrupted.get()) + } else { + assert(cleanup.head.isInstanceOf[ExecutionException] && !callerInterrupted.get()) + } + intercept[CancellationException](results.hasNext) + } + } + } + } + + test("one blocking request abort cannot stall cancellation of another task") { + val monitorEntered = new CountDownLatch(1) + val monitorRelease = new CountDownLatch(1) + val armed = new AtomicBoolean() + val firstContext = new TaskContextTestSupport.ControlledContext(() => { + if (armed.compareAndSet(true, false)) { + monitorEntered.countDown() + assert(monitorRelease.await(10, TimeUnit.SECONDS)) + } + }) + val secondContext = new TaskContextTestSupport.ControlledContext(() => ()) + val abortEntered = new CountDownLatch(1) + val abortRelease = new CountDownLatch(1) + val firstStarted = new CountDownLatch(1) + val secondStarted = new CountDownLatch(1) + val secondStopped = new CountDownLatch(1) + val workerRelease = new Semaphore(0) + val firstWorker = new AtomicReference[Thread]() + val secondWorker = new AtomicReference[Thread]() + val failure = new AtomicReference[Throwable]() + val post = new HttpPost("http://example.invalid/test") { + override def abort(): Unit = { + abortEntered.countDown() + assert(abortRelease.await(10, TimeUnit.SECONDS)) + super.abort() + } + } + Using.resource(VerifyBoundedRequests.client(() => { + firstWorker.set(Thread.currentThread()) + firstStarted.countDown() + workerRelease.acquireUninterruptibly() + VerifyBoundedRequests.response(200, new StringEntity("released"), () => ()) + })) { http => + val firstHandler: HandlingUtils.HandlerFunc = (_, data) => HandlingUtils.bounded()(http, data) + val secondHandler: HandlingUtils.HandlerFunc = (_, _) => { + secondWorker.set(Thread.currentThread()) + secondStarted.countDown() + try { + new CountDownLatch(1).await() + HTTPSchema.stringToResponse("released", 200, "OK") + } finally secondStopped.countDown() + } + Using.resource(new SingleThreadedHTTPClient(firstHandler, 1000)) { first => + Using.resource(new SingleThreadedHTTPClient(secondHandler, 1000)) { second => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val blockingRequest = new HTTPRequestData(post) { + override def toHTTPCore: HttpPost = post + } + val firstResults = TaskContextTestSupport.withContext(firstContext) { + first.sendCancellableRequestsWithContext(Iterator(first.RequestWithContext(Some(blockingRequest), None)), + 1, 30.seconds) + } + val secondResults = TaskContextTestSupport.withContext(secondContext) { + second.sendCancellableRequestsWithContext(Iterator(second.RequestWithContext(Some(request), None)), + 1, 30.seconds) + } + val consumer = new Thread(() => { + try firstResults.next() catch { + case error: Throwable => failure.set(error) + } + }) + try { + assert(firstResults.hasNext && secondResults.hasNext) + assert(firstStarted.await(5, TimeUnit.SECONDS) && secondStarted.await(5, TimeUnit.SECONDS)) + armed.set(true) + firstContext.interrupt() + assert(monitorEntered.await(5, TimeUnit.SECONDS)) + secondContext.interrupt() + monitorRelease.countDown() + consumer.start() + assert(abortEntered.await(5, TimeUnit.SECONDS)) + assert(secondStopped.await(2, TimeUnit.SECONDS), "Shared cancellation monitor waited for request abort") + assert(consumer.isAlive, "The first scope must still be blocked in request abort") + } finally { + monitorRelease.countDown() + abortRelease.countDown() + workerRelease.release() + consumer.join(10000) + firstContext.complete() + secondContext.complete() + Seq(firstWorker, secondWorker).foreach { worker => + Option(worker.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + assert(!consumer.isAlive) + assert(Option(failure.get()).collect { + case error: TaskKilledException => error.reason + }.contains("Synthetic task cancellation")) + } + } + } + } + + test("one draining scope cannot block cancellation of another task on the shared monitor") { + val monitorEntered = new CountDownLatch(1) + val monitorRelease = new CountDownLatch(1) + val armed = new AtomicBoolean() + val firstContext = new TaskContextTestSupport.ControlledContext(() => { + if (armed.compareAndSet(true, false)) { + monitorEntered.countDown() + assert(monitorRelease.await(10, TimeUnit.SECONDS)) + } + }) + val secondContext = new TaskContextTestSupport.ControlledContext(() => ()) + val firstStarted = new CountDownLatch(1) + val secondStarted = new CountDownLatch(1) + val secondStopped = new CountDownLatch(1) + val release = new Semaphore(0) + val firstWorker = new AtomicReference[Thread]() + val secondWorker = new AtomicReference[Thread]() + val failure = new AtomicReference[Throwable]() + val firstHandler: HandlingUtils.HandlerFunc = (_, _) => { + firstWorker.set(Thread.currentThread()) + firstStarted.countDown() + release.acquireUninterruptibly() + HTTPSchema.stringToResponse("released", 200, "OK") + } + val secondHandler: HandlingUtils.HandlerFunc = (_, _) => { + secondWorker.set(Thread.currentThread()) + secondStarted.countDown() + try { + new CountDownLatch(1).await() + HTTPSchema.stringToResponse("released", 200, "OK") + } finally secondStopped.countDown() + } + Using.resource(new SingleThreadedHTTPClient(firstHandler, 1000)) { first => + Using.resource(new SingleThreadedHTTPClient(secondHandler, 1000)) { second => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val firstResults = TaskContextTestSupport.withContext(firstContext) { + first.sendCancellableRequestsWithContext(Iterator(first.RequestWithContext(Some(request), None)), + 1, 30.seconds) + } + val secondResults = TaskContextTestSupport.withContext(secondContext) { + second.sendCancellableRequestsWithContext(Iterator(second.RequestWithContext(Some(request), None)), + 1, 30.seconds) + } + val consumer = new Thread(() => { + try firstResults.next() catch { + case error: Throwable => failure.set(error) + } + }) + try { + assert(firstResults.hasNext && secondResults.hasNext) + assert(firstStarted.await(5, TimeUnit.SECONDS) && secondStarted.await(5, TimeUnit.SECONDS)) + armed.set(true) + firstContext.interrupt() + assert(monitorEntered.await(5, TimeUnit.SECONDS)) + consumer.start() + eventually(timeout(Span(5, Seconds))) { + assert(consumer.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.ThreadPoolExecutor" && + frame.getMethodName == "awaitTermination")) + } + secondContext.interrupt() + monitorRelease.countDown() + assert(secondStopped.await(2, TimeUnit.SECONDS), "Shared cancellation monitor waited for another scope") + assert(consumer.isAlive, "The first scope must still be draining its uncooperative worker") + } finally { + monitorRelease.countDown() + release.release() + consumer.join(10000) + firstContext.complete() + secondContext.complete() + Seq(firstWorker, secondWorker).foreach { worker => + Option(worker.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + assert(!consumer.isAlive) + assert(Option(failure.get()).collect { + case error: TaskKilledException => error.reason + }.contains("Synthetic task cancellation")) + } + } + } + + test("nested iterator failures drain uncooperative HTTP workers only once") { + val started = new CountDownLatch(1) + val release = new Semaphore(0) + val worker = new AtomicReference[Thread]() + val original = new IllegalStateException("Synthetic input failure") + val handler: HandlingUtils.HandlerFunc = (_, _) => { + worker.set(Thread.currentThread()) + started.countDown() + release.acquireUninterruptibly() + HTTPSchema.stringToResponse("released", 200, "OK") + } + Using.resource(new SingleThreadedHTTPClient(handler, 1000)) { production => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val first = production.RequestWithContext(Some(request), None) + val requests = Iterator(0, 1).map { index => + if (index == 1) { + assert(started.await(5, TimeUnit.SECONDS)) + throw original + } + first + } + val results = production.sendCancellableRequestsWithContext(requests, 2, 30.seconds) + try { + val before = System.nanoTime() + val error = intercept[IllegalStateException](results.next()) + val elapsed = (System.nanoTime() - before).nanos + assert(error eq original) + assert(error.getSuppressed.length == 1) + assert(error.getSuppressed.head.getMessage.contains("request quiescence is not confirmed")) + assert(elapsed < 8.seconds, s"Cleanup repeated its worker drain: $elapsed") + val retry = System.nanoTime() + val cancelled = intercept[CancellationException](results.hasNext) + assert(cancelled.getSuppressed.isEmpty && (System.nanoTime() - retry).nanos < 2.seconds) + } finally { + release.release() + Option(worker.get()).foreach { thread => + thread.join(5000) + assert(!thread.isAlive) + } + } + } + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyFabricRetryAfter.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyFabricRetryAfter.scala new file mode 100644 index 00000000000..fb4e67a5797 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyFabricRetryAfter.scala @@ -0,0 +1,274 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.io.split1 + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.io.http.{HTTPRequestData, HTTPResponseData, HTTPTransformer, HandlingUtils} +import com.sun.net.httpserver.{HttpExchange, HttpServer} +import org.apache.commons.io.IOUtils +import org.apache.http.client.methods.{HttpGet, HttpPost} +import org.apache.http.client.utils.DateUtils +import org.apache.http.entity.StringEntity +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.StructType +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} + +import java.io.IOException +import java.net.InetSocketAddress +import java.nio.charset.StandardCharsets +import java.util.concurrent.{ConcurrentLinkedQueue, Executors, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} +import scala.collection.JavaConverters._ +import scala.util.Using + +class VerifyFabricRetryAfter extends TestBase { + import VerifyBoundedRequests.{client, response} + + test("bounded responses reject oversized Retry-After hints before expensive numeric parsing") { + val sends = new AtomicInteger() + val closes = new AtomicInteger() + val hint = "9" * (1024 * 1024) + Using.resource(client(() => { + val first = sends.incrementAndGet() == 1 + val result = response(if (first) 503 else 200, new StringEntity("result"), + () => { closes.incrementAndGet(); () }) + if (first) result.setHeader("Retry-After", hint) + result + })) { http => + val request = new HTTPRequestData(new HttpPost("http://example.invalid/test")) + val start = System.nanoTime() + val result = HandlingUtils.bounded(25)(http, request) + val elapsed = System.nanoTime() - start + assert(result.statusLine.statusCode == 200 && sends.get() == 2 && closes.get() == 2) + assert(elapsed >= TimeUnit.MILLISECONDS.toNanos(25) && elapsed < TimeUnit.SECONDS.toNanos(5)) + assert(HandlingUtils.retryAfterMillis("0." + ("0" * 61) + "1", 0).contains(1L)) + assert(HandlingUtils.retryAfterMillis("0." + ("0" * 62) + "1", 0).isEmpty) + } + } + + test("authentication callbacks require the bounded retry policy before any request is sent") { + val sends = new AtomicInteger() + val refreshes = new AtomicInteger() + Using.resource(client(() => { + sends.incrementAndGet() + response(401, new StringEntity("rejected"), () => ()) + })) { http => + val request = new HttpPost("http://example.invalid/test") + try { + val error = intercept[IllegalArgumentException] { + HandlingUtils.sendWithRetryPolicy(http, request, Array(0), Set.empty, 0, + boundedPolicy = false, quiet = true, refreshAuth = Some { _ => + refreshes.incrementAndGet() + throw new IOException("Synthetic refresh failure") + }) + } + assert(error.getMessage.contains("Authentication replay requires the bounded retry policy")) + assert(sends.get() == 0 && refreshes.get() == 0) + } finally request.releaseConnection() + } + } + + test("bounded Fabric refresh honors fractional Retry-After within the shared retry budget") { + val times = new ConcurrentLinkedQueue[Long]() + val refreshes = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + times.add(System.nanoTime()) + val first = times.size() == 1 + val result = response(if (first) 401 else 503, new StringEntity("rejected"), + () => { closes.incrementAndGet(); () }) + if (first) result.setHeader("Retry-After", "1.0005") + result + })) { http => + val request = new HttpPost("http://example.invalid/test") + request.setHeader("Authorization", "MwcToken stale") + try { + val result = HandlingUtils.convertAndClose(HandlingUtils.sendWithRetryPolicy( + http, request, Array(0), Set.empty, 0, boundedPolicy = true, quiet = true, + refreshAuth = Some { rejected => + assert(rejected == "MwcToken stale") + refreshes.incrementAndGet() + "MwcToken fresh" + })) + val attempts = times.asScala.toVector + assert(attempts.size == 2 && refreshes.get() == 1 && closes.get() == 2) + assert(attempts(1) - attempts.head >= TimeUnit.MILLISECONDS.toNanos(1001)) + assert(result.statusLine.statusCode == 503, "Authentication replay consumes the only retry") + } finally request.releaseConnection() + } + } + + test("bounded Fabric refresh keeps long Retry-After rejections without changing legacy replay") { + Seq("60.0001", "120", DateUtils.formatDate(new java.util.Date(System.currentTimeMillis() + 120000))) + .foreach { hint => + Seq(false, true).foreach { bounded => + val sends = new AtomicInteger() + val refreshes = new AtomicInteger() + val closes = new AtomicInteger() + Using.resource(client(() => { + val first = sends.incrementAndGet() == 1 + val result = response(if (first) 401 else 200, new StringEntity(if (first) "rejected" else "retried"), + () => { closes.incrementAndGet(); () }) + if (first) result.setHeader("Retry-After", hint) + result + })) { http => + val request = new HttpPost("http://example.invalid/test") + try { + val refresh: String => String = _ => { + refreshes.incrementAndGet() + "MwcToken fresh" + } + val (received, sent) = if (bounded) { + HandlingUtils.sendWithRetryPolicy(http, request, Array(0), Set.empty, 0, + boundedPolicy = true, quiet = true, refreshAuth = Some(refresh)) -> request + } else { + HandlingUtils.sendWithFabricAuthRetries(http, new HTTPRequestData(request), Array(0), + getAuthHeader = () => "MwcToken stale", refreshAuthHeader = refresh) + } + try { + val result = HandlingUtils.convertAndClose(received) + if (bounded) assert(Option(received.getFirstHeader("Retry-After")).exists(_.getValue == hint)) + assert(result.statusLine.statusCode == (if (bounded) 401 else 200)) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == + (if (bounded) "rejected" else "retried")) + assert(sends.get() == (if (bounded) 1 else 2) && closes.get() == sends.get()) + assert(refreshes.get() == (if (bounded) 0 else 1)) + } finally sent.releaseConnection() + } finally request.releaseConnection() + } + } + } + } + + test("cancellable HTTP dispatch gates Fabric refresh and honors server hints on the real transport") { + val cases = Seq( + ("120", true, true, "MwcToken stale", false), + ("0.05", true, true, "MwcToken stale", true), + ("0.05", false, true, "MwcToken stale", false), + ("0.05", true, false, "MwcToken stale", false), + ("0.05", true, true, "Bearer explicit", false)) + val budgets = cases.map { case (hint, marked, trusted, authorization, replay) => + (hint, marked, trusted, authorization, Seq(0, 0), replay) + } :+ ("0.05", true, true, "MwcToken stale", Seq.empty[Int], false) + budgets.foreach { case (hint, marked, trusted, authorization, budget, replay) => + val times = new ConcurrentLinkedQueue[Long]() + val authHeaders = new ConcurrentLinkedQueue[String]() + val refreshHeaders = new ConcurrentLinkedQueue[String]() + val cookies = new ConcurrentLinkedQueue[String]() + val markers = new ConcurrentLinkedQueue[String]() + val server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + val pool = Executors.newFixedThreadPool(2) + server.setExecutor(pool) + server.createContext("/", (exchange: HttpExchange) => { + IOUtils.toByteArray(exchange.getRequestBody) + exchange.getRequestBody.close() + val (code, body) = if (exchange.getRequestURI.getPath == "/refresh") { + refreshHeaders.add(exchange.getRequestHeaders.getFirst("Authorization")) + 200 -> "MwcToken fresh" + } else { + times.add(System.nanoTime()) + authHeaders.add(exchange.getRequestHeaders.getFirst("Authorization")) + Option(exchange.getRequestHeaders.getFirst("Cookie")).foreach(cookies.add) + Option(exchange.getRequestHeaders.getFirst(HTTPRequestData.FabricAuthMarkerHeader)).foreach(markers.add) + if (times.size() == 1) { + exchange.getResponseHeaders.add("Retry-After", hint) + exchange.getResponseHeaders.add("Set-Cookie", "synthetic=1") + 401 -> "rejected" + } else 200 -> "accepted" + } + val bytes = body.getBytes(StandardCharsets.UTF_8) + try { + exchange.sendResponseHeaders(code, bytes.length) + exchange.getResponseBody.write(bytes) + } finally exchange.close() + }) + server.start() + try { + val address = s"http://127.0.0.1:${server.getAddress.getPort}/test" + val refreshAddress = s"http://127.0.0.1:${server.getAddress.getPort}/refresh" + val request = new HttpPost(address) + request.setHeader("Authorization", authorization) + if (marked) request.setHeader(HTTPRequestData.FabricAuthMarkerHeader, "true") + val input = HTTPRequestData.makeToRowConverter(new HTTPRequestData(request)) + request.releaseConnection() + val data = spark.createDataFrame(Seq(Row(input)).asJava, + new StructType().add("request", HTTPRequestData.schema)) + val stage = new HTTPTransformer().setInputCol("request").setOutputCol("response") + .setConcurrency(1).setConcurrentTimeout(10) + stage.setHandler { (http, requestData) => + HandlingUtils.boundedWithFabricAuth(budget, + uri => trusted && uri == address, + rejected => { + val refreshRequest = new HttpGet(refreshAddress) + refreshRequest.setHeader("Authorization", rejected) + try { + val result = HandlingUtils.convertAndClose(http.execute(refreshRequest)) + new String(result.entity.get.content, StandardCharsets.UTF_8) + } finally refreshRequest.releaseConnection() + })(http, requestData) + } + val result = HTTPResponseData.makeFromRowConverter( + stage.transformCancellable(data).head().getAs[Row]("response")) + val attempts = times.asScala.toVector + assert(result.statusLine.statusCode == (if (replay) 200 else 401)) + assert(new String(result.entity.get.content, StandardCharsets.UTF_8) == + (if (replay) "accepted" else "rejected")) + assert(attempts.size == (if (replay) 2 else 1)) + assert(refreshHeaders.asScala.toVector == (if (replay) Vector("MwcToken stale") else Vector.empty)) + assert(authHeaders.asScala.toVector == + (if (replay) Vector("MwcToken stale", "MwcToken fresh") else Vector(authorization))) + assert(cookies.isEmpty && markers.isEmpty) + if (replay) assert(attempts(1) - attempts.head >= TimeUnit.MILLISECONDS.toNanos(50)) + else assert(result.headers.exists(h => h.name.equalsIgnoreCase("Retry-After") && h.value == hint)) + } finally { + server.stop(0) + pool.shutdownNow() + } + } + } + + test("bounded Fabric refresh can be cancelled during Retry-After without refreshing or replaying") { + val sends = new AtomicInteger() + val refreshes = new AtomicInteger() + val closes = new AtomicInteger() + val failure = new AtomicReference[Throwable]() + Using.resource(client(() => { + val first = sends.incrementAndGet() == 1 + val result = response(if (first) 401 else 200, new StringEntity("rejected"), + () => { closes.incrementAndGet(); () }) + if (first) result.setHeader("Retry-After", "30") + result + })) { http => + val request = new HttpPost("http://example.invalid/test") + val caller = new Thread(() => { + try { + val result = HandlingUtils.sendWithRetryPolicy(http, request, Array(0), Set.empty, 0, + boundedPolicy = true, quiet = true, refreshAuth = Some { _ => + refreshes.incrementAndGet() + "MwcToken fresh" + }) + result.close() + } catch { + case error: Throwable => failure.set(error) + } + }) + caller.start() + try { + eventually(timeout(Span(5, Seconds))) { + assert(caller.getState == Thread.State.TIMED_WAITING) + } + caller.interrupt() + caller.join(5000) + assert(!caller.isAlive && failure.get().isInstanceOf[InterruptedException]) + assert(sends.get() == 1 && refreshes.get() == 0 && closes.get() == 1) + } finally { + caller.interrupt() + caller.join(5000) + request.releaseConnection() + assert(!caller.isAlive, "The cancelled authentication wait must not retain its worker") + } + } + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilities.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilities.scala index 5046c3d5c6f..19c9fbfb378 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilities.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilities.scala @@ -265,6 +265,8 @@ object DatabricksUtilities { .filterNot(_.getAbsolutePath.contains("Snow Leopard Detection")) // TODO Remove this exclusion .filterNot(_.getAbsolutePath.contains("Flooding Risk")) // Azure Maps Spatial API retired 9/30/2025 .filterNot(_.getAbsolutePath.contains("Geospatial Services")) // Azure Maps Spatial API retired 9/30/2025 + // Opt-in model calls are qualified separately from automatic notebook jobs. + .filterNot(_.getName == "AIPrompt.ipynb") // Split CPU notebooks into 5 partitions for parallel ADO matrix jobs. // Each partition creates its own cluster, so all 5 run simultaneously. diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilitiesSuite.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilitiesSuite.scala index 6983f9e0a79..9b7fe64aa9b 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilitiesSuite.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilitiesSuite.scala @@ -10,11 +10,26 @@ import org.scalatest.funsuite.AnyFunSuite import spray.json.DefaultJsonProtocol._ import spray.json._ +import java.io.File import java.time.Instant import scala.collection.mutable class DatabricksUtilitiesSuite extends AnyFunSuite { + test("Exclude the opt-in AIPrompt notebook from automatic Databricks inference coverage") { + val examples = DatabricksUtilities.NotebookFiles.filter(_.getName == "AIPrompt.ipynb") + assert(examples.length == 1) + assert(examples.forall(file => !DatabricksUtilities.CPUNotebooks.contains(file))) + assert(DatabricksUtilities.CPUNotebooks.nonEmpty) + } + + test("Exclude the opt-in AIPrompt notebook without constructing a Synapse cloud runner") { + val examples = Array( + new File("ExploreAlgorithmsOpenAIAIPrompt.py"), + new File("OnePlusOne.py")) + assert(SynapseTests.selectNotebooks(examples).toSeq == Seq(examples.last)) + } + test("Accept only the trusted Databricks AAD workspace") { val environment = Map( "MML_ADB_WORKSPACE_HOST" -> DatabricksUtilities.AadWorkspaceHost, diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/SynapseTests.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/SynapseTests.scala index 0ac173972be..8b65398e5a0 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/SynapseTests.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/SynapseTests.scala @@ -45,8 +45,8 @@ class SynapseTestCleanup extends TestBase { } -class SynapseTests extends TestBase { - final val excludedNotebooks: Set[String] = Set( +object SynapseTests { + private[nbtest] val ExcludedNotebooks: Set[String] = Set( "Finetune", // Excluded by design task 1829306 "GPU", "PhiModel", @@ -56,6 +56,8 @@ class SynapseTests extends TestBase { "DocumentQuestionandAnsweringwithPDFs", // Wait for Synapse fix "SetupCognitive", // No code to run "CreateaSparkCluster", // No code to run + // Opt-in model calls are qualified separately from automatic notebook jobs. + "ExploreAlgorithmsOpenAIAIPrompt", "Deploying", // New issue "TuningHyperOpt", // New issue "IsolationForests", // New issue @@ -64,13 +66,18 @@ class SynapseTests extends TestBase { "ExploreAlgorithmsDeepLearningQuickstartONNXModelInference" // ONNX package issues ) + private[nbtest] def selectNotebooks(files: Array[File]): Array[File] = + files.filterNot(file => ExcludedNotebooks.exists(excluded => file.getAbsolutePath.contains(excluded))) + .sortBy(_.getAbsolutePath) +} + +class SynapseTests extends TestBase { + final val excludedNotebooks: Set[String] = SynapseTests.ExcludedNotebooks + val generatedNotebooks = SharedNotebookE2ETestUtilities.generateNotebooks() println(s"Found ${generatedNotebooks.length} notebooks in ${SharedNotebookE2ETestUtilities.NotebooksDir}") - val selectedPythonFiles: Array[File] = generatedNotebooks - .filterNot(file => excludedNotebooks.exists(excluded => file.getAbsolutePath.contains(excluded))) - // .filter(file => file.getName().contains("OnePlusOne")) - .sortBy(_.getAbsolutePath) + val selectedPythonFiles: Array[File] = SynapseTests.selectNotebooks(generatedNotebooks) val expectedPoolCount: Int = selectedPythonFiles.length diff --git a/core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala b/core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala new file mode 100644 index 00000000000..a909943095e --- /dev/null +++ b/core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala @@ -0,0 +1,31 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark + +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.TaskMemoryManager +import org.apache.spark.metrics.MetricsSystem + +import java.util.Properties + +object TaskContextTestSupport { + class ControlledContext(onMonitorCheck: () => Unit) + extends TaskContextImpl(0, 0, 0, 0, 0, 1, Option.empty[TaskMemoryManager].orNull, + new Properties(), Option.empty[MetricsSystem].orNull, TaskMetrics.empty, 1) { + override def isInterrupted(): Boolean = { + if (Thread.currentThread().getName == "synapseml-http-cancellation-monitor") onMonitorCheck() + super.isInterrupted() + } + + def interrupt(): Unit = markInterrupted("Synthetic task cancellation") + + def complete(): Unit = markTaskCompleted(None) + } + + def withContext[T](context: TaskContext)(action: => T): T = { + val previous = TaskContext.get() + TaskContext.setTaskContext(context) + try action finally TaskContext.setTaskContext(previous) + } +} diff --git a/docs/Explore Algorithms/OpenAI/AIPrompt.ipynb b/docs/Explore Algorithms/OpenAI/AIPrompt.ipynb new file mode 100644 index 00000000000..56076af0cc5 --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/AIPrompt.ipynb @@ -0,0 +1,941 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# AIPrompt: one Spark transformer for AI inference\n", + "\n", + "`AIPrompt` extends the familiar `OpenAIPrompt` interface with explicit API selection, JSON request options and native Microsoft Decision-1 results. This notebook demonstrates the supported public workflows using small, synthetic datasets. It complements the [OpenAI notebook](https://github.com/microsoft/SynapseML/blob/master/docs/Explore%20Algorithms/OpenAI/OpenAI.ipynb) and the [Unified AI prompting guide](./Unified%20AI%20prompting.md).\n", + "\n", + "![AIPrompt routes a Spark DataFrame through five explicit API contracts, with shared controls, distinct outputs and compatibility limits.](https://github.com/user-attachments/assets/bfcb52f9-b019-455b-9e23-598a44fae229)\n", + "\n", + "Choose the API your endpoint implements, not just a model name. The five routes share an interface, but preserve their own request and response contracts. API selection does not imply provider or model availability.\n", + "\n", + "| API | Examples in this notebook |\n", + "| --- | --- |\n", + "| `chat_completions` | Templates, messages, generation options and structured output |\n", + "| `responses` | Text, request-body extensions, structured output and response chaining |\n", + "| `embeddings` | Vectors, dimensions and per-row options |\n", + "| `systemone` | Decision-1 through `AIPrompt` and the direct `AISystemOne` class |\n", + "| `decisions` | Commented reference only: OpenAIDecisions is not supported in Azure OpenAI/Foundry |\n", + "\n", + "The examples also cover output/error/usage columns, parameter inspection and copying, SparkML persistence, pipelines and pandas interoperability. Recovery and durable Delta output are a separate feature, not part of this notebook.\n", + "\n", + "### Compatibility and migration\n", + "\n", + "`AIPrompt` preserves supported scalar/column parameter bindings (the last setter wins), copying and save/load, but it is **not a drop-in replacement for every `OpenAIPrompt` behavior**. Existing legacy stages keep their own contracts. Inheriting a setter does not make that option valid for every API.\n", + "\n", + "| Area | What to expect from AIPrompt |\n", + "| --- | --- |\n", + "| Request options | Explicitly unsupported typed options fail before HTTP. Chat-only options include `n`, `stop`, penalties and `seed`; Responses-only options include `store` and `previousResponseId`; embeddings supports `dimensions` and `user`. The provider still decides which options its model accepts. |\n", + "| Wire mapping | Token-limit aliases map to `max_completion_tokens` for Chat and `max_output_tokens` for Responses. Responses maps `responseFormat` to `text.format`, `verbosity` to `text.verbosity`, and `reasoningEffort` to `reasoning.effort`. |\n", + "| Messages | `messagesCol` is input-only, unlike legacy usage that could name generated messages. Choose exactly one input binding; do not combine it with a prompt template. |\n", + "| Responses sampling | Configured `temperature` and `topP` are forwarded even for GPT-5-named deployments, rather than silently suppressed by name as in the legacy stages. Configure only options your model supports; `seed` is not supported on Responses. |\n", + "| Native APIs | `AISystemOne` and `OpenAIDecisions` lock their respective API selectors and reuse AIPrompt's implementation. SystemOne uses named questions and a lossless answer map; Decisions uses ordered questions and typed ordered answers. Their outputs remain native structs, not text postprocessing results. |\n", + "| Native configuration | Provide an explicit native endpoint, model and authentication. There is no default model or implicit reuse of global Chat routing/credentials. Foundry Microsoft Decision-1 uses SystemOne, not OpenAI Decisions. Decisions is a public OpenAI beta contract, is rejected for Azure OpenAI/Foundry hosts, and has not been live-qualified here. |\n", + "| Transport | The default transport bounds requests and retries and does not automatically follow redirects or retain cookies. An explicit custom handler opts out of this default policy. |\n", + "\n", + "See the [guide's full parameter capability table and migration notes](./Unified%20AI%20prompting.md) before porting an existing stage. `extraBody` supports provider extensions but cannot overwrite explicitly configured typed fields or request identity." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and service configuration\n", + "\n", + "Use a Spark environment with matching SynapseML Python and JVM artifacts that include `AIPrompt` and `AISystemOne`. See the [installation guide](https://microsoft.github.io/SynapseML/docs/Get%20Started/Install%20SynapseML/). Import this notebook into your Spark platform, as with `OpenAI.ipynb`.\n", + "\n", + "Provide the following environment variables through your platform's secret/configuration mechanism. No credentials are embedded or printed. Missing configuration is an error, not a skipped example.\n", + "\n", + "- `AZURE_OPENAI_ENDPOINT`: your resource root, such as `https://.openai.azure.com`.\n", + "- `AZURE_OPENAI_CHAT_DEPLOYMENT`: a deployment supporting Chat, Responses, structured output and vision; for example, a GPT-4.1-family deployment. The generation examples use sampling parameters supported by that family. Reasoning models may require different options.\n", + "- `AZURE_OPENAI_EMBEDDING_DEPLOYMENT`: a `text-embedding-3` deployment supporting the `dimensions` option.\n", + "- `AZURE_OPENAI_TOKEN`: a current Entra bearer token for Azure AI services, with inference permissions.\n", + "- `SYSTEMONE_URL`: the complete Foundry route ending in `/providers/microsoft/v1/systemone`.\n", + "- `SYSTEMONE_DEPLOYMENT`: your Microsoft Decision-1 deployment name.\n", + "- `SYSTEMONE_TOKEN`: a current Entra token authorized for that resource.\n", + "\n", + "The examples use `setAADToken`. API-key callers can use `setSubscriptionKey` instead. Do not put credentials in a prompt, `extraBody`, notebook output or a saved stage. No OpenAI Python SDK is required.\n", + "\n", + "**Cost and replay:** these are real model calls. `transform` is lazy; an action such as `collect`, `show` or `toPandas` sends requests. Repeated actions, task retries or cache eviction can repeat paid requests. Each example below materializes its inference result once. Neither Spark caching nor the provider's `store` option guarantees exactly-once execution." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "from pathlib import Path\n", + "from tempfile import TemporaryDirectory\n", + "\n", + "from pyspark.ml import Pipeline, PipelineModel\n", + "from pyspark.ml.linalg import Vector\n", + "from pyspark.sql import Row, SparkSession, functions as F\n", + "from synapse.ml.services.ai import AIPrompt, AISystemOne\n", + "\n", + "spark = SparkSession.builder.getOrCreate()\n", + "\n", + "\n", + "def required_setting(name):\n", + " value = os.environ.get(name)\n", + " if not value or not value.strip():\n", + " raise ValueError(f\"Set {name} before running this notebook\")\n", + " return value.strip()\n", + "\n", + "\n", + "openai_endpoint = required_setting(\"AZURE_OPENAI_ENDPOINT\").rstrip(\"/\")\n", + "chat_deployment = required_setting(\"AZURE_OPENAI_CHAT_DEPLOYMENT\")\n", + "embedding_deployment = required_setting(\"AZURE_OPENAI_EMBEDDING_DEPLOYMENT\")\n", + "openai_token = required_setting(\"AZURE_OPENAI_TOKEN\")\n", + "systemone_url = required_setting(\"SYSTEMONE_URL\")\n", + "systemone_deployment = required_setting(\"SYSTEMONE_DEPLOYMENT\")\n", + "systemone_token = required_setting(\"SYSTEMONE_TOKEN\")\n", + "\n", + "completed_examples = []\n", + "\n", + "\n", + "def prompt_for(api_type):\n", + " route = {\n", + " \"chat_completions\": \"chat/completions\",\n", + " \"responses\": \"responses\",\n", + " \"embeddings\": \"embeddings\",\n", + " }[api_type]\n", + " deployment = embedding_deployment if api_type == \"embeddings\" else chat_deployment\n", + " stage = (\n", + " AIPrompt()\n", + " .setApiType(api_type)\n", + " .setUrl(f\"{openai_endpoint}/openai/v1/{route}\")\n", + " .setDeploymentName(deployment)\n", + " .setAADToken(openai_token)\n", + " .setOutputCol(\"answer\")\n", + " .setErrorCol(\"error\")\n", + " .setUsageCol(\"usage\")\n", + " .setRawResponseCol(\"raw_response\")\n", + " .setConcurrency(1)\n", + " .setTimeout(120)\n", + " )\n", + " if api_type == \"responses\":\n", + " stage.setStore(False)\n", + " return stage\n", + "\n", + "\n", + "def validate_example(name, rows):\n", + " assert rows, f\"{name}: expected at least one result\"\n", + " errors = [row[\"error\"] for row in rows if row[\"error\"] is not None]\n", + " assert not errors, f\"{name}: {errors}\"\n", + " assert all(row[\"answer\"] is not None for row in rows), f\"{name}: missing answer\"\n", + " assert all(\n", + " json.loads(row[\"raw_response\"]) for row in rows\n", + " ), f\"{name}: missing raw JSON\"\n", + " completed_examples.append(name)\n", + " print(f\"Passed: {name} ({len(rows)} rows)\")\n", + " return rows\n", + "\n", + "\n", + "def run_example(name, stage, data):\n", + " return validate_example(name, stage.transform(data).collect())\n", + "\n", + "\n", + "countries = spark.createDataFrame([(\"France\",)], [\"country\"])\n", + "tickets = spark.createDataFrame(\n", + " [(\"Checkout is unavailable. Every customer receives an error.\",)], [\"text\"]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prompt templates and generation options\n", + "\n", + "Select the API explicitly. Template placeholders bind to input columns. A complete URL is used as supplied rather than having another route appended. For Chat, the normal output is the first choice's text; `rawResponseCol` retains the full response, including all choices.\n", + "\n", + "Set only options supported by your chosen API and model. Explicit options are not silently dropped." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "chat = (\n", + " prompt_for(\"chat_completions\")\n", + " .setSystemPrompt(\"Answer factual questions briefly.\")\n", + " .setPromptTemplate(\"What is the capital of {country}?\")\n", + " .setMaxTokens(128)\n", + " .setTemperature(0.0)\n", + " .setTopP(1.0)\n", + " .setUser(\"aiprompt-notebook\")\n", + ")\n", + "planned = chat.transform(countries)\n", + "planned.printSchema() # Schema planning does not call the model.\n", + "chat_rows = run_example(\"chat/template\", chat, countries)\n", + "print(chat_rows[0][\"answer\"])\n", + "print(chat_rows[0][\"usage\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Column-driven settings, constructor parameters and `setParams`\n", + "\n", + "Service parameters support scalar and column bindings. The last setter for the same parameter wins. An explicit deployment or deployment column takes precedence over applicable global defaults; a missing per-row value is an error, not an instruction to switch models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "per_row = countries.withColumn(\"deployment\", F.lit(chat_deployment)).withColumn(\n", + " \"temperature\", F.lit(0.0)\n", + ")\n", + "column_chat = (\n", + " chat.copy({})\n", + " .setDeploymentNameCol(\"deployment\")\n", + " .setTemperatureCol(\"temperature\")\n", + " .setParams(promptTemplate=\"Name the capital of {country} in one word.\")\n", + ")\n", + "run_example(\"chat/column-parameters\", column_chat, per_row)\n", + "\n", + "constructed = AIPrompt(apiType=\"responses\", promptTemplate=\"Say hello to {country}.\")\n", + "constructed.setParams(extraBody={\"metadata\": {\"example\": \"constructor\"}})\n", + "assert constructed.getApiType() == \"responses\"\n", + "assert constructed.getExtraBody()[\"metadata\"][\"example\"] == \"constructor\"\n", + "assert constructed.hasParam(\"extraBody\")\n", + "assert constructed.isSet(constructed.apiType)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Supply messages directly\n", + "\n", + "Use `setMessagesCol` when you already have role/content messages. Use an explicit Spark schema. Do not combine raw messages with a prompt template, `inputCol`, `systemPrompt` or template-attachment settings; put system messages in the array instead." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "messages = spark.createDataFrame(\n", + " [([{\"role\": \"user\", \"content\": \"Explain Spark laziness in one sentence.\"}],)],\n", + " \"messages ARRAY>\",\n", + ")\n", + "for api_type in (\"chat_completions\", \"responses\"):\n", + " message_prompt = prompt_for(api_type).setMessagesCol(\"messages\")\n", + " run_example(f\"{api_type}/raw-messages\", message_prompt, messages)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Responses and request-body extensions\n", + "\n", + "`setExtraBody` accepts a Python dictionary or a JSON-object string. Keys use the selected provider API's wire names, not Spark setter names. A typed option that you explicitly set owns its JSON subtree: a conflicting body extension is rejected. Identity, credentials, routing and asynchronous/streaming controls cannot be overridden through this parameter.\n", + "\n", + "Responses text concatenates assistant text parts in response order. Raw JSON preserves non-text outputs. AIPrompt does not execute tools or run a tool loop." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "responses = (\n", + " prompt_for(\"responses\")\n", + " .setPromptTemplate(\n", + " \"Explain the customer impact of this ticket in one sentence: {text}\"\n", + " )\n", + " .setMaxTokens(256)\n", + " .setStore(False)\n", + " .setExtraBody({\"metadata\": {\"example\": \"aiprompt-notebook\"}})\n", + ")\n", + "response_rows = run_example(\"responses/dictionary-extra-body\", responses, tickets)\n", + "print(response_rows[0][\"answer\"])\n", + "\n", + "json_extension = responses.copy({}).setExtraBody(\n", + " '{\"metadata\":{\"example\":\"json-string\"}}'\n", + ")\n", + "run_example(\"responses/json-string-extra-body\", json_extension, tickets)\n", + "\n", + "row_extensions = tickets.withColumn(\n", + " \"body_options\", F.lit('{\"metadata\":{\"example\":\"per-row\"}}')\n", + ")\n", + "column_extension = responses.copy({}).setExtraBodyCol(\"body_options\")\n", + "run_example(\"responses/column-extra-body\", column_extension, row_extensions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Structured output and Spark postprocessing\n", + "\n", + "`setResponseSchema` configures the model's JSON Schema. `setPostProcessing(\"json\")` and its `jsonSchema` option configure the separate Spark SQL output type. Neither setting substitutes for the other. Continue checking errors, refusals and incomplete output.\n", + "\n", + "Use `setResponseFormat` for an existing named envelope or the `text`/`json_object` selectors. Do not pass the bare string `json_schema`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "answer_schema = {\n", + " \"type\": \"object\",\n", + " \"properties\": {\"answer\": {\"type\": \"string\"}},\n", + " \"required\": [\"answer\"],\n", + " \"additionalProperties\": False,\n", + "}\n", + "\n", + "for api_type in (\"chat_completions\", \"responses\"):\n", + " structured = (\n", + " prompt_for(api_type)\n", + " .setPromptTemplate(\"What is the capital of {country}?\")\n", + " .setResponseSchema(answer_schema, name=\"capital\", strict=True)\n", + " .setPostProcessing(\"json\")\n", + " .setPostProcessingOptions({\"jsonSchema\": \"answer STRING\"})\n", + " )\n", + " structured_rows = run_example(\n", + " f\"{api_type}/schema-and-postprocessing\", structured, countries\n", + " )\n", + " assert isinstance(structured_rows[0][\"answer\"][\"answer\"], str)\n", + "\n", + "formatted = (\n", + " prompt_for(\"chat_completions\")\n", + " .setPromptTemplate(\n", + " \"Return a JSON object with an answer naming the capital of {country}.\"\n", + " )\n", + " .setResponseFormat({\"type\": \"json_object\"})\n", + " .setPostProcessing(\"json\")\n", + " .setPostProcessingOptions({\"jsonSchema\": \"answer STRING\"})\n", + ")\n", + "run_example(\"chat/json-object-format\", formatted, countries)\n", + "\n", + "envelope = formatted.copy({}).setResponseFormat(\n", + " {\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\"name\": \"capital\", \"strict\": True, \"schema\": answer_schema},\n", + " }\n", + ")\n", + "run_example(\"chat/explicit-schema-envelope\", envelope, countries)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Delimiter and regular-expression postprocessing\n", + "\n", + "`csv` splits the trimmed response using a Spark regular-expression delimiter; it is not an RFC CSV parser and does not interpret quoted fields. `regex` extracts the requested capture group. `regexGroup` is required. Supplying only `delimiter`, `regex` or `jsonSchema` options infers the corresponding mode; do not combine different option families." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "list_input = spark.createDataFrame(\n", + " [(\"Return exactly one;two;three with no extra text or formatting.\",)],\n", + " [\"instruction\"],\n", + ")\n", + "delimited = (\n", + " prompt_for(\"responses\")\n", + " .setInputCol(\"instruction\")\n", + " .setPostProcessing(\"csv\")\n", + " .setPostProcessingOptions({\"delimiter\": \";\"})\n", + ")\n", + "delimited_rows = run_example(\"postprocessing/delimiter\", delimited, list_input)\n", + "assert delimited_rows[0][\"answer\"] == [\"one\", \"two\", \"three\"]\n", + "\n", + "code_input = spark.createDataFrame(\n", + " [(\"Return exactly CODE: 42 with no other text.\",)], [\"instruction\"]\n", + ")\n", + "regex_prompt = (\n", + " prompt_for(\"responses\")\n", + " .setInputCol(\"instruction\")\n", + " .setPostProcessingOptions({\"regex\": r\"CODE:\\s*(\\d+)\", \"regexGroup\": \"1\"})\n", + ")\n", + "regex_rows = run_example(\"postprocessing/regex\", regex_prompt, code_input)\n", + "assert regex_rows[0][\"answer\"] == \"42\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Images and files\n", + "\n", + "For template attachments, mark a column as `path` with `setColumnTypes` and include it in the prompt template. `setFileSizeLimitMB` bounds attachment reads, not model response size. Paths can be files, HTTPS URLs or self-contained data URLs. These examples use a synthetic blue square and a one-page PDF containing `Invoice total: 42 dollars.`; no files are downloaded or uploaded to a provider file store.\n", + "\n", + "Raw message parts are API-specific: Chat uses an `image_url` object with a `url`, whereas Responses uses an `image_url` string and supports `input_file` parts. Use a vision/file-capable deployment. Do not send arbitrary binary files to Chat." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "image_data_url = \"data:image/png;base64,\" + (\n", + " \"iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAeklEQVR4nO3PUQkAIBTAwJfEYGY3\"\n", + " \"jCH8OITBAtxm7fN1wwUNaEEDWtCAFjSgBQ1oQQNa0IAWNKAFDWhBA1rQgBY0oAUNaEEDWtCA\"\n", + " \"FjSgBQ1oQQNa0IAWNKAFDWhBA1rQgBY0oAUNaEEDWtCAFjSgBQ1oQQNa0IAWPHYBWyZBPPuE\"\n", + " \"Fn4AAAAASUVORK5CYII=\"\n", + ")\n", + "file_data_url = \"data:application/pdf;base64,\" + (\n", + " \"JVBERi0xLjQKMSAwIG9iago8PCAvVHlwZSAvQ2F0YWxvZyAvUGFnZXMgMiAwIFIgPj4KZW5kb2Jq\"\n", + " \"CjIgMCBvYmoKPDwgL1R5cGUgL1BhZ2VzIC9LaWRzIFszIDAgUl0gL0NvdW50IDEgPj4KZW5kb2Jq\"\n", + " \"CjMgMCBvYmoKPDwgL1R5cGUgL1BhZ2UgL1BhcmVudCAyIDAgUiAvTWVkaWFCb3ggWzAgMCA2MTIg\"\n", + " \"NzkyXSAvUmVzb3VyY2VzIDw8IC9Gb250IDw8IC9GMSA0IDAgUiA+PiA+PiAvQ29udGVudHMgNSAw\"\n", + " \"IFIgPj4KZW5kb2JqCjQgMCBvYmoKPDwgL1R5cGUgL0ZvbnQgL1N1YnR5cGUgL1R5cGUxIC9C\"\n", + " \"YXNlRm9udCAvSGVsdmV0aWNhID4+CmVuZG9iago1IDAgb2JqCjw8IC9MZW5ndGggNTcgPj4Kc3Ry\"\n", + " \"ZWFtCkJUIC9GMSAxOCBUZiA1MCA3NTAgVGQgKEludm9pY2UgdG90YWw6IDQyIGRvbGxhcnMuKSBU\"\n", + " \"aiBFVAplbmRzdHJlYW0KZW5kb2JqCnhyZWYKMCA2CjAwMDAwMDAwMDAgNjU1MzUgZiAKMDAwMDAw\"\n", + " \"MDAwOSAwMDAwMCBuIAowMDAwMDAwMDU4IDAwMDAwIG4gCjAwMDAwMDAxMTUgMDAwMDAgbiAKMDAw\"\n", + " \"MDAwMDI0MSAwMDAwMCBuIAowMDAwMDAwMzExIDAwMDAwIG4gCnRyYWlsZXIKPDwgL1NpemUgNiAv\"\n", + " \"Um9vdCAxIDAgUiA+PgpzdGFydHhyZWYKNDE4CiUlRU9GCg==\"\n", + ")\n", + "\n", + "for name, instruction, asset, expected in (\n", + " (\"image\", \"Name the main color of this image.\", image_data_url, \"blue\"),\n", + " (\"file\", \"What is the invoice total in this file?\", file_data_url, \"42\"),\n", + "):\n", + " assets = spark.createDataFrame(\n", + " [(instruction, asset)], \"instruction STRING, asset STRING\"\n", + " )\n", + " attachment_prompt = (\n", + " prompt_for(\"responses\")\n", + " .setPromptTemplate(\"{instruction} {asset}\")\n", + " .setColumnTypes({\"asset\": \"path\"})\n", + " .setFileSizeLimitMB(2.0)\n", + " )\n", + " asset_rows = run_example(f\"attachments/template-{name}\", attachment_prompt, assets)\n", + " assert expected in asset_rows[0][\"answer\"].lower()\n", + "\n", + "chat_image = spark.createDataFrame(\n", + " [\n", + " (\n", + " [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"Name the main color of this image.\"},\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\"url\": image_data_url, \"detail\": \"low\"},\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " )\n", + " ],\n", + " \"messages ARRAY>>>>\",\n", + ")\n", + "chat_image_rows = run_example(\n", + " \"attachments/raw-chat-image\",\n", + " prompt_for(\"chat_completions\").setMessagesCol(\"messages\"),\n", + " chat_image,\n", + ")\n", + "assert \"blue\" in chat_image_rows[0][\"answer\"].lower()\n", + "\n", + "response_assets = spark.createDataFrame(\n", + " [\n", + " (\n", + " [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"input_text\",\n", + " \"text\": \"Name the image's main color and the invoice total.\",\n", + " },\n", + " {\"type\": \"input_image\", \"image_url\": image_data_url},\n", + " {\n", + " \"type\": \"input_file\",\n", + " \"file_data\": file_data_url,\n", + " \"filename\": \"invoice.pdf\",\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " )\n", + " ],\n", + " \"messages ARRAY>>>\",\n", + ")\n", + "response_asset_rows = run_example(\n", + " \"attachments/raw-responses-image-and-file\",\n", + " prompt_for(\"responses\").setMessagesCol(\"messages\"),\n", + " response_assets,\n", + ")\n", + "assert \"blue\" in response_asset_rows[0][\"answer\"].lower()\n", + "assert \"42\" in response_asset_rows[0][\"answer\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Responses storage and conversation chaining\n", + "\n", + "`setStore(True)` stores a response at the provider, not a Spark result table. `setResponseIdCol` exposes its identifier. Supply that identifier with `setPreviousResponseId` or the per-row `setPreviousResponseIdCol` to continue a conversation. This requires provider-side storage to be enabled for your resource.\n", + "\n", + "This example stores only synthetic content and deletes its own stored response in `finally`. It never lists or deletes other responses. The continuation calls use `store=False`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from urllib.parse import quote\n", + "import requests\n", + "import sys\n", + "\n", + "stored = (\n", + " prompt_for(\"responses\")\n", + " .setPromptTemplate(\n", + " \"Remember the code name ALPHA. Reply with a brief acknowledgement.\"\n", + " )\n", + " .setStore(True)\n", + " .setResponseIdCol(\"response_id\")\n", + ")\n", + "created_response_ids = []\n", + "deleted_response_ids = []\n", + "try:\n", + " first_rows = stored.transform(countries).collect()\n", + " created_response_ids.extend(\n", + " row[stored.getResponseIdCol()]\n", + " for row in first_rows\n", + " if row[stored.getResponseIdCol()]\n", + " )\n", + " validate_example(\"responses/store-and-response-id\", first_rows)\n", + " response_id = first_rows[0][stored.getResponseIdCol()]\n", + " assert response_id\n", + "\n", + " continued = (\n", + " prompt_for(\"responses\")\n", + " .setStore(False)\n", + " .setPreviousResponseId(response_id)\n", + " .setPromptTemplate(\n", + " \"What code name did I give you? Reply with just the code name.\"\n", + " )\n", + " )\n", + " continued_rows = run_example(\"responses/scalar-previous-id\", continued, countries)\n", + " assert \"ALPHA\" in continued_rows[0][\"answer\"].upper()\n", + "\n", + " with_ids = countries.withColumn(\"previous_id\", F.lit(response_id))\n", + " column_continuation = continued.copy({}).setPreviousResponseIdCol(\"previous_id\")\n", + " column_rows = run_example(\n", + " \"responses/column-previous-id\", column_continuation, with_ids\n", + " )\n", + " assert \"ALPHA\" in column_rows[0][\"answer\"].upper()\n", + "finally:\n", + " original_error = sys.exc_info()[1]\n", + " cleanup_errors = []\n", + " for response_id in created_response_ids:\n", + " try:\n", + " deleted = requests.delete(\n", + " f\"{openai_endpoint}/openai/v1/responses/{quote(response_id, safe='')}\",\n", + " headers={\"Authorization\": f\"Bearer {openai_token}\"},\n", + " timeout=60,\n", + " )\n", + " deleted.raise_for_status()\n", + " if deleted.status_code != 204:\n", + " confirmation = deleted.json()\n", + " if (\n", + " not isinstance(confirmation, dict)\n", + " or confirmation.get(\"deleted\") is not True\n", + " ):\n", + " raise ValueError(\"Provider did not confirm response deletion\")\n", + " deleted_response_ids.append(response_id)\n", + " except (requests.RequestException, ValueError) as error:\n", + " cleanup_errors.append((response_id, error))\n", + " print(f\"Deleted {len(deleted_response_ids)} response(s) created by this example.\")\n", + " if cleanup_errors:\n", + " failed_ids = [identifier for identifier, _ in cleanup_errors]\n", + " message = (\n", + " f\"Could not delete stored response(s) {failed_ids}; retry these deletions.\"\n", + " )\n", + " if original_error is None:\n", + " raise RuntimeError(message) from cleanup_errors[0][1]\n", + " print(message, file=sys.stderr)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Embeddings\n", + "\n", + "Select `embeddings` and bind text with `setTextCol`. Results are Spark ML vectors, not JSON strings. The `dimensions` option requires a model that supports it. AIPrompt supports float embeddings; it rejects unsupported encodings before sending a request. Both scalar and per-row `user` values are sent to the embedding API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "embedding = (\n", + " prompt_for(\"embeddings\")\n", + " .setTextCol(\"text\")\n", + " .setDimensions(64)\n", + " .setUser(\"aiprompt-notebook\")\n", + " .setExtraBody({\"encoding_format\": \"float\"})\n", + ")\n", + "embedding_rows = run_example(\"embeddings/scalar-options\", embedding, tickets)\n", + "assert isinstance(embedding_rows[0][\"answer\"], Vector)\n", + "assert len(embedding_rows[0][\"answer\"]) == 64\n", + "\n", + "embedding_input = tickets.withColumn(\"dimensions\", F.lit(32)).withColumn(\n", + " \"request_user\", F.lit(\"notebook-row\")\n", + ")\n", + "row_embedding = (\n", + " embedding.copy({}).setDimensionsCol(\"dimensions\").setUserCol(\"request_user\")\n", + ")\n", + "row_embedding_rows = run_example(\n", + " \"embeddings/column-options\", row_embedding, embedding_input\n", + ")\n", + "assert len(row_embedding_rows[0][\"answer\"]) == 32" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Microsoft Decision-1 with `AISystemOne` and `AIPrompt`\n", + "\n", + "`AISystemOne` is the directly usable native transformer. The facade uses the same protocol with `apiType=\"systemone\"`. This is not a Chat emulation and is distinct from OpenAIDecisions.\n", + "\n", + "Use the complete Foundry System One route without a dated Azure OpenAI API-version parameter. Supply your deployment explicitly; the native class does not invent a default model. Questions form a named object with `noul`, `choice` and `score` entries.\n", + "\n", + "The result's `answers` field is a map of complete JSON-encoded answer values. Decode the values without assuming that every answer has the same undocumented probability or confidence fields." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "questions = {\n", + " \"urgent\": {\"type\": \"noul\", \"instructions\": \"Is checkout unavailable?\"},\n", + " \"team\": {\n", + " \"type\": \"choice\",\n", + " \"instructions\": \"Which team should handle this request?\",\n", + " \"criteria\": {\"billing\": \"Billing problems\", \"engineering\": \"Software outages\"},\n", + " },\n", + " \"severity\": {\n", + " \"type\": \"score\",\n", + " \"instructions\": \"How severe is the issue?\",\n", + " \"criteria\": [\"Cosmetic\", \"Workaround available\", \"Service unavailable\"],\n", + " },\n", + "}\n", + "\n", + "\n", + "def configure_systemone(stage):\n", + " return (\n", + " stage.setUrl(systemone_url)\n", + " .setDeploymentName(systemone_deployment)\n", + " .setAADToken(systemone_token)\n", + " .setStateCol(\"text\")\n", + " .setQuestions(questions)\n", + " .setOutputCol(\"answer\")\n", + " .setErrorCol(\"error\")\n", + " .setUsageCol(\"usage\")\n", + " .setRawResponseCol(\"raw_response\")\n", + " .setConcurrency(1)\n", + " .setTimeout(120)\n", + " )\n", + "\n", + "\n", + "systemone = configure_systemone(AISystemOne())\n", + "native_rows = run_example(\"systemone/direct\", systemone, tickets)\n", + "decoded_answers = {\n", + " name: json.loads(value)\n", + " for name, value in native_rows[0][\"answer\"][\"answers\"].items()\n", + "}\n", + "assert set(decoded_answers) == set(questions)\n", + "print(decoded_answers)\n", + "\n", + "systemone_facade = configure_systemone(AIPrompt().setApiType(\"systemone\"))\n", + "run_example(\"systemone/facade\", systemone_facade, tickets)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### JSON state and per-row questions\n", + "\n", + "`setStateMode(\"json\")` interprets the bound string as JSON rather than literal text. `setQuestionsCol` takes a JSON-string column. Scalar question bindings also accept JSON strings. As with other service parameters, the last scalar/column setter wins." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "json_tickets = spark.createDataFrame(\n", + " [(json.dumps({\"service\": \"checkout\", \"available\": False}), json.dumps(questions))],\n", + " [\"state_json\", \"questions_json\"],\n", + ")\n", + "json_systemone = (\n", + " systemone.copy({})\n", + " .setStateCol(\"state_json\")\n", + " .setStateMode(\"json\")\n", + " .setQuestionsCol(\"questions_json\")\n", + ")\n", + "run_example(\"systemone/json-state-and-column-questions\", json_systemone, json_tickets)\n", + "\n", + "string_questions = systemone.copy({}).setQuestions(json.dumps(questions))\n", + "run_example(\"systemone/json-string-questions\", string_questions, tickets)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Row errors and lossless observations\n", + "\n", + "Malformed per-row body extensions become row errors. Valid rows can still succeed. Invalid stage-wide configuration instead fails preflight. HTTP errors, provider failures, refusals, incomplete output and parse failures must not be interpreted as successful answers.\n", + "\n", + "`rawResponseCol` contains response JSON, not authentication headers. It can contain sensitive model output: apply your normal access and retention controls." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mixed_input = spark.createDataFrame(\n", + " [(\"valid\", \"Say hello.\", \"{}\"), (\"invalid\", \"Say hello.\", \"[]\")],\n", + " [\"case\", \"text\", \"body_options\"],\n", + ")\n", + "row_validation = (\n", + " prompt_for(\"chat_completions\")\n", + " .setPromptTemplate(\"{text}\")\n", + " .setExtraBodyCol(\"body_options\")\n", + ")\n", + "mixed_rows = {\n", + " row[\"case\"]: row for row in row_validation.transform(mixed_input).collect()\n", + "}\n", + "assert mixed_rows[\"valid\"][\"error\"] is None\n", + "assert mixed_rows[\"valid\"][\"answer\"] is not None\n", + "assert mixed_rows[\"invalid\"][\"error\"] is not None\n", + "assert mixed_rows[\"invalid\"][\"answer\"] is None\n", + "completed_examples.append(\"validation/per-row-error\")\n", + "print(\"Passed: valid row succeeds; malformed extension is an explicit row error\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Copy, clear and inspect parameters\n", + "\n", + "These are ordinary SparkML stages. `copy` keeps configuration without mutating the original. `clear` removes an explicit binding. Do not print `extractParamMap()` on an authenticated stage because its parameters can include credentials." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "copied = chat.copy({chat.temperature: 0.1})\n", + "assert copied.getTemperature() == 0.1\n", + "assert chat.getTemperature() == 0.0\n", + "copied.clear(copied.temperature)\n", + "copied.setTemperature(0.0)\n", + "run_example(\"sparkml/copy-and-clear\", copied, countries)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Save/load and pipelines\n", + "\n", + "Save a credential-free copy, then rebind authentication after loading with the same concrete class. This saves the stage configuration, not model results. The example uses a temporary local directory; on a distributed cluster, use a fresh caller-owned path accessible to all Spark nodes instead.\n", + "\n", + "A `Pipeline` composes AIPrompt with other SparkML stages. Fitting this transformer-only pipeline does not perform inference; transforming and materializing its result does." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "saveable = chat.copy({})\n", + "saveable.clear(saveable.AADToken)\n", + "\n", + "with TemporaryDirectory(prefix=\"aiprompt-example-\") as temporary:\n", + " stage_path = str(Path(temporary) / \"stage\")\n", + " saveable.write().save(stage_path)\n", + " loaded = AIPrompt.load(stage_path).setAADToken(openai_token)\n", + " run_example(\"sparkml/save-load\", loaded, countries)\n", + "\n", + " read_loaded = AIPrompt.read().load(stage_path).setAADToken(openai_token)\n", + " assert read_loaded.getPromptTemplate() == chat.getPromptTemplate()\n", + "\n", + " pipeline_path = str(Path(temporary) / \"pipeline\")\n", + " pipeline = Pipeline(stages=[saveable]).fit(countries)\n", + " pipeline.write().save(pipeline_path)\n", + " restored_pipeline = PipelineModel.load(pipeline_path)\n", + " restored_pipeline.stages[0].setAADToken(openai_token)\n", + " pipeline_rows = restored_pipeline.transform(countries).collect()\n", + " assert len(pipeline_rows) == 1 and pipeline_rows[0][\"error\"] is None\n", + " assert pipeline_rows[0][\"answer\"] is not None\n", + " completed_examples.append(\"sparkml/pipeline-save-load\")\n", + "\n", + " native_path = str(Path(temporary) / \"systemone\")\n", + " native_saveable = systemone.copy({})\n", + " native_saveable.clear(native_saveable.AADToken)\n", + " native_saveable.save(native_path)\n", + " native_loaded = AISystemOne.load(native_path).setAADToken(systemone_token)\n", + " run_example(\"systemone/save-load\", native_loaded, tickets)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## pandas interoperability\n", + "\n", + "AIPrompt transforms Spark DataFrames. Convert a small pandas input to Spark and collect the small result back to pandas. Do not collect a large distributed dataset to the driver." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "pandas_input = pd.DataFrame({\"country\": [\"France\"]})\n", + "pandas_result = chat.transform(spark.createDataFrame(pandas_input)).toPandas()\n", + "assert len(pandas_result) == 1\n", + "assert pandas_result[\"error\"].isna().all()\n", + "assert pandas_result[\"answer\"].str.len().gt(0).all()\n", + "completed_examples.append(\"pandas/round-trip\")\n", + "print(pandas_result[[\"country\", \"answer\"]])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAIDecisions: commented reference only\n", + "\n", + "**Do not run this section against Azure OpenAI or Microsoft Foundry. OpenAIDecisions is not supported there.** The Foundry Decision-1 examples above use `AISystemOne`, a different API. A deployed model name alone does not establish Decisions endpoint support.\n", + "\n", + "The example below is intentionally entirely commented out. It illustrates the direct class and facade for a separately qualified OpenAI Decisions endpoint; it is not part of this notebook's executed examples. Native protocol tests do not establish live service availability." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# OpenAIDecisions is not supported in Azure OpenAI/Foundry. Do not run these examples there.\n", + "# from synapse.ml.services.openai import OpenAIDecisions\n", + "#\n", + "# decision_questions = [\n", + "# {\"type\": \"predicate\", \"name\": \"urgent\", \"instructions\": \"Is checkout unavailable?\"}\n", + "# ]\n", + "# decisions = (\n", + "# OpenAIDecisions()\n", + "# .setUrl(\"https://api.openai.com/v1/decisions\")\n", + "# .setModel(\"\")\n", + "# .setSubscriptionKey(\"\")\n", + "# .setInputCol(\"text\")\n", + "# .setQuestions(decision_questions)\n", + "# .setSafetyIdentifier(\"aiprompt-notebook\")\n", + "# .setOutputCol(\"answer\")\n", + "# .setErrorCol(\"error\")\n", + "# .setRawResponseCol(\"raw_response\")\n", + "# )\n", + "# decisions.transform(tickets).show(truncate=False)\n", + "#\n", + "# facade = (\n", + "# AIPrompt().setApiType(\"decisions\")\n", + "# .setUrl(\"https://api.openai.com/v1/decisions\")\n", + "# .setModel(\"\")\n", + "# .setSubscriptionKey(\"\")\n", + "# .setInputCol(\"text\")\n", + "# .setQuestions(decision_questions)\n", + "# .setOutputCol(\"answer\")\n", + "# .setErrorCol(\"error\")\n", + "# )\n", + "# facade.transform(tickets).show(truncate=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Execution summary\n", + "\n", + "Every active inference example above checks its result. The commented Decisions section performs no imports, creates no stage and sends no requests. Notebook outputs are intentionally cleared in source control so credentials, endpoint details and model responses are not published." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "assert len(completed_examples) == len(set(completed_examples))\n", + "print(f\"Completed {len(completed_examples)} supported examples.\")\n", + "for name in completed_examples:\n", + " print(name)\n", + "print(\n", + " \"OpenAIDecisions: intentionally not executed; unsupported in Azure OpenAI/Foundry.\"\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/Explore Algorithms/OpenAI/Unified AI prompting.md b/docs/Explore Algorithms/OpenAI/Unified AI prompting.md new file mode 100644 index 00000000000..328564e39f4 --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/Unified AI prompting.md @@ -0,0 +1,917 @@ +--- +title: Unified AI prompting +hide_title: true +sidebar_label: Unified AI prompting +description: Call five model APIs from Spark with typed inputs, native questions, and explicit diagnostics. +--- + +# Unified AI prompting + +`AIPrompt` applies a synchronous model API to a Spark DataFrame, using Spark +partitions and bounded concurrent HTTP requests. It does not require the OpenAI +Python SDK. Existing `OpenAIPrompt` and `OpenAIEmbedding` stages remain available +with their existing contracts. Migrating to `AIPrompt` is not a drop-in replacement: +for example, Responses forwards configured sampling options even for GPT-5-named +deployments instead of applying the legacy name-based suppression described below. + +![AIPrompt routes a Spark DataFrame through five explicit API contracts, with shared controls, distinct outputs and compatibility limits.](https://github.com/user-attachments/assets/bfcb52f9-b019-455b-9e23-598a44fae229) + +The diagram summarizes routing, not universal provider support. Inherited setters +are validated against the selected API; a model can impose further restrictions. +Foundry Microsoft Decision-1 uses SystemOne, not the public OpenAI Decisions beta. +Decisions is rejected for Azure OpenAI/Foundry hosts and has not been live-qualified +here. + +| `apiType` | Input | Output | +| --- | --- | --- | +| `chat_completions` (default) | Prompt template, text column, or messages | First choice's text, or configured Spark postprocessing type | +| `responses` | Prompt template, text column, or messages | Assistant text parts concatenated in response order, or configured postprocessing type | +| `systemone` | Text/JSON state and named questions | Struct with returned model, lossless answer map and provider usage JSON | +| `decisions` | Text or documented user messages and ordered questions | Struct with returned model, typed ordered answers and normalized usage | +| `embeddings` | Text column, scalar text or prompt template | Spark ML vector | + +Use `AISystemOne` and `OpenAIDecisions` directly when only that protocol is needed. +They share the facade's request validation, transport, diagnostics and persistence. +They have **no default model**. Model availability is provider-specific; selecting +an API never silently redirects to another API based on a model or hostname. + +Inference is lazy: `show`, `collect`, repeated actions, recomputation and Spark +task retries can repeat paid requests. These stages do not persist input or +result datasets and do not require Delta. Provider `store=true` is unrelated to +persisting a Spark result. Plan how your application will materialize outcomes; +caching alone is not an exactly-once inference guarantee. + +See the [AIPrompt notebook](AIPrompt.md) for examples. Its runnable source is +[AIPrompt.ipynb](https://github.com/microsoft/SynapseML/blob/master/docs/Explore%20Algorithms/OpenAI/AIPrompt.ipynb). + +## PySpark and Scala + +```python +import os +from synapse.ml.services.ai import AIPrompt, AISystemOne +from synapse.ml.services.openai import OpenAIDecisions + +prompt = ( + AIPrompt() + .setApiType("responses") + .setUrl(os.environ["MODEL_ENDPOINT"]) + .setDeploymentName(os.environ["MODEL_DEPLOYMENT"]) + .setSubscriptionKey(os.environ["MODEL_API_KEY"]) + .setPromptTemplate("Summarize this ticket: {ticket}") + .setOutputCol("answer") + .setErrorCol("error") + .setUsageCol("usage") + .setRawResponseCol("raw") + .setExtraBody({"text": {"verbosity": "low"}}) +) +result = prompt.transform(tickets) # Lazy: no model calls yet. +``` + +```scala +import com.microsoft.azure.synapse.ml.services.ai.AIPrompt + +val prompt = new AIPrompt() + .setApiType("responses") + .setUrl(sys.env("MODEL_ENDPOINT")) + .setDeploymentName(sys.env("MODEL_DEPLOYMENT")) + .setSubscriptionKey(sys.env("MODEL_API_KEY")) + .setPromptTemplate("Summarize this ticket: {ticket}") + .setOutputCol("answer").setErrorCol("error") + .setUsageCol("usage").setRawResponseCol("raw") + .setExtraBody("""{"text":{"verbosity":"low"}}""") +val result = prompt.transform(tickets) +``` + +Do not print credentials or save a credential-bearing stage into a shared model +directory. Standard SparkML stage persistence preserves configured Params, +including explicitly supplied authentication Params. +Use the corresponding stage's `load`/`read` API to restore a saved `AIPrompt`, +`AISystemOne` or `OpenAIDecisions`, or the standard Spark `PipelineModel` loader +for a saved pipeline. Copy and save/load preserve native JSON values, parameter +defaults and explicit bindings. Saving a stage stores its configuration, not +inference outcomes. +`Pipeline.fit`, cross-validation and repeated scoring can evaluate inference +more than once. Account for those model calls before placing a stage in a +training or tuning pipeline. + +Choose exactly one input binding: `promptTemplate`, `messagesCol`, `stateCol`, +`inputCol`, or `text`/`textCol`. Bind native JSON strings with `inputMode="json"` +(`setStateMode("json")` is an alias); the default treats them as text. +For System One, `stateCol` and `inputCol` are supported but distinct persisted +Params, not aliases. +For Responses and Decisions, message structs may include `type="message"` or +a null/absent `type`; other non-null message types are rejected before HTTP. +Decisions JSON messages follow the same rule; a null `type` is omitted on the wire. +For inline Decisions images, null `file_id` and `detail` options are also omitted; +non-null file IDs remain unsupported. Decisions JSON input preserves provider +extension members; `messagesCol` uses the narrower Responses-compatible conversion. +Chat message structs do not accept this discriminator. +Spark's resolver selects the top-level `messagesCol` column; nested protocol keys +are case-sensitive JSON names such as `role`, `content`, `type` and `name`. +Duplicate field names anywhere within a bound message struct fail schema planning, +including for empty inputs, rather than silently discarding one of the values. +Explicit messages already contain their system instructions and attachments: +`messagesCol` cannot be combined with an explicitly set `systemPrompt`, +`fileSizeLimitMB`, `dropPrompt`, or non-empty `columnTypes`. For Decisions it +also cannot be combined with an explicitly set `inputMode`; message structs +are already structured input. These conflicts fail before HTTP. +Unlike legacy `OpenAIPrompt`, `AIPrompt` uses an explicitly set `messagesCol` +only as an input binding, not as the name of generated output messages. +With a prompt template and `dropPrompt=False`, retained messages use a UID-based, +collision-safe name reported by `transformSchema`; there is no separate naming +Param. Rename that column on the returned DataFrame if a stable consumer-facing +name is needed. Existing `OpenAIPrompt` behavior is unchanged. +Prompt templates must render a Spark string. Convert array or binary expressions +explicitly, for example with `to_json` or `cast(... as string)`, before templating. +A template's outer projection must be row-local: +it cannot filter, aggregate or expand the input rows. Compute window expressions +(including `window(...)`), outer aggregates and generators as input columns first. +Templates may contain scalar SQL subqueries, including aggregates inside the +subquery. These queries can read other data again on each action. +Row-local templates also support command-result DataFrames such as `SHOW TABLES` +and `SHOW COLUMNS`. +Direct attachment placeholders follow Spark's column resolver and use the +attachment filename. Wildcards such as `{to_json(struct(*))}` still see only +the public input columns with their original values, not generated scratch fields. +Escaped braces and SQL expressions retain their usual template meaning. +Native and embedding input cells distinguish missing values from whitespace: +null or empty strings produce `missing_input`. System One and embeddings send +nonempty text unchanged, including whitespace-only text. Decisions requires +nonblank text, including text inside JSON messages; whitespace-only text produces +`invalid_request` without HTTP. Nonblank text is sent unchanged, not trimmed. +JSON mode requires a JSON value, so a whitespace-only JSON document also produces +`invalid_request`. Normalize or filter blank input before inference when your +application treats it as missing. + +### Native question types + +System One uses a named object, not a chat prompt or the Decisions question array: + +```python +system_questions = { + "repeat_contact": { + "type": "noul", + "instructions": "Has the customer contacted support before?", + }, + "team": { + "type": "choice", + "instructions": "Which team should handle this request?", + "criteria": { + "billing": "Charges, invoices and refunds", + "engineering": "Bugs, errors and outages", + }, + }, + "severity": { + "type": "score", + "instructions": "How severe is this issue?", + "criteria": ["Cosmetic", "Major", "Critical"], + }, +} +systemone = ( + AISystemOne() + .setUrl(os.environ["FOUNDRY_ENDPOINT"]) + .setDeploymentName(os.environ["SYSTEMONE_DEPLOYMENT"]) + .setSubscriptionKey(os.environ["FOUNDRY_API_KEY"]) + .setStateCol("ticket") + .setQuestions(system_questions) + .setOutputCol("decision") + .setErrorCol("error") + .setRawResponseCol("raw") +) +``` + +The route is `/providers/microsoft/v1/systemone`, without an +Azure OpenAI dated API version. `Decision-1` is an example deployment name, not +a built-in default. Returned model identity can differ from the deployment alias. +`decision.answers` is a **SynapseML** `map` whose values are complete +JSON encodings. A numeric answer remains `"1.7"` as JSON; a provider string is +quoted JSON; objects, booleans, arrays and nulls remain lossless. This does not +assert an undocumented provider answer schema. Decode values according to your +application's schema, or retain them with the raw response. + +Decisions uses the public OpenAI beta contract and an ordered question array: + +```python +decision_questions = [ + {"type": "predicate", "name": "repeat_contact", + "instructions": "Has the customer contacted support before?"}, + {"type": "choice", "name": "team", + "instructions": "Which team should handle this request?", + "choices": [ + {"value": "billing", "description": "Charges, invoices and refunds"}, + {"value": "engineering", "description": "Bugs, errors and outages"}, + ]}, + {"type": "score", "name": "severity", + "instructions": "How severe is this issue?", + "levels": [ + {"label": "Cosmetic", "description": "Appearance only"}, + {"label": "Major", "description": "A task fails but a workaround exists"}, + {"label": "Critical", "description": "A task is fully blocked"}, + ]}, +] +decisions = ( + OpenAIDecisions() + .setUrl("https://api.openai.com/v1") + .setModel(os.environ["DECISIONS_MODEL"]) + .setSubscriptionKey(os.environ["OPENAI_API_KEY"]) + .setInputCol("ticket") + .setQuestions(decision_questions) + .setOutputCol("decision") + .setErrorCol("error") +) +``` + +This is **not** a Foundry endpoint. An available System One deployment does not +establish Decisions availability. Decisions preflight rejects Azure resource +hosts in the public and sovereign clouds (`*.azure.com`, `*.azure.us` and +`*.azure.cn`), including complete URL overrides. Live Decisions qualification is tracked in +[microsoft/SynapseML#2760](https://github.com/microsoft/SynapseML/issues/2760); +local protocol fixtures do not establish live availability. +Choice values use a tagged Spark struct +(`type`, `stringValue`, `booleanValue`); `true` and `"true"` remain distinct. +Score probability entries preserve numeric index and label. Scores must be +between zero and the last probability-entry index, inclusive; bounds are checked +on the original JSON number before conversion to Spark `DoubleType`. Out-of-range +scores produce a `parse_error`. Refusals remain +typed answers alongside any other answers. Inline images are supported in the +documented Decisions input form; hosted-image URLs are rejected. Image data URLs +accept case-insensitive scheme, media type and `base64` markers, plus MIME +parameters before `;base64`. Parameter values with special characters must use +URL escapes. The original URL is forwarded unchanged; the payload must decode +to nonempty base64 data. Syntax validation does not allocate a decoded image. +Callers must still bound inline payloads to their executor memory budget and the +provider's request limits; `fileSizeLimitMB` does not apply to Decisions input. + +Native JSON validation checks the supported question and input forms, including +recognized image fields; it is not a complete provider-schema allowlist. +Additional provider-owned members are forwarded without rewriting them, and +the provider determines whether they are valid. This permits supported options +such as Decisions `prompt_cache_breakpoint` on input parts. `messagesCol` uses +the narrower Responses-compatible message converter. + +The same three-question workflows are available directly from Scala: + +```scala +import com.microsoft.azure.synapse.ml.services.ai.AISystemOne +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions + +val systemQuestions = """{ + "repeat_contact": {"type":"noul", "instructions":"Has the customer contacted support before?"}, + "team": {"type":"choice", "instructions":"Which team should handle this request?", + "criteria":{"billing":"Charges, invoices and refunds", "engineering":"Bugs, errors and outages"}}, + "severity": {"type":"score", "instructions":"How severe is this issue?", + "criteria":["Cosmetic", "Major", "Critical"]} +}""" +val systemone = new AISystemOne() + .setUrl(sys.env("FOUNDRY_ENDPOINT")) + .setDeploymentName(sys.env("SYSTEMONE_DEPLOYMENT")) + .setSubscriptionKey(sys.env("FOUNDRY_API_KEY")) + .setStateCol("ticket").setQuestions(systemQuestions) + .setOutputCol("decision").setErrorCol("error").setRawResponseCol("raw") +val systemResults = systemone.transform(tickets) // Lazy; an action sends requests. + +val decisionQuestions = """[ + {"type":"predicate", "name":"repeat_contact", + "instructions":"Has the customer contacted support before?"}, + {"type":"choice", "name":"team", "instructions":"Which team should handle this request?", + "choices":[{"value":"billing", "description":"Charges, invoices and refunds"}, + {"value":"engineering", "description":"Bugs, errors and outages"}]}, + {"type":"score", "name":"severity", "instructions":"How severe is this issue?", + "levels":[{"label":"Cosmetic", "description":"Appearance only"}, + {"label":"Major", "description":"A workaround exists"}, + {"label":"Critical", "description":"A task is fully blocked"}]} +]""" +val decisions = new OpenAIDecisions() + .setUrl("https://api.openai.com/v1") + .setModel(sys.env("DECISIONS_MODEL")) + .setSubscriptionKey(sys.env("OPENAI_API_KEY")) + .setInputCol("ticket").setQuestions(decisionQuestions) + .setOutputCol("decision").setErrorCol("error").setRawResponseCol("raw") +val decisionResults = decisions.transform(tickets) +``` + +### Validate a native deployment + +The opt-in live suites use synthetic support tickets and exercise both the +protocol-specific stage and `AIPrompt`. Configure credentials only in the test +process environment, never in a notebook checked into source control. + +| API | Required environment variables | Authentication | +| --- | --- | --- | +| System One | `SYNAPSEML_SYSTEMONE_ENDPOINT`, `SYNAPSEML_SYSTEMONE_DEPLOYMENT` | `SYNAPSEML_SYSTEMONE_TOKEN` or `SYNAPSEML_SYSTEMONE_KEY` | +| Decisions | `SYNAPSEML_DECISIONS_ENDPOINT`, `SYNAPSEML_DECISIONS_MODEL` | `SYNAPSEML_DECISIONS_TOKEN` or `SYNAPSEML_DECISIONS_KEY` | + +Use an endpoint that exposes the selected native API and a model authorized for +it. Model discovery or a successful chat request does not prove native API +availability. Configure the native endpoint and its authentication on the stage; +native APIs do not inherit process-wide OpenAI URL or API-key defaults, and the +facade does not infer an implicit Fabric native route. + +```bash +sbt "cognitive/testOnly com.microsoft.azure.synapse.ml.services.ai.AISystemOneLiveSuite com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisionsLiveSuite" +python -m unittest discover -s cognitive/src/test/python/synapsemltest/services -p test_AIPromptLive.py -v +``` + +For the Python command, install this build's generated Python packages and +configure Spark with the matching JARs first. The Python tests check the public +wrappers, runtime schema, null-input errors, native answers, usage, raw responses, +and conversion of the cached results to pandas. Each configured test makes real +inference requests and can incur charges. Missing endpoint, model, or credentials +explicitly skip a live test; a configured endpoint's failure fails the test. +Skipped tests and local protocol fixtures are not live qualification evidence. + +## Typed option inventory and body extensions + +Common configuration includes endpoint/complete URL, deployment/model, scalar or +column authentication, custom headers, concurrency, HTTP timeout, concurrent +timeout, output/error/usage/raw columns, and SparkML copy/save/load. An explicit +deployment/model binding wins over operation-specific global defaults. A null +explicit deployment column is a row error, not permission to use another model. +When implicit defaults select a deployment, unused model defaults do not affect +input-column validation or requests. The embedding-specific deployment +default takes precedence over the general deployment default. +Native System One and Decisions APIs require an explicit model binding and do not +inherit global chat deployment, Foundry model, endpoint or API-key defaults. +Configure credentials for the selected native provider on the stage rather than +reusing process-wide OpenAI defaults. +Endpoint URLs and complete `customUrlRoot` overrides must use a DNS host name or +IP address that Java's URI parser accepts as a server host. Host names containing +underscores are not supported by these new stages, even when a legacy stage or +self-hosted gateway accepts them. Configure a DNS-compatible alias instead. +Invalid hosts fail preflight; credentials and +fragments in endpoint URLs are separately rejected. An explicit endpoint port +must be between 1 and 65535; omit it to use the scheme's default. +The selected endpoint must contain correctly paired UTF-16 surrogates, including +in its path and query. Malformed strings fail before HTTP requests. +Valid supplementary characters and U+FFFD remain supported; +an unused `url` does not override or invalidate an explicit `customUrlRoot`. +On OpenAI-hosted endpoints (`api.openai.com` and `*.api.openai.com`, including +regional endpoints), subscription keys use the `Authorization: Bearer` header +for Chat Completions, Responses and embeddings. Decisions uses bearer keys on +every supported endpoint. Pass the raw API key to `setSubscriptionKey`, without +an `Authorization` scheme prefix. Use custom authentication headers when supplying +a complete header value. Host comparisons ignore case and one DNS trailing dot; +the supplied URL itself is retained. These rules do not classify lookalike domains +such as `api.openai.com.example.org` as OpenAI hosts. +For these OpenAI hosts, a bare `setUrl` root (with or without trailing slashes) +automatically uses `/v1` for Chat Completions, Responses, embeddings and Decisions. +For example, `setUrl("https://api.openai.com")` routes Responses requests to +`https://api.openai.com/v1/responses`. Dated API-version defaults are ignored and +explicit `apiVersion` settings are rejected, just as with an explicit `/v1` URL. +Existing non-root paths, complete `customUrlRoot` overrides and other hosts are +not rewritten. The saved URL parameter is unchanged. +All new stages send one `Content-Type: application/json` header, replacing +caller or telemetry content-type headers regardless of casing. Other header +values retain their existing precedence; legacy stages are unchanged. +Before sending a request, these stages validate their resolved credentials, +custom headers and telemetry headers, including scalar or column credentials +after bearer conversion. Header names +must be nonempty ASCII HTTP tokens. Values may contain horizontal tabs, printable +ASCII or Latin-1 characters in U+00A0-U+00FF; other characters are rejected rather +than silently replaced by the HTTP client. Invalid resolved scalar credentials or +headers fail preflight before processing, without including +the header name or value in the diagnostic. Correct the configuration and run again. +When a higher-priority credential column or a telemetry-header column could override +a scalar value, validation of that conditional fallback waits until the row is resolved. +Null or blank credentials and missing telemetry entries can expose an invalid scalar +fallback; only those rows fail, while rows with valid overrides can run. +Invalid column-bound values instead produce an `invalid_request` row error without +sending that row; valid sibling rows can still run. Credential precedence and ignored +header values are unchanged. Schema checks do not fetch implicit Fabric credentials; +those credentials are validated when building requests. Legacy OpenAI stages are unchanged. +The shared transport adds its platform `User-Agent` afterward; this existing +header, including `MMLSPARK_PLATFORM_INFO`, is outside the stage's validation. +Use representable characters in that environment value to avoid the HTTP +client's existing character replacement. +Their shared `questions` configuration must be non-null and valid before +execution; null cells in `questionsCol` instead produce row-level `missing_input` +errors without sending requests. +Global body-option defaults that do not apply to the selected API are ignored, +including their column bindings. Global API-version defaults +are likewise ignored for native APIs, `/v1` routes and complete `customUrlRoot` +overrides. Dated routes still use their scalar or column-bound version defaults. +An unused global URL is also ignored when `customUrlRoot` supplies the endpoint. +A null cell in a bound API-version column produces `missing_input` without HTTP; +it does not fall back to the built-in version. +Explicit `apiVersion` and `apiVersionCol` bindings are unsupported on native and +`/v1` routes and with complete `customUrlRoot` overrides. For a complete override, +include any required version directly in the URL query instead. Explicit +unsupported body options still fail validation. +Deployment names used in Azure URL paths are encoded as one path segment. They +cannot contain `.` or `..` path components separated by `/` or `\`, or percent +signs that could introduce another decoding layer. The same check applies after +decoding the deployment segment in a complete operation URL. Ordinary encoded +characters remain supported; model names sent only in native or `/v1` request +bodies are not subject to this path restriction. +As with legacy OpenAI stages, v1-base detection ignores case. +AIPrompt also recognizes complete API routes without regard to case, using the raw +path: percent-encoded separators such as `%2F` remain data, not route boundaries. +Classic Azure deployment routes remain dated even when the deployment is named +`v1`. For Azure, supply the resource endpoint plus `deploymentName`, or the +complete operation URL. An incomplete `/openai/deployments/` base passed +to `setUrl` is rejected rather than interpreted as an operation. +`setUrl` removes trailing slashes before adding a route, but preserves path casing; +the endpoint must accept the supplied path. Complete API URLs never receive another route suffix. A +deployment in a complete Azure URL must match the effective deployment: scalar +contradictions fail preflight, and column-value contradictions produce row errors +without HTTP. A complete `/models/chat/completions` URL uses `model`, not +`deploymentName`. These checks also apply to recognizable routes in +`customUrlRoot`; the complete override itself retains its path and query bytes, +including trailing slashes, and never receives an appended route or API version. +In Fabric, Chat Completions, Responses and embeddings automatically use the +Fabric endpoint when no stage URL, complete override or global URL is configured. +This execution-time default is not saved into the stage or global settings. +Native APIs still require a configured endpoint. Outside Fabric, configure an +endpoint with `setUrl` or `setCustomUrlRoot`. +Automatic Fabric credentials are available only for the +trusted Fabric destination, never an unrelated global URL or explicit custom endpoint. +On the implicit Fabric endpoint, requests retain SynapseML's background/flex +workload headers, including when explicit credentials are supplied. Custom and +telemetry headers cannot override that classification, regardless of casing or +scalar/column binding. Explicit endpoints retain their caller-supplied headers. +Unlike legacy OpenAI stages, AIPrompt treats even an explicit `setUrl` or +`setCustomUrlRoot` pointing to the exact managed Fabric endpoint as caller-owned: +configure credentials and workload headers yourself, or leave the endpoint +implicit to use automatic Fabric authentication and classification. +Scalar models use the existing Fabric model-status preflight before inference. +Column-selected models remain distributed and are validated by +the inference endpoint; they are not collected to the driver for preflight. +Schema planning itself never performs that service call. + +HTTP `timeout` is measured in seconds and must be finite and at least `0.001` +(one millisecond); smaller values are rejected before inference. +Very large values saturate at the HTTP client's limit of 2,147,483,647 milliseconds +(about 24.9 days), rather than wrapping into a disabled deadline. +When set, `concurrentTimeout` must be finite and at least `0.000000001` seconds +(one nanosecond). It limits the wait for each next result, starting when the +iterator begins waiting for that result, not when its request was submitted. +Requests overlap, so this is not a per-request wall-clock deadline or a deadline +for the entire batch. Expiry cancels all in-flight workers for that task and fails +the Spark task; it does not produce a `transport_error` row. Spark task retries +may resend requests from the failed task, including ones whose external calls +already completed. If retries are exhausted, the Spark job fails. +The default bounded HTTP handler propagates Spark task cancellation and concurrent +timeouts, including at concurrency one, to its own requests and retry workers, +including during Retry-After backoff. A shared executor-local monitor observes +Spark cancellation even when downstream row processing is blocked, interrupting +workers and starting a scope-owned abort pass without waiting for another row. +A slow abort cannot block the shared monitor or cancellation of unrelated tasks. +Task cleanup waits up to five seconds in total for its abort pass and HTTP workers +to stop. A failed abort wait still attempts to drain HTTP workers within the same +budget; timeout or interruption also interrupts the abort pass. Failure to confirm +quiescence is an execution error, not a successful result. An already-sent provider +request may still complete or incur charges. +Non-fatal abort callback errors are logged but do not fail task cleanup when both +the abort pass and all HTTP workers have stopped. Interruption and fatal callback +failures still fail cleanup after worker draining is attempted. +Explicitly calling `setHandler` with any handler, or replacing the handler +default in a subclass, opts out of +the default transport protections: it selects the legacy HTTP client, whose +automatic retries, redirects and cookie management remain enabled and whose requests are not +registered for cancellation aborts. Leave `handler` unset and its default unchanged +to retain the bounded transport guarantees below. Legacy stages keep their existing transport behavior. +Caller `Content-Length` and `Transfer-Encoding` headers are rejected case-insensitively: +the HTTP client owns request framing. Scalar custom/telemetry headers fail preflight; +column-bound headers produce `invalid_request` for the affected row before HTTP. +The same restriction applies to outgoing authentication headers introduced by subclass +header-name hooks, after credential selection and bearer conversion. Unused hooks do not +invalidate requests, and a column-dependent credential fallback is checked per row. +Other valid caller headers, including `Host`, `Connection`, `Expect`, `TE` and +`Upgrade`, retain their existing behavior. +Saving an AIPrompt stage with a replaced handler default is rejected, including +defaults replaced in subclass constructors. To persist a custom handler, leave the +original default unchanged and use `setHandler`; the explicit binding survives +copy/save/load. Loading a persisted custom handler requires +`spark.synapseml.legacy.allowUnsafeJavaDeserialization=true`. Enable this only for +trusted artifacts: Java deserialization can execute code. Stages using the +unchanged bounded default do not require this opt-in. +Rejecting replaced defaults avoids silently restoring different executable +behavior on load. All ordinary parameter metadata, including defaults hidden by +explicit values, is encoded before a direct stage overwrite deletes an existing +artifact. This includes native JSON parameters and scalar service parameters. +The metadata and parameter decoders must also accept the saved values under the +runtime's Spark JSON parsing limits, so standard Spark pipeline loaders can read +them. Oversized values fail with an actionable error before a direct stage overwrite. +Metadata containing literal unpaired UTF-16 surrogates also fails before saving: +Spark must not replace an invalid string setting with a valid but different prompt. +This covers explicit values and defaults, including shadowed defaults. Valid +supplementary characters, U+FFFD and literal JSON escape text remain unchanged. +Obtaining the writer (`stage.write` in Scala or `stage.write()` in Python) checks +handler defaults and parser configuration; metadata encoding happens when saving. +The guard does not make storage writes atomic or protect an enclosing Spark +`Pipeline` or `PipelineModel` overwrite. To preserve an existing artifact against +storage or complex-parameter serialization failures, save to a new path instead +of overwriting it. +The internal bounded handler is not a public retry-tuning API. +The bounded handler shares its finite retry budget across connection failures, +retryable HTTP statuses and one implicit Fabric authentication refresh. With no +retries configured, a 401 is returned without refreshing. This is intentionally +different from the legacy handler's additional authentication replay. +The default AIPrompt transport disables Apache's internal retries and automatic +redirects. Redirect responses (`300` through `399`) are terminal: they are neither +followed nor retried, and become `http_error` rows with their original status. +Once a redirect status is received, even an I/O failure reading its body or metadata +preserves that status instead of replaying the request. A sanitized body-unavailable +message is logged and exposed through `rawResponseCol`, if enabled. +This avoids repeating an accepted POST or forwarding authentication headers +to a different destination. +Automatic cookie management is also disabled: response `Set-Cookie` headers +are not stored or replayed on later rows or retries. Requests +therefore do not acquire hidden session state from another row's credentials. +An explicitly configured caller `Cookie` header is still sent unchanged. +Legacy HTTP clients retain their existing retry, redirect and cookie behavior. +If token refresh fails, the original 401 becomes a terminal `http_error` row +without another inference request. Interruption and cancellation still propagate. +Successful HTTP responses (`200` through `299`) are terminal once their body and +HTTP metadata have been read completely. Malformed JSON in a fully read response +produces a row diagnostic rather than another inference request. An I/O failure +while reading the body or HTTP metadata, even after a successful status arrives, +can resend the POST within the shared retry budget. The provider may already +have processed it, so these retries can duplicate work and charges. Unlike this +read-failure case, an available redirect status is preserved without replay. + +| Typed request options | Supported APIs | +| --- | --- | +| `maxTokens`, `maxCompletionTokens` (mutually exclusive aliases) | Chat: `max_completion_tokens`; Responses: `max_output_tokens` | +| `temperature`, `topP`, `user`, `verbosity`, `reasoningEffort`, `responseFormat` | Chat, Responses | +| `n`, `stop`, `presencePenalty`, `frequencyPenalty`, `seed` | Chat | +| `store`, `previousResponseId` | Responses | +| `dimensions`, `user` | Embeddings | +| `safetyIdentifier` | Decisions | +| `questions` | System One, Decisions | +| `stateCol` | System One | +| `systemPrompt`, `columnTypes`, `fileSizeLimitMB`, `dropPrompt` | Chat, Responses with generated messages, not `messagesCol` | +| `postProcessing`, `postProcessingOptions` | Chat, Responses | +| Legacy `echo`, `cacheLevel`, `bestOf`, `logProbs` | Explicitly rejected; not silently dropped | + +When configured explicitly or through a parameter default, `fileSizeLimitMB` must +be positive. An explicit value overrides its default. Invalid limits, including NaN, +are rejected during preflight in Scala and Python before HTTP. +A file exceeding a valid limit still produces a row error. For +filesystem attachments, `AIPrompt` reads at most the configured byte limit plus +one byte to detect overflow; it does not first buffer the entire file. Unset or +positive-infinite limits remain unlimited. `AIPrompt` captures the input Spark +session's Hadoop configuration, including session overrides, when `transform` +builds the lazy plan. Configure filesystem access before calling `transform`; +later configuration or credential rotation requires a new `transform` call. +Legacy `OpenAIPrompt` continues using its executor-default Hadoop configuration. +Null, empty or whitespace-only attachment paths contribute no attachment; their +direct filename placeholders render as empty strings. Other usable attachments +and prompt text are preserved when at least one configured attachment is usable. +If every configured attachment is absent, the row produces `missing_input` even +when its prompt text is populated. Null non-path template values retain Spark's +null behavior. + +**Token-limit compatibility:** Both `AIPrompt` and legacy `OpenAIPrompt` Chat +send `maxTokens` as `max_completion_tokens`, not `max_tokens`; +`maxCompletionTokens` uses the same wire field. Responses maps either alias to +`max_output_tokens`. An older dated +API or compatible endpoint may require `max_tokens` instead. For such an +endpoint, leave both typed token options unset and use +`setExtraBody({"max_tokens": 32})` in Python or +`setExtraBody(Map("max_tokens" -> 32))` in Scala. This wire-name override cannot +be combined with an explicitly set typed token option, including a column binding. + +**Responses migration:** Legacy `OpenAIPrompt`/`OpenAIResponses` drops `temperature`, +`top_p` and `seed` when the deployment name contains `gpt-5` (case-insensitive). +`AIPrompt` does not infer model capabilities from deployment names: Responses +forwards explicitly configured `temperature`/`topP` and applicable global defaults +as `temperature`/`top_p`, including for GPT-5-named deployments. `seed` remains +unsupported for Responses and is rejected when explicitly set. +Some reasoning models reject sampling options. Leave unsupported options unset, +and clear inherited sampling defaults with `OpenAIDefaults.resetTemperature()` / +`resetTopP()` in Scala or `reset_temperature()` / `reset_top_p()` on a Python +`OpenAIDefaults` instance before executing. A provider rejection becomes a row +error. + +`AISystemOne` and `OpenAIDecisions` specialize the shared `AIPrompt` implementation +and inherit its setters; they do not expose a separately narrowed Param surface. +Their API selector is locked to their native protocol. The capability table +still applies: setting an inherited generative, embedding, or other-protocol +option fails validation before inference rather than silently ignoring it. +The direct stages keep their documented package names: `AISystemOne` is in +`services.ai`, while the OpenAI-branded `OpenAIDecisions` is in `services.openai`. +Both belong to the cognitive module and reuse the shared facade rather than +independent provider implementations. + +ServiceParams retain last-setter-wins scalar/column semantics. `setModelCol`, +`setStoreCol` and `setResponseFormatCol` are also available. Native outputs and +embedding vectors are not fed through generative-text postprocessing. +`setResponseSchema` configures the wire JSON schema; `setPostProcessing` and its +Spark DDL options configure local text parsing. These are separate contracts. +`setResponseFormat` accepts a schema/format map or the string shorthand `"text"` +or `"json_object"` in both Scala and Python. Shorthands ignore case and surrounding +whitespace; a blank string leaves the current binding unchanged. Pass scalar +schemas as maps, not JSON-document strings. Python constructors and `setParams` +follow the same rules. +For Chat and Responses, `responseFormat` and `responseFormatCol` preserve format +descriptions and provider extension fields while translating between the nested +Chat and flat Responses shapes. Put format metadata inside `json_schema` when +using the nested shape; mixing that wrapper with other top-level format fields, +or putting `type` inside it, is rejected rather than silently discarding fields. +Descriptions in a raw JSON Schema remain inside the schema. +`responseFormatCol` accepts a JSON string or a map with string keys throughout +and JSON-compatible scalar, array and map values. Non-string map keys and +struct, binary, date/time, interval or UDT values are rejected during schema validation, +before inference. Convert richer SQL shapes to an explicit +JSON string first. Invalid format contents within a supported column type remain +per-row request errors. +JSON postprocessing follows Spark's nullable parsing semantics, including when +the supplied DDL contains `NOT NULL`. Planned and actual output schemas make +nested struct fields, array elements and map values nullable; `NOT NULL` is not +an output validation constraint. +JSON object order is preserved, including structured-output property order. +Use a Python dictionary, Scala `ListMap`, or Java `LinkedHashMap` +when the provider should follow a particular property order; an unordered map +cannot establish that order. Copying and saving/loading a stage preserve it. +JSON-string columns and parameters that accept JSON text also preserve object order. +This includes top-level `extraBody` fields, named native questions and their +criteria. JSON decimal precision is preserved for both explicitly set values +and JVM-defined defaults. `setResponseSchema` uses the same JSON value rules as +`setResponseFormat`, including exact JVM decimal values and Python integers +larger than 64 bits through the setter, constructor and `setParams`. Python +dictionaries require string keys at every level. The new APIs accept finite Python +`decimal.Decimal` values in `extraBody`, `responseFormat`, `questions` and +`setResponseSchema`, including constructor and `setParams` inputs where applicable. +Their JSON parameter getters preserve integers as `int` and return `float` when +its JSON representation preserves the number; otherwise they return `Decimal`. +This preserves exact values when passing a getter result back to a setter. +Numeric spelling can change: for example, `2.50` becomes `2.5` and `1e2` +becomes `100.0`. Providers can distinguish the **rendered** integer and decimal +JSON forms, such as `32` and `32.0`. +The JVM stores numbers as `BigDecimal`; a stored scale of zero renders as an +integer. Consequently, input tokens such as `1e0` and `3.2e1` render as `1` and +`32`, and their Python getters return `int`. JVM floating-point conversion can +also yield scale zero: Scala `12345678.0` renders as `12345678`, unlike the Python +float input `12345678.0`. These rendered forms, not the original language type +or token spelling, are sent to the provider. +Use `Decimal("0.12345678901234567890123456789")` for precise Python inputs, or +pass JSON text to a parameter that accepts it, +for example `setExtraBody('{"threshold":0.12345678901234567890123456789,"limit":1e400}')`. +The JVM parses that text without converting its numbers through Python floats. +Converting a value to `float` before passing it can already lose precision; +`float("1e400")` is infinity and is rejected. +Non-finite `Decimal` values are also rejected. Standard Python `json.dumps` +does not serialize `Decimal`; pass getter results directly to these setters +rather than converting them through `float`. Legacy OpenAI stages retain their +existing Python numeric conversion behavior. +Scala and Java `Float` values retain their decimal +representation rather than being widened to a `Double` before JSON encoding. +Loaded Python stages also retain ordinary parameter +defaults in their Python form, including empty dictionaries for `columnTypes` +and `postProcessingOptions`. Attachment parts from +`columnTypes` follow input DataFrame column order, not map insertion order; +reorder the input columns or provide explicit messages to choose a different order. +`columnTypes` keys must also be distinct under Spark's column resolver. With +case-insensitive resolution, `asset` and `ASSET` cannot both be keys, even when +their kinds agree: preflight rejects the duplicate before HTTP. +With case-sensitive resolution, distinct input columns may use those names. +`responseIdCol` exposes an observed response ID; Responses `store` can enable +the default ID column but does not persist a Spark result. + +`setExtraBody` accepts a Python dictionary or JSON-object string, Scala map/`JsObject`/JSON-object +string, or Java map. `setExtraBodyCol` accepts a JSON-object string column; null +means no extension for that row. Unknown provider-owned fields are forwarded +unchanged. Use **wire names**, not Spark setter names. +`extraBody` and `responseFormat` reject null scalar objects and null column names, +including through generic JVM parameter assignment. Nested JSON nulls and null +cells in these optional object columns remain supported. +Inference also accepts an empty-named input column, whether retained or bound +to an input option. Output-role names must be nonempty. + +An extension can replace an implicit default, but cannot override a subtree +owned by an explicitly set typed option, even with an equal value. Disjoint +object siblings merge recursively. Arrays, scalars and nulls are atomic. +Token aliases conflict semantically, not only by spelling. Invalid scalar +configuration fails preflight; invalid row extensions produce row errors. +Duplicate keys, cycles, non-finite numbers and non-object extensions are rejected. +JSON parameter values are limited to 100 nesting levels below their root, +including scalar leaves. Object and question parameters preserve that same limit across Scala/Java +containers, typed JSON values, JSON strings and save/load; the persistence +envelope does not consume a payload level. Scalar response formats are stored +in Chat's normalized shape: the `json_schema` and `schema` wrappers consume two +levels, leaving 98 levels inside the schema. Deeper schemas are rejected before +replacing an existing scalar or column binding, including with `setResponseSchema`. +Responses reply parsing reserves three additional envelope levels for an echoed +`text.format.schema`, without increasing request parameter limits. Duplicate-key, +numeric-token and malformed-JSON checks still apply to the whole reply. +Native request JSON parsing allows strings up to the JVM string-size limit rather +than a runtime-dependent smaller Jackson default, without changing legacy stages' +parsers. Stage persistence is different: saved parameter metadata must satisfy +Spark's JSON parser limits (Jackson's default string limit is 20,000,000 characters). +Oversized parameter metadata is rejected before a direct stage overwrite; keep +large per-row values in input columns rather than saved scalar parameters. +Available memory and provider payload limits still apply during inference. +The native parser's 1,000-character numeric-token +constraint also applies to serialized numbers supplied through Scala/Java +containers and typed JSON setters, so accepted values remain decodable during +save/load. The exponent in scientific-notation serialization must not exceed +2,147,483,647; a short token such as `100e2147483647` is therefore rejected too. +Parser-limit diagnostics distinguish those limits from malformed JSON. + +Identity/input, route, authentication, headers, retry and storage settings cannot +be overridden through the body. Streaming and background lifecycles are rejected. +Embeddings support only absent or `"float"` `encoding_format`, not base64. +Tool requests require `rawResponseCol`; the stage does not execute tools or run +a tool loop. Provider rejection is reported rather than silently removing fields. +Responses output-only `include` options, such as `message.output_text.logprobs` +and `reasoning.encrypted_content`, can be supplied through body extensions; +their returned fields remain available in `rawResponseCol`. A string `file_data` +is literal content, not a path to download, even when it resembles a URL. +Supply the encoding required by the provider; preserving the string does not +guarantee provider acceptance. + +## Output and errors + +Input columns remain in order. A requested output can replace an existing input +column in place. Output, error, usage, raw and response-ID roles must be distinct +under the configured Spark case resolver. Scratch columns never replace inputs +or public outputs. `transform` uses the input dataset's Spark session for planning; +`transformSchema`, which has no dataset, uses the active session and is side-effect free. +Use the execution session as the active session for schema checks so column +resolution and template analysis use the same SQL configuration as execution. +Schema checks do not run Spark jobs or make model calls. Model availability and +row-dependent checks occur during execution. + +`rawResponseCol` contains the complete UTF-8 response body before typed projection, +including unknown fields and all chat choices, except that bodies containing +invalid UTF-8 are discarded. +It never contains HTTP headers. Invalid UTF-8 in a successful HTTP response produces +a terminal `parse_error` without retrying; the answer, raw body, usage and response +ID are null rather than silently repaired text. Non-success HTTP responses retain +their `http_error` status even when the body cannot be decoded. Valid Unicode, +including a correctly encoded replacement character (U+FFFD), is preserved. +Raw responses and retained prompts can contain sensitive user or provider data. +Apply appropriate access controls before displaying or saving them. +JSON strings and object keys must also contain correctly paired UTF-16 surrogates, +including characters decoded from JSON escapes. Lone surrogates are rejected, not +replaced: malformed scalar JSON parameters fail validation, malformed row-bound JSON +produces `invalid_request` without HTTP, and malformed provider JSON produces +`parse_error` with no decoded answer, usage or response ID. An otherwise valid UTF-8 +raw response remains available for diagnosis. Valid supplementary characters are +preserved through requests, responses and persistence. +Prompt templates are checked before Spark constructs string literals. +A literal lone surrogate fails template validation instead +of being repaired and sent as a different prompt or JSON state. +Dated-route API versions are checked before URL encoding: malformed scalar or +default values fail preflight, while malformed column values produce `invalid_request` +without HTTP. Valid Unicode is encoded unchanged; unused version defaults on +non-dated routes remain ignored. +The same checks apply to model or deployment names embedded in deployment URLs: +selected scalar and default values fail preflight, and column values fail per row. +Names used only in request bodies retain per-row validation, and overridden +defaults remain ignored. +URL-related Unicode errors name the selected parameter (`url`, `customUrlRoot`, +`model`, `deploymentName` or `apiVersion`) without including its value. +JSON postprocessing also checks decoded strings and keys, including nested fields +outside the configured schema. Invalid escapes produce `parse_error` rather than a +successful, silently repaired answer. As with other postprocessing conversion errors, +Spark's partial parsed output may remain present; check `error` before using it. +The original raw response is retained, and the existing Spark JSON syntax and read +options are unchanged. +Chat projects the first choice; Responses concatenates assistant text, not tool +arguments. Diagnostics in `error.response` contain stable JSON `code` values: +request/input validation, transport/HTTP failure, provider failure, refusal, +incomplete result, no-text result and parse failure are distinct. Missing text +does not imply safety filtering. Available usage and partial/raw content are +retained; unknown usage is null rather than invented zeroes. +For lazy path attachments in `AIPrompt`, read failures (including filesystem or security +exceptions) produce a sanitized `input_io_error`; +invalid content or attachment settings produce `invalid_request`. Neither sends +an inference request. A storage-client cancellation fails the `AIPrompt` task with a +fixed-message `CancellationException`, without the client's original message or cause. +This is a task failure, not a Spark task kill: Spark retries can repeat earlier +inference in that lazy partition. Legacy `OpenAIPrompt` instead retains its existing +sanitized attachment-error row when neither the thread nor the Spark task is interrupted. +Both stages stop on `InterruptedException` and `ClosedByInterruptException`. +Unlike earlier legacy handling of non-fatal read failures, they also stop when +a read fails while the thread is interrupted or the Spark task is canceled. +Actual thread interruption is preserved; a Spark task-only cancellation does not +set the executor thread's interrupt flag, so cleanup honors `interruptOnCancel=false`. +HTTP and attachment task-kill exceptions preserve Spark's recorded cancellation +reason rather than replacing it with a generic message. The bounded handler also +rechecks cancellation after a rejected worker submission and before reporting +an expired result wait. +Null entries in optional usage-detail maps are omitted without discarding other +counters. When supplied, the main input, output and total token counters must still +be non-negative integers; explicit null or invalid numeric values are parse errors. +Non-null Responses `incomplete_details` marks partial output incomplete even +when the top-level status is absent or inconsistent. +Answer, usage and response ID are validated independently: malformed optional +metadata does not discard a valid answer or other valid observations. Existing +provider diagnostics take precedence over metadata parse errors. This includes +optional `model` metadata for Chat Completions, Responses and embeddings. Native +responses still require a valid model in their typed output. Non-success HTTP +responses retain valid usage and identity, but keep the HTTP status diagnostic. +An established provider failure remains `provider_error` even when its Chat +`choices` or Responses `output` is null or malformed. Valid partial text is +retained in order where it can be decoded; Chat still projects only the first +choice and never substitutes a later choice for an invalid first one. Malformed +answer fields in an otherwise successful response yield a null answer and +`parse_error`, even if another choice or content part was valid. +Optional embedding `object` tags may be absent or null; non-null tags and the +vector's index, values and shape are still validated. +Malformed optional error objects, chat finish/refusal fields and Responses status +fields likewise retain valid answers. Non-null `incomplete_details` still reports +incomplete output even if its details are malformed; successful answer fields remain strict. +Response-body and HTTP metadata read failures preserve a received redirect or +error status. The same status policy still applies: ordinary `4xx` errors are +terminal, while `429`, `5xx` and the managed Fabric authentication refresh use +their existing bounded retry rules. An unreadable body is replaced with a +sanitized diagnostic, not a successful response. Read failures after a `2xx` +status or before a usable status consume the connection retry budget and can +replay an already-processed, billable request. Exhausting those retries produces +`transport_error`. Once a complete response has been buffered, a +connection-close IO failure is logged without discarding the answer or resending +the request. Legacy HTTP handlers retain their existing behavior. +JSON postprocessing also reports `parse_error` for malformed JSON or incompatible +struct fields, even when Spark's permissive parser returns a partially populated +struct. After removing supported code fences, text must contain exactly one +complete JSON value; trailing fragments or another value are errors. +Spark's root-object conversion to a singleton array of structs is preserved, with +the same field-conversion checks as an explicit array. +This check uses Spark's JSON read syntax, not the stricter validation used for +provider request JSON: single-quoted keys and strings, `NaN`/`Infinity` in floating +fields, and duplicate keys (last value wins for struct fields) retain Spark's +default behavior. +A valid object containing only null fields is not a parse failure. +Partial parsed values and raw output remain available; an existing provider +diagnostic is not replaced. +Explicitly disabling postprocessing with `setPostProcessing("")` survives +copy and save/load. Clear that Param before supplying options that infer a mode. + +## Scaling responsibly + +Before inference, tune partition count and `concurrency` to the deployment's rate +limit. Repartition the input DataFrame deliberately; do not collect a production +dataset or loop over its rows on the driver. +Approximate in-flight requests are active Spark partitions times per-partition +concurrency, subject to available executor slots and rows in the batch. +With `handler` unset and its default unchanged, each active partition has its own HTTP worker pool, which +can create up to `concurrency` threads as rows are submitted. The default is `1`; +explicit positive values have no additional library-imposed ceiling and are not +silently clamped to an executor-wide budget. Budget roughly `concurrency` times +simultaneously running partitions per executor for HTTP workers, in addition to +Spark's own threads. Each partition also buffers up to `concurrency` pending +request/response entries; inline media and large responses can make that buffer +expensive. With `dropPrompt=true` (the default), generated messages are not kept +as an extra copy in that buffer. Retaining them with `dropPrompt=false` adds their +payload, including attachments, alongside the request and retained input. +High settings can exhaust container thread/process limits or native +and heap memory, causing executor loss and possible request replay. Do not copy +large values from legacy `OpenAIPrompt` without measuring this per-partition +resource cost. Start at the default, then increase only within provider quota +and executor limits. Responses are buffered in memory; budget for their size as +well as the number of concurrent requests. + +Retained vector map keys are encoded around inference's +row conversions without an extra join. This +preserves SQL-distinct dense and sparse keys even when their mathematical vector +values are equal. Use SQL `map_entries` or `explode` to inspect such keys: +collecting a whole map into a language-level dictionary can apply that +language's key equality and collapse entries independently of Spark storage. +Other UDT map keys, including Spark ML matrices nested in structs or arrays, +are rejected during lazy schema preflight rather than risk silently losing keys. +This does not restrict ordinary retained UDT columns outside map keys. + +The adapter prepares scalar JSON options once and resolves column positions +outside the row loop. The facade uses the shared HTTP transport and one bounded +retry policy, not a second outer retry loop. When a valid Retry-After hint is +present, the retry waits at least that long. Hints longer than 64 characters are +treated as unusable before numeric or date parsing to bound parsing work. +Decimal-second hints are rounded up to the next whole millisecond. Hints above 60 seconds terminate the +bounded request with the original HTTP error rather than sleeping indefinitely +or retrying early. This includes implicit Fabric authentication repair after a +401: a usable hint delays token refresh and replay, and a hint above the cap +returns the rejection without refreshing. Without a usable hint, authentication +repair is immediate. It happens at most once and consumes the same retry budget. +If an error response body cannot be read, a received usable +hint still applies: the HTTP status and hint are preserved with a sanitized +body-read diagnostic, rather than retrying as though no response was received. +Without a usable hint, 429 responses use bounded exponential +backoff with jitter; other retryable failures use the 100/500/1,000 ms schedule. +Unlike the legacy HTTP policy, 429 retries consume this finite retry budget. +There is no failure-rate circuit breaker; other rows can continue after row +errors, so an outage or quota limit can leave many rows without model output. +Inspect the error column rather than treating a completed Spark action as proof +that every inference succeeded. Any retry must account for possible duplicate +provider work when the original request's outcome was uncertain. +Do not claim a throughput improvement without comparing the same workload, +provider latency, partition count, concurrency, errors and physical sends. + +The opt-in `AIPromptBenchmarkSuite` compares the public chat paths using 4,096 +synthetic rows, varying text lengths of 128-1,023 characters, and a recording +HTTP server with 5 ms response latency. One local Spark run measured: + +| Partitions / concurrency per partition | `OpenAIPrompt` rows/s | `AIPrompt` rows/s | +| --- | ---: | ---: | +| 4 / 1 | 56.79 | 81.46 | +| 16 / 4 | 336.74 | 1,074.55 | + +Both paths produced zero row errors and exactly 4,096 physical sends per +measurement. This measures local serialization/transport overhead, not provider +quota or remote-network performance. It is not a +production throughput guarantee. The new path constructs Spark row converters +once per partition, rather than resolving encoders for each request and response. +Reproduce with `SYNAPSEML_AI_BENCHMARK=true` and +`sbt "cognitive/testOnly com.microsoft.azure.synapse.ml.services.ai.AIPromptBenchmarkSuite"`. diff --git a/pipeline.yaml b/pipeline.yaml index 7abb14cb7b0..5ad22545a59 100644 --- a/pipeline.yaml +++ b/pipeline.yaml @@ -830,6 +830,9 @@ jobs: language: PACKAGE: "services.language" FLAKY: "true" + ai: + PACKAGE: "services.ai" + FLAKY: "true" openai: PACKAGE: "services.openai" FLAKY: "true" diff --git a/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-claude-opus-5.5-slice1.md b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-claude-opus-5.5-slice1.md new file mode 100644 index 00000000000..4ba74c7bc2a --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-claude-opus-5.5-slice1.md @@ -0,0 +1,70 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 166 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T10:32:50Z" +outcome: "blocked" +reviewer_model: "claude-opus-5.5" +input_sha256: "869ce6ea6713bde32f7f769eba6a3ab10d19dca1f31a9b11c6c383b463fcbcf7" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: claude-opus-5.5 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-claude-opus-5.5-slice1.md +- **Issues Found**: 0 +- **Verdict**: BLOCKED + +## Evidence Checklist +- [ ] **Blocked: required coverage of the assigned frozen diff was not achieved.** I tried to read the frozen diff with bounded `view` calls for lines 1–150, 151–300, 301–450, 451–600, 601–750 and 751–881, then retried 1–150 once. The tool permission gate rejected all seven calls ("no interactive user response was available"), so no frozen-diff lines came back. Without those reads, the required line 1–881 coverage audit can't pass, and this result can't count as clean. As the plan requires, I didn't try any other way to read it (no search, listing or shell). Next step: re-run this slice with read access to the assigned frozen diff. +- [x] **Slice identity:** slice 1/24, "Cancellation lifecycle and regression coverage". It has four new files (`new file mode`, before is `/dev/null`): + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (+248) + - `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (+10) + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (+568) + - `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (+31) + + The embedded excerpt adds up to 857 content lines plus 24 header lines (6 per file) = 881, which matches the stated length. This is only a consistency check; it doesn't replace the required reads. +- [x] **Worktree copies match the embedded diff:** I read all four files in the worktree in full, and they match the embedded diff line for line. Anchors: + - `CancellableRequests.scala`: object at L20, library context classloader at L31 and L153, kill reason at L94. + - `TaskContextUtils.scala`: L8–10. + - `TaskContextTestSupport.scala`: L13–30. + - The test file has all 568 lines. +- [x] **`CancellableRequests.scala`: provisional review found no concrete defect.** + - **No lost abort:** each request is registered before the cancellation check (L55–57), and `cancelled` is set before the request snapshot (L136–139). So an in-flight request is either aborted or rejected at its check. + - **No hang on dropped work:** `shutdownNow` (L138) drops queued futures, but `await` can't hang on them because `checkTask` runs on every 100 ms poll (L117). + - **Shared monitor never waits on an abort:** the abort runs on its own thread, outside the Scope lock (L141–154). + - **Bounded, single cleanup:** `close` runs once under `cleanupLock`. The abort wait and the worker drain share one 5 s budget. It clears and then restores the thread's interrupt flag, and attaches later failures to the first one as suppressed exceptions (L160–199). + - **Nested failures drain once:** nested `guarded` calls run the drain only once (L208–214). + - **Kill reason wins:** Spark's kill reason takes priority over generic cancellation (L90–100, L123–125). + - **No deadlock:** I found no lock-order cycle among the Scope lock, `cleanupLock` and the executor's internal locks. + - **No sensitive logging:** the abort-failure log includes only the exception class name (L145). +- [x] **`TaskContextUtils.scala`:** a narrow accessor for Spark's `private[spark]` `getKillReason()`. It doesn't handle a null context itself, but every caller passes a non-null one: + - `CancellableRequests.scala` L92–94 uses the captured task context. + - `cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala` L42–44 checks `Option(TaskContext.get())` first. +- [x] **Tests in `VerifyCancellableRequests.scala` (9 test bodies, 11 cases):** anchors at L58, L86, L126, L177, L235, L296, L394/398, L477/481 and L547. Each case would fail if the property it names were broken: + - The monitor and abort threads use the library classloader, while request workers inherit the caller's. + - An iterator closed before any request is submitted sends nothing. + - Closing an abandoned iterator drains every worker. + - Cleanup fails only when workers stay active (cooperative vs. uncooperative abort). + - An unresponsive abort gets exactly one abort pass. + - Workers are still drained after the abort wait is interrupted or the abort fails. + - The shared monitor isn't stalled by a blocking abort or by another scope's drain. + - Nested failures drain only once. + + Every blocking point is released in `finally`, and every thread join has a time limit. The classloader case doesn't depend on test order: `build.sbt` has no `fork` setting and sets `Test / parallelExecution := false` (L282). +- [x] **`TaskContextTestSupport.scala`:** the arguments match Spark 3.5.0's `TaskContextImpl` constructor (L14–15; Spark version at `build.sbt` L31). Because the file is in package `org.apache.spark`, it can call `markInterrupted`, `markTaskCompleted` and `setTaskContext`. The monitor hook (L17) matches the production monitor thread name (`CancellableRequests.scala` L29). +- [x] **Related code outside the slice that I read:** + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala`: L84–90 (`sendCancellableRequestsWithContext`); L113–129, L194, L329–335, L353, L449–454 and L607–636. Those last ranges cover the `check` and `withRequest` call sites, and confirm the bounded retry path re-checks for cancellation before every attempt. + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala` L123–154: the only production consumer, inside `mapPartitions`. Legacy `transform` stays non-cancellable. + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala` L178–213: `toHTTPCore` builds a new request on each call. + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala` L34–52: the `client` and `response` test helpers. + - `scala.util.Using` already compiles on Scala 2.12.17: `core/src/test/scala/com/microsoft/azure/synapse/ml/fabric/FabricCertificateBasedAuthenticator.scala` L31 imports it. +- [x] **Independence:** I didn't open or search any review artifacts, review-run directories, session state, prompts or validation content. I read only product source, tests and `build.sbt`. My only access outside the worktree was the denied bounded reads of the assigned frozen diff. +- [ ] **Nothing executed:** this is a read-only review. I noted the supplied Java 8/17 compile, test-suite and scalastyle results but did not reproduce them. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md new file mode 100644 index 00000000000..314df3f235e --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md @@ -0,0 +1,86 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 166 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T10:17:12Z" +outcome: "blocked" +reviewer_model: "gpt-6.1-sol" +input_sha256: "b351ba1494910185dd36c5d94041b04c463fcec66ea1b94c56c486c0b93b68bb" +--- + +--- +producer: review-code +schema_version: 1 +tier: Max +contract: direct +execution_backend: copilot-cli +attempt: 166 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T10:16:07.050Z" +outcome: blocked +reviewer_model: gpt-6.1-sol +input_sha256: 0f393e4addc8a2fcde73366bf264714f8f79a5a4f851b756d3c4dbcd3354d30f +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: gpt-6.1-sol +- **Artifact**: reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md +- **Issues Found**: 0 +- **Verdict**: BLOCKED + +## Evidence Checklist +- [x] Identified slice 1/24: four added files and 881 assigned frozen diff lines. +- [ ] Mandatory frozen diff reads for ranges 1–150, 151–300, 301–450, 451–600, 601–750, and 751–881 all returned permission denied. Complete before/after coverage could not be established; the embedded excerpt cannot replace this requirement. +- [ ] Coverage remains unverified for all assigned files: `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala`, `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala`, `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala`, and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala`. +- [x] No commands, edits, delegation, publication, or test execution performed. Supplied validation claims were not independently reproduced. + +The required immutable input must be readable before this slice can receive a completed review. Zero reported issues does not indicate a clean result. + +## Driver decision and stage outcome + +- **Effective plan:** Max, pre-commit attempt 166, six ordered stages. Repository + `.github/skills/synapseml-pr-loop/references/review-gauntlet.md` requires six + complete clean rounds. The broad public API, HTTP concurrency, cancellation, + serialization and generated Python surface justify Max; no numeric complexity + score is claimed. +- **Scope and overrides:** 86 files and 1,041,176 patch bytes. The driving human + authorized 24 complete full-file slices, the original 40-minute deadline per + invocation and the existing unlimited-cycle extension. Three reviewers run + concurrently on one slice; all 72 qualified clean results are required per + stage. Any substantive source fix restarts the complete six-stage pass. +- **Identity:** Primary `gpt-6-astra`. The live catalog resolved + `gpt-6.1-sol/max`, `grok-4.7/xhigh` and `claude-opus-5.5/max`; all three runtime + identities matched their assignments. +- **Frozen input:** Source HEAD `5250eb2874612851f9f2b77dbc70dc558cd19b2a`, + target `d3f3af6bd3e4b7d191f51c9eb29f9aeffe788193`, manifest + `4378b4a5fe291a37997a920681e1eccd5e2a67dc280ab98111b7e88a4aabff58`. + Full and 24-slice disposable-index roundtrips passed. +- **Provenance clarification:** The outer skill-owned header records the verified + full-input hash. The preserved nested reviewer header is self-authored and is + not harness provenance. +- **Stage 1, slice 1:** All three results are blocked. The required bounded views + of the assigned frozen diff were denied by the CLI path-permission gate. + Complete embedded text or current-source reads do not replace that requirement. + The three reports contain no concrete product findings; their provisional + analysis does not count as a qualified clean result. +- **Additional scope rejection:** Grok also attempted a dependency-source glob + in two local dependency caches outside the authorized product scope. The tool + denied that request; no cache content was returned. The independence audit + rejects those requested paths rather than treating the denied read as allowed. +- **Disposition:** All responses and transport are retained, and the completed + processes released their evidence files. No further slice was dispatched and + no stage completed. The driving human explicitly approved a dedicated + input-only Git repository containing frozen source diffs, using the installed + runner's `AdditionalRepo` option for future invocations. No broad filesystem + permission, yolo mode, installed-runner edit, model change or deadline change + is authorized. A new transport-only attempt will retain the same source + fingerprint, fix-cycle count and mandatory full-diff coverage audit. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-grok-4.7-slice1.md b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-grok-4.7-slice1.md new file mode 100644 index 00000000000..e5f03ae6166 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-grok-4.7-slice1.md @@ -0,0 +1,42 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 166 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T10:35:19Z" +outcome: "blocked" +reviewer_model: "grok-4.7" +input_sha256: "857fb319631f69521c68b0d14e2ffc7256660d71b52744f58e4a4c646b1e8a36" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: grok-4.7 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-166-stage-attempt-1-review-1-grok-4.7-slice1.md +- **Issues Found**: 0 +- **Verdict**: BLOCKED + +## Evidence Checklist +- [x] Slice 1/24 scope identified from the assignment: `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala`, `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala`, `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala`, and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala`. +- [x] The embedded assigned diff is complete, not truncated: four new-file hunks (`+248`, `+10`, `+568`, `+31`) plus 24 diff headers total 881 lines, matching the stated length. +- [x] Worktree copies of those four files were read and match that embedded before/after text. Related contracts were read in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala`, `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala`, and the cancellation section of `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [ ] Mandatory bounded views of the assigned frozen diff, lines 1–881 in chunks of at most 150, did not succeed. Each view was permission-denied with no interactive approval, so the audited frozen-diff reads are incomplete. Current sources and the embedded diff do not satisfy that coverage gate. This is not a coverage-qualified clean result. + +## Resolution Log +_Updated by the driving agent as findings are addressed._ + +### Driver audit + +This result remains blocked and does not qualify as clean. In addition to the +failed mandatory frozen-diff views, the transport records a dependency-source +glob in two local dependency caches outside the authorized product scope. The +tool denied the request and returned no cache content. The audit rejects those +requested paths; the evidence checklist above does not waive this restriction. +The approved future input-repository access covers only frozen source diffs, +not dependency caches or review feedback. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-claude-opus-5.5-slice1.md b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-claude-opus-5.5-slice1.md new file mode 100644 index 00000000000..4565bc414e8 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-claude-opus-5.5-slice1.md @@ -0,0 +1,75 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 167 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T11:10:27Z" +outcome: "clean" +reviewer_model: "claude-opus-5.5" +input_sha256: "66e57cc13e47ca37afb67ac2efaf3ddf1f0b301ae1ad2b978cd71ef675deb5d1" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: claude-opus-5.5 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-claude-opus-5.5-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] **Diff coverage:** I read slice 1/24 of the assigned frozen diff (881 lines) in six bounded views: lines 1–150, 151–300, 301–450, 451–600, 601–750 and 751–881. No view was truncated. +- [x] **Assigned files:** All four files are `new file mode 100644`, so the before side is `/dev/null` (empty). I read every added line on the after side: + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (+248) + - `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (+10) + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (+568) + - `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (+31) + + The diff has 857 added lines plus 4 × 6 header lines, which equals 881. +- [x] **Worktree check:** The current `CancellableRequests.scala` matches the diff, where the file line number is the diff line number minus 6. I spot-checked lines 60–100 and 155–248. The grep hits at lines 31, 94 and 153 also line up. +- [x] **Callers and dependencies checked:** + - `HTTPClients.scala`: + - 84–90: `sendCancellableRequestsWithContext` calls `CancellableRequests.map`. + - 112–127: bounded `retryPause` calls `check()`. + - 329 and 335: bounded checks. + - 607–637: `bounded` and `boundedWithFabricAuth` use `withRequest`; an `IOException` leads to `check()`. + - `HTTPTransformer.scala` 125–153: the only caller, inside executor `mapPartitions`. + - `AIPrompt.scala` 113–123 and 543: the default bounded handler is what routes to `transformCancellable`. + - `AIRequestAdapter.scala` 228–232: requires concurrency > 0 and a finite `concurrentTimeout` of at least 1e-9, so `await` never gets a non-finite or negative-infinite duration. + - `OpenAIAttachmentUtils.scala` 42–44: checks for a null `TaskContext` before calling `getKillReason`. + - `HTTPSchema.scala` 165–188: `HTTPRequestData` is not final, `toHTTPCore` can be overridden, and the `this(HttpRequestBase)` constructor exists. + - `VerifyBoundedRequests.scala` 34–51: the `private[split1]` `client` and `response` helpers. +- [x] **Concurrency and lifecycle:** + - **Registration vs. abort race:** A request is added to the set before `check()`, and `cancelled` is set before the abort snapshot is taken. A request registered after the snapshot therefore sees the cancellation and throws. + - **Lock order:** `close` holds `cleanupLock` and then briefly takes the scope lock. The monitor takes only the scope lock, and aborts run on their own thread, so there is no inversion. + - **Monitor watch:** `observeCancellation` respects `closed`; `cancel` cancels the watch, and remove-on-cancel drops it from the queue. + - **`close()`:** It clears and later restores the thread's interrupt flag. The abort wait and worker drain share one 5-second budget. + - **Suppressed exceptions:** Nested `guarded` calls add exactly one suppressed cleanup error, and an exception is never suppressed into itself. + - **`await` timing:** It always makes one final zero-wait poll before reporting a timeout. + - **Unstarted work:** Futures dropped by `shutdownNow` never complete, but `checkTask()` throws before anyone waits on them. +- [x] **Behavior matches the stated scope:** + - The task-kill reason is kept, with `"HTTP request task canceled"` used only when no reason is recorded (line 94). + - The monitor thread (line 31) and abort threads (line 153) use the library classloader. Request workers still inherit the caller's classloader. + - `checkTask()` still checks thread interruption first, so that behavior is unchanged. +- [x] **Version compatibility:** + - `build.sbt` uses Scala 2.12.17 and Spark 3.5.0. `scala.util.Using` already appears in existing code (`FabricCertificateBasedAuthenticator.scala:31`), so a backport is on the classpath. + - `build.sbt:282` sets `Test / parallelExecution := false`, so the test hooks on the JVM-wide monitor cannot affect other suites. + - The 11-argument `TaskContextImpl` constructor passes `cpus` explicitly, avoiding any need for `SparkEnv`. + - `getKillReason` is already used at `HTTPSinkV2.scala:142`. +- [x] **Tests:** + - There are 11 cases (1+1+1+2+1+2+1+1+1), which matches the claimed 11/11. + - Timing margins hold: the drain is at most 5 s against 8 s assertions, and the `eventually` windows of 2 s and 5 s sit inside the drain budget. + - Every blocked thread is released in a `finally` block, and threads are joined before their liveness is checked. + - The tests are not vacuous: the abort- and monitor-classloader assertions would fail without the explicit `setContextClassLoader`, and the second-scope cancellation would miss its 2 s limit if the monitor did the abort or drain itself. +- [x] **Security:** The abort-failure log records only the exception's simple class name. Timeout and cancellation messages contain no URLs, headers or credentials, and nothing secret was added. +- [x] **Concerns considered and rejected:** + - A `RejectedExecutionException` race between `checkTask` and `submit`: Spark still reports the task as killed using its recorded reason. + - `require` raising `IllegalArgumentException` for the drain failure: this only affects the message wording. + - Aborts running one at a time: HttpClient forces an immediate (RST) close on abort, so a single abort does not block for long. + - The monitor thread never exits: this would only matter with per-session classloaders, which SynapseML's py4j-based deployments don't use. +- [ ] **Not run:** I executed no builds or tests in this read-only review. I did not re-run the supplied Java 8 and port compile/test/style evidence. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md new file mode 100644 index 00000000000..fe626f1f828 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md @@ -0,0 +1,99 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 167 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T11:01:35Z" +outcome: "clean" +reviewer_model: "gpt-6.1-sol" +input_sha256: "12b5dcbf44641e29df8ff80a18a48f9ddfe1b73f293633dea7cb7333e1e18790" +--- + +--- +producer: review-code +schema_version: 1 +tier: Max +contract: direct +execution_backend: copilot-cli +attempt: 167 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T10:56:12.728Z" +outcome: clean +reviewer_model: gpt-6.1-sol +input_sha256: 0f393e4addc8a2fcde73366bf264714f8f79a5a4f851b756d3c4dbcd3354d30f +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: gpt-6.1-sol +- **Artifact**: reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Reviewed slice 1/24 completely using bounded reads of the assigned frozen diff: lines 1–150, 151–300, 301–450, 451–600, 601–750, and 751–881. All four files are additions; complete before/after coverage includes their `/dev/null` baselines and full added contents. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` lines 1–248: request-registration races, cancellation monitoring, ordered bounded execution, timeout accounting, interruption preservation, abort isolation, cleanup budgeting, and classloader handling. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` lines 1–10 and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` lines 1–31, including recorded kill-reason access and controlled task lifecycle. +- [x] Reviewed `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` lines 1–568, including all 11 generated test cases covering abandonment, abort failures, interrupted cleanup, shared-monitor isolation, and classloader inheritance. +- [x] Cross-checked necessary contracts in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala`, `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala`, and `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala`. +- [x] Checked cancellation-reason integration against `cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala`, public cancellation assertions in `cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptCancellationSuite.scala`, and `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [x] Considered the supplied validation summary without treating reported results as independently verified. No commands or tests were executed. This clean verdict applies only to slice 1/24. + +## Driver decision and stage outcome + +- **Effective plan:** Max, pre-commit attempt 167, six ordered stages. Repository + `.github/skills/synapseml-pr-loop/references/review-gauntlet.md` requires six + complete clean rounds. The broad public API, HTTP concurrency, cancellation, + serialization and generated Python surface justify Max; no numeric complexity + score is claimed. +- **Scope and overrides:** 86 files and 1,041,176 patch bytes. The driving human + authorized 24 complete full-file slices, the original 40-minute deadline per + invocation, at most three concurrent reviewers, and the existing unlimited-cycle + extension. All three models review every slice; a stage requires 72 qualified + clean results. Source fixes restart the complete six-stage pass. +- **Identity:** Primary `gpt-6-astra`. The live catalog resolved + `gpt-6.1-sol/max`, `grok-4.7/xhigh` and `claude-opus-5.5/max`; all three runtime + identities matched their assignments. +- **Frozen input:** Source HEAD `5250eb2874612851f9f2b77dbc70dc558cd19b2a`, + target `d3f3af6bd3e4b7d191f51c9eb29f9aeffe788193`, manifest + `4378b4a5fe291a37997a920681e1eccd5e2a67dc280ab98111b7e88a4aabff58`. + Full and 24-slice disposable-index roundtrips passed. +- **Input access:** This was a same-source transport retry. The driving human + approved a dedicated input-only Git repository containing frozen source diffs, + accessed through the installed runner's `AdditionalRepo` option. Only the exact + assigned diff was authorized outside the product worktree. No yolo, broad + filesystem access, installed-runner edit, tool expansion or deadline change + was used. +- **Provenance clarification:** The outer skill-owned header records the verified + full-input hash. The preserved nested reviewer header is self-authored and is + not harness provenance. +- **Stage 1, slice 1:** All three reviewers successfully read all 881 frozen diff + lines in bounded views. GPT and Opus were independently audited clean, with + 29 and 28 read-only tool calls respectively. Grok reported two findings, but + its checkout-wide mixed Scala/Java/Markdown search returned the path of a prior + review artifact. Its independence audit therefore failed; its result does not + qualify for advancement. +- **Disposition:** All responses and transport are retained, and all completed + processes released their evidence files. The driver independently reproduced + the submission-rejection finding through the production HTTP iterator and + corrected it. The timeout finding's claimed positive-wait control flow and + Spark retry consequence were not supported; its original report contains the + detailed adjudication. This source fix consumes attempt 167. No second slice + or stage was completed, and no result from this attempt qualifies the new patch. + +### Driver follow-up after the source fix + +The corrected submission path passed the primary Java 8 baseline and both Java 17 +Spark 4 compatibility replays: core/cognitive main and test compilation, +12/12 core cancellation tests, 2/2 minimum-timeout controls, 11/11 cognitive +cancellation tests, and all four style gates. No selected tests were skipped. +The original finding artifact contains the exact commands and before/after +regression evidence. Current-head CI and a new complete review pass remain +required; these local results do not establish merge readiness. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-grok-4.7-slice1.md b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-grok-4.7-slice1.md new file mode 100644 index 00000000000..4ef4e115368 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-grok-4.7-slice1.md @@ -0,0 +1,157 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 167 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T11:29:50Z" +outcome: "issues_found" +reviewer_model: "grok-4.7" +input_sha256: "753030fa273a6e46ed48640bf5440a53ca05323dc435b38f0dd34bb47601a17b" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: grok-4.7 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-167-stage-attempt-1-review-1-grok-4.7-slice1.md +- **Issues Found**: 2 +- **Verdict**: ISSUES_FOUND + +## Evidence Checklist +- [x] Slice 1/24 covered from the assigned frozen diff, lines 1-881, in bounded reads of at most 150 lines. All four files are additions; before is empty and after is the full diff. +- [x] `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` reviewed for check, submit, await, cancel, close, and fill, including kill-reason and classloader behavior. +- [x] `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala`, `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala`, and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` reviewed against that lifecycle. +- [x] Cross-file contracts checked in `HTTPClients.scala` (`sendCancellableRequestsWithContext`, `withRequest`, bounded retry checks), `HTTPTransformer.scala` (`transformRequests`), the unified prompting cancellation and kill-reason sections, and `OpenAIAttachmentUtils.scala` use of `TaskContextUtils.getKillReason`. +- [ ] Tests were not executed; this pass is read-only and does not treat claimed prior runs as new proof. + +## Issues + +### Issue 1: Pool shutdown can surface as RejectedExecutionException +- **Severity**: Medium +- **File**: core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala +- **Line(s)**: 102-113, 138, 227-232 +- **Description**: `submit` calls `checkTask()` and then `pool.submit` with no lock held across both steps. `cancel` sets `cancelled` and calls `pool.shutdownNow()` on the monitor thread. The fixed pool uses an unbounded queue, so `submit` rejects only after that shutdown. If the caller is descheduled after `checkTask()` returns, `pool.submit` throws `RejectedExecutionException`. `interruptOnCancel=false` does not interrupt the caller, and neither `submit` nor `fill` maps that rejection back to `TaskKilledException` or `CancellationException`. A cancellation during `input.next()` is still caught by the next `checkTask()`; the hole is the check-then-submit gap, which `fill` hits while prefetching. +- **Risk**: A killed task can fail as a normal task error. Spark may retry it and resend HTTP requests the cancellation was supposed to stop, and the recorded kill reason is lost. Driver close during the same gap becomes `RejectedExecutionException` instead of `CancellationException`. +- **Suggested Fix**: Catch `RejectedExecutionException` in `submit`, call `checkTask()`, and if the scope is cancelled or the pool is shut down throw that cancellation exception rather than the rejection. Add a regression that shuts the scope down between the check and `submit`. + +### Issue 2: Expired wait reports timeout without rechecking cancellation +- **Severity**: Low +- **File**: core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala +- **Line(s)**: 116-125 +- **Description**: `await` calls `checkTask()` before `future.get`, but the `remaining <= 0` timeout branch throws `TimeoutException` immediately. The `ExecutionException` branch calls `checkTask()` before propagating the worker failure. A task-only kill during that final wait, up to `PollNanos` (100ms), is therefore reported as a timeout. +- **Risk**: A cancellation that overlaps timeout expiry becomes a retriable timeout instead of `TaskKilledException` with Spark's recorded reason. The guide treats those as different outcomes, and only the timeout is documented as a retryable task failure. +- **Suggested Fix**: Call `checkTask()` before throwing the timeout, matching the `ExecutionException` path, and cover a kill that arrives during the final poll with a regression. + +## Resolution Log +_Updated by the driving agent as findings are addressed._ + +### Issue 1 +- **Status**: Open +- **What changed**: pending +- **Why**: pending +- **How verified**: pending + +### Issue 2 +- **Status**: Open +- **What changed**: pending +- **Why**: pending +- **How verified**: pending + +## Driver audit and adjudication + +The original response above is preserved. The bounded-view audit confirmed all +881 assigned diff lines were read. However, the checkout-wide +`**/*.{scala,java,md}` search returned the path of +`reviews/pr-2811/task-unknown-attempt-165-stage-attempt-1-review-1-claude-opus-5.5-slice1.md`. +That violates the explicit prior-feedback search exclusion. This result is not +an independent qualifying review, regardless of the findings below. + +### Issue 1: Confirmed diagnostic race; corrected + +The driver reproduced a `RejectedExecutionException` from the production +`sendCancellableRequestsWithContext` iterator when task-only cancellation shut +down its pool during worker construction. The new regression pauses that +construction with an inheritable-thread-local callback and uses a second watch +on the shared monitor to observe completion of the first scope's shutdown. +It does not replace the executor, access private fields, or issue HTTP requests. + +`Scope.submit` now catches only `RejectedExecutionException`, rechecks the task +and scope cancellation state, and otherwise rethrows the original rejection. +The regression checks the recorded `TaskKilledException` reason, zero handler +calls, and the caller's unchanged interrupt flag. All iterator/context cleanup +runs in `finally`. The unified prompting guide documents this cancellation race. + +The broader claim that Spark necessarily reports this as an ordinary retryable +task failure is not established. Spark 3.5.0's +[`Executor.TaskRunner.run`](https://github.com/apache/spark/blob/v3.5.0/core/src/main/scala/org/apache/spark/executor/Executor.scala) +classifies `InterruptedException` or any `NonFatal` failure as `TaskKilled` +when the task has a recorded kill reason. The demonstrated defect and fix concern +the exception emitted by the HTTP iterator, not a proven additional Spark retry. + +Before the fix, on Temurin Java 8: + +```text +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyCancellableRequests -- -z "worker submission" +FAILED: 1/1; expected TaskKilledException, received RejectedExecutionException. +``` + +After the fix, on the same baseline: + +```text +core/compile +core/Test/compile +cognitive/compile +cognitive/Test/compile +PASSED. + +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyCancellableRequests +PASSED: 12/12. + +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyBoundedRequests -- -z "minimum concurrent timeout" +PASSED: 2/2. + +cognitive/testOnly com.microsoft.azure.synapse.ml.services.ai.AIPromptCancellationSuite com.microsoft.azure.synapse.ml.services.openai.OpenAIPromptMultimodalRequestSuite com.microsoft.azure.synapse.ml.services.ai.AIPromptAttachmentSuite -- -z "cancellation" -z "CancellationException" +PASSED: 11/11 across three suites. + +core/scalastyle +core/Test/scalastyle +cognitive/scalastyle +cognitive/Test/scalastyle +PASSED: zero errors and warnings. +``` + +No selected tests were skipped. Port compatibility validation is recorded +separately after completion; these local results are not current-head CI. + +### Issue 2: Not accepted as reported + +The reported final-wait control flow is incorrect. In `Scope.await`, a positive +`waitNanos` implies `remaining > 0`. Its `TimeoutException` takes the `None` path +and recurses through `checkTask()` before another wait or timeout decision. +The branch that throws the concurrent-timeout exception has `remaining <= 0`, +so its `future.get` uses a zero wait, not a final wait of up to 100 ms. + +There is still an ordinary race after the last cancellation observation, but +this does not establish the reported missed positive-wait poll. The claim that +a recorded Spark task kill becomes a retryable task failure is also contradicted +by the executor's recorded-kill fallback described above. No timeout behavior +was changed. Both existing minimum-timeout controls passed: a buffered result +remains available, while pending work times out and is drained. + +This submission fix consumes attempt 167; none of this attempt's results +qualifies the corrected frozen patch. + +### Completed port compatibility replay + +The two changed Scala files matched the preceding frozen source in both isolated +port checkouts before replay, and matched the corrected source after validation. +On Java 17, both Spark 4.0.1/Scala 2.13.16 and Spark 4.1.1/Scala 2.13.17 passed +the exact compile, test and style commands listed above: four compile gates, +12/12 cancellation tests, 2/2 minimum-timeout controls, 11/11 cognitive +cancellation tests, and four style gates with zero errors or warnings. +No tests were skipped. These are focused compatibility replays, not full port CI. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-claude-opus-5.5-slice1.md b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-claude-opus-5.5-slice1.md new file mode 100644 index 00000000000..5ec698c3f3a --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-claude-opus-5.5-slice1.md @@ -0,0 +1,72 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 168 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:02:39Z" +outcome: "clean" +reviewer_model: "claude-opus-5.5" +input_sha256: "6d6b45048e764405770f769d8e031301b4fb20f3e49eb8b909ad280c2af3769d" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: claude-opus-5.5 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-claude-opus-5.5-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] **Coverage of slice 1/24 (Cancellation lifecycle and regression coverage).** I read the assigned frozen diff in seven bounded views: lines 1–150, 151–300, 301–450, 451–600, 601–750, 751–900 and 901–938. Every line from 1 to 938 came back without truncation. The total matches the file contents: 4 files × 6 header lines, plus 252 + 10 + 621 + 31 added lines, equals 938. +- [x] **Before and after content.** All four assigned files are new (`new file mode 100644`, `--- /dev/null`, hunks of the form `@@ -0,0 +1,N @@`). There is no "before" version, and I reviewed each file's complete "after" content: + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (252 lines) + - `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (10 lines) + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (621 lines) + - `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (31 lines) +- [x] **Contracts with other files (read-only).** + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala`: + - lines 62–90: the cancellable client and `sendCancellableRequestsWithContext`, including the comment that callers outside a Spark task must close the iterator. + - line 120: the cancellation check in `retryPause`. + - lines 329, 335 and 353: the checks in `sendWithRetryPolicy`. + - lines 618, 624 and 630: `boundedWithFabricAuth`, which calls `withRequest` and `check`. + - `HTTPTransformer.scala` lines 107–152: `transformCancellable` and `concurrentWait`. + - `Clients.scala`: `RequestWithContext` and the ordering used by the older `AsyncClient`. + - `VerifyBoundedRequests.scala` lines 34–51: the `client` and `response` test fixtures. + - `cognitive/.../openai/OpenAIAttachmentUtils.scala` lines 37–48: this caller of `TaskContextUtils` only runs when `TaskContext` is non-null. + - `build.sbt`: Spark 3.5.0, Scala 2.12.17, `Test / parallelExecution := false`. +- [x] **Concurrency and lifecycle in `CancellableRequests.scala`.** + - **Lock ordering:** `observeCancellation` (lines 80–88) and `cancel` (137–162) take the scope lock. `close` (164–203) takes `cleanupLock` and holds the scope lock only briefly inside `cancel`. The bounded wait for workers to stop never holds the scope lock, so one scope that is shutting down cannot stall the shared monitor for other scopes. + - **Listener vs. watch race:** if the task-completion listener runs immediately on an already-completed task, the `closed` flag prevents the watch from being scheduled (lines 209–210). + - **Request registration vs. abort snapshot:** `withRequest` (53–60) registers the request and then checks for cancellation. `cancel` sets `cancelled` before calling `shutdownNow` and taking the snapshot (140–143). Under the memory model, every in-flight request is either in the snapshot or fails its own check. + - **Queued work that never starts:** the `checkTask` poll every 100 ms in `await` (119–135) stops the wait on futures that `shutdownNow` drained. + - **Bounded cleanup:** abort and worker shutdown share one 5-second budget. + - **Interrupt flag:** `close` clears the caller's interrupt flag at line 168 and restores it at line 200. +- [x] **Kill reason and rejected submissions.** + - `checkTask` (90–100) checks in this order: + - a thread interrupt raises `InterruptedException`, which Spark maps back to the recorded kill reason; + - a task kill raises `TaskKilledException` carrying `getKillReason`, with the generic message only as a fallback; + - a scope cancellation raises `CancellationException`. + - `await` repeats this check before rethrowing the cause of an `ExecutionException` (127–129). + - The `RejectedExecutionException` path in `submit` (102–117) is the race the latest change targets. I traced it through `ThreadPoolExecutor.execute` and `addWorker`. Worker construction happens before `mainLock` is taken, so a concurrent `shutdownNow` leads to `addWorkerFailed` and then `reject`. The path then rechecks cancellation and otherwise rethrows the original rejection, as the spec requires. +- [x] **Regression tests.** I traced all 12 tests (10 `test` blocks, two of them parameterized ×2) against the production code. The count matches the supplied 12/12 claim. + - **Submission race (lines 78–129):** `InheritableThreadLocal.childValue` runs inside `Executors.defaultThreadFactory().newThread`. Because the monitor is single-threaded, the observer watch is guaranteed to run after `cancel()` finishes. + - **Abandoned iterator (lines 150–200):** a fixed pool always adds a new worker while below its core size, so `workers.size() == 2` is deterministic. The bounded join at line 186 runs before the liveness assertions. + - **Uncooperative workers:** a cleanup failure surfaces through the message of `TaskCompletionListenerException`. + - **Thread cleanup:** every worker, abort and consumer thread is joined or shut down in a `finally` block. +- [x] **Spark API use in the test support and accessor.** + - `TaskContextTestSupport.scala` uses the 11-argument Spark 3.5 `TaskContextImpl` constructor and passes `cpus = 1` explicitly, so `SparkEnv` is not needed. It also uses `markInterrupted`, `markTaskCompleted` and `setTaskContext`; these are Spark-internal and are only called from code in the `org.apache.spark` package. + - `TaskContextUtils.scala` line 9 follows the existing injection objects (`UDFUtils`, `SConf`, `RegressionUtils`, `BlockManagerUtils`). `HTTPSinkV2.scala` already calls `getKillReason()` in the same way. +- [x] **Security.** Abort failures log only the exception's class name (no URL, headers or body). All threads are daemon threads. The change adds no credentials, endpoints or RDD use. The only new public JVM API is the `TaskContextUtils` accessor. +- [x] **Considered and dismissed, with no reproducible defect:** + - **(a) Errors from workers skip cleanup in `guarded`:** throwables that are not `Exception`s bypass it, but inside a task the completion listener still closes the scope, and the driver close requirement is documented in `HTTPClients.scala`. + - **(b) Mismatched exception types:** a failed worker shutdown raises `IllegalArgumentException` (from `require`), while a failed abort raises `IllegalStateException`. The former only appears suppressed or wrapped on cleanup paths. + - **(c) Permanent monitor thread:** the shared monitor is created lazily and lives as a permanent daemon thread on the library classloader. This is the intended design, and it does not retain the task classloader. + - **(d) Overflow in `await`:** the arithmetic only overflows for timeouts near −2^63 ns, which the validated callers cannot pass. +- [ ] **Not applicable — no execution evidence from this review.** Under the read-only contract I ran no build, test or scalastyle command. I cross-checked the supplied Java 8 and Spark port results against the source (for example, that 12 tests exist) but did not rerun them. +- [x] **Independence.** I did not access `reviews/`, session state or other reviewers' output. All searches were limited to `core/src` and `cognitive/src`, plus targeted reads of `build.sbt`. The only reads outside the worktree were bounded views of the assigned frozen diff. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md new file mode 100644 index 00000000000..dacca732572 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md @@ -0,0 +1,105 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 168 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T11:48:44Z" +outcome: "clean" +reviewer_model: "gpt-6.1-sol" +input_sha256: "807962a438b7c409f88ff4fcb9566da93960c2499689da44a0687da7e065b76f" +--- + +--- +producer: review-code +schema_version: 1 +tier: Max +contract: direct +execution_backend: copilot-cli +attempt: 168 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T11:44:43.245Z" +outcome: clean +reviewer_model: gpt-6.1-sol +input_sha256: f677a5fdd658816937af7c3465db15a2cacf78f966dfff72a5bb61d5831927a5 +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: gpt-6.1-sol +- **Artifact**: reviews\pr-2811\task-unknown-attempt-168-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Read the complete assigned frozen diff, slice 1/24, through bounded ranges 1–150, 151–300, 301–450, 451–600, 601–750, 751–900, and 901–938. All four files are additions; complete before/after coverage established against `/dev/null`. +- [x] Reviewed `core\src\main\scala\com\microsoft\azure\synapse\ml\io\http\CancellableRequests.scala` and `core\src\main\scala\org\apache\spark\injections\TaskContextUtils.scala`, including submission races, cancellation reasons, interruption preservation, request registration, bounded cleanup, and infrastructure classloaders. +- [x] Reviewed all 12 regression cases in `core\src\test\scala\com\microsoft\azure\synapse\ml\io\split1\VerifyCancellableRequests.scala` and the complete `core\src\test\scala\org\apache\spark\TaskContextTestSupport.scala`. +- [x] Traced necessary contracts through HTTP client dispatch, transformer iteration, native AIPrompt validation, attachment cancellation, related regression tests, and `docs\Explore Algorithms\OpenAI\Unified AI prompting.md`. No concrete correctness, security, or specification defect identified. +- [x] Considered supplied validation evidence without independently executing commands or tests. This clean verdict qualifies only slice 1/24, not the complete PR. + +## Driver decision and stage outcome + +- **Effective plan:** Max, pre-commit attempt 168, six ordered stages. Repository + `.github/skills/synapseml-pr-loop/references/review-gauntlet.md` requires six + complete clean rounds. The broad public API, HTTP concurrency, cancellation, + serialization and generated Python surface justify Max; no numeric complexity + score is claimed. +- **Scope and overrides:** 86 files and 1,043,725 patch bytes. The driving human + authorized 24 complete full-file slices, the original 40-minute deadline per + invocation, at most three concurrent reviewers, and the existing unlimited-cycle + extension. All three models review every slice; a stage requires 72 qualified + clean results. Source fixes restart the complete six-stage pass. +- **Identity:** Primary `gpt-6-astra`. The live catalog resolved + `gpt-6.1-sol/max`, `grok-4.7/xhigh` and `claude-opus-5.5/max`; all three runtime + identities matched their assignments. +- **Frozen input:** Source HEAD `5250eb2874612851f9f2b77dbc70dc558cd19b2a`, + target `d3f3af6bd3e4b7d191f51c9eb29f9aeffe788193`, manifest + `ff746b509c02dc5a452398374537eb81254984cef2087af817427ce88ce21123`. + Full and 24-slice disposable-index roundtrips passed. +- **Input access:** The human-approved input-only Git repository contained frozen + source diffs, accessed through the installed runner's `AdditionalRepo` option. + Only the exact assigned diff was authorized outside the product worktree. + No yolo, broad filesystem access, installed-runner edit, tool expansion or + deadline change was used. +- **Provenance clarification:** The outer skill-owned header records the verified + full-input hash. The preserved nested reviewer header is self-authored and is + not harness provenance. +- **Stage 1, slice 1:** All three reviewers read the complete 938-line frozen + diff. GPT and Opus were independently audited clean, with 37 and 22 read-only + tool calls respectively. Grok reported one finding but made five source-only + or exact-build-file searches without the explicitly required subdirectory + paths. Its stricter search-scope audit failed. No prior-feedback result was + identified in those five searches, but the recorded requirement is not + retroactively waived. +- **Disposition:** All original responses and transport are retained, and all + completed processes released their evidence files. The driver independently + reproduced the narrowed expired-wait diagnostic race and corrected it. + This source fix consumes attempt 168. No second slice or stage completed; + no result from this attempt qualifies the corrected patch. + +### Driver follow-up after the source fix + +The corrected expired-wait path passed the primary Java 8 baseline and both +Java 17 Spark 4 compatibility replays: four core/cognitive compile gates, +13/13 core cancellation tests, 2/2 minimum-timeout controls, 11/11 cognitive +cancellation tests, and all four style gates. No selected tests were skipped. +The original finding artifact contains exact commands and before/after evidence. +Current-head CI and a new complete review pass remain required. + +For future passes, the driver restores the earlier bounded root-query policy: +fixed product-subtree patterns, exact root filenames, or source-only filename +extensions may qualify. Broad Markdown searches remain forbidden, as do reads +outside the product worktree except for the exact assigned frozen diff. The +transport audit also checks every search result for review or validation artifact +paths. This focuses independence checks on permitted source scope and actual +returned content rather than requiring the optional `paths` argument for every +source-only lookup. Attempt 168's stricter recorded contract and failed result +are unchanged. Coverage, models, reasoning, deadlines, concurrency, stage gates +and filesystem permissions are unchanged. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-grok-4.7-slice1.md b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-grok-4.7-slice1.md new file mode 100644 index 00000000000..09cacc2aec5 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-grok-4.7-slice1.md @@ -0,0 +1,132 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 168 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:07:09Z" +outcome: "issues_found" +reviewer_model: "grok-4.7" +input_sha256: "77a2aafe10abdfffb92389c64ab0000afe77db7bea94917f038e1a8b0ee2c9c8" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: grok-4.7 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-168-stage-attempt-1-review-1-grok-4.7-slice1.md +- **Issues Found**: 1 +- **Verdict**: ISSUES_FOUND + +## Evidence Checklist +- [x] Slice 1/24 assigned frozen diff read in full (lines 1–150, 151–300, 301–450, 451–600, 601–750, 751–900, 901–938); before/after coverage is a new-file addition for all four paths. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (scope lifecycle, submit rejection, await, abort, drain, iterator) against the cancellation guide in `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` and confirmed `getKillReason()` is the same public `TaskContext` accessor already used by `core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/HTTPSinkV2.scala`. +- [x] Reviewed `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` for classloader, submission-race, drain, shared-monitor, and kill-reason assertions. +- [x] Checked cross-file contracts in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala` (`sendCancellableRequestsWithContext`, `withRequest`) and `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala` (per-partition iterator plus task-completion close). +- [ ] Tests were not executed in this read-only review. + +## Issues + +### Issue 1: Expired concurrent-timeout wait drops a task kill +- **Severity**: Low +- **File**: core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala +- **Line(s)**: 120-129 +- **Description**: `await` calls `checkTask()` and then, when `future.get` times out with `remaining <= 0`, throws `TimeoutException` without checking again. The `ExecutionException` branch immediately below does recheck. A task kill or scope cancellation that becomes visible after that pre-wait check and before the throw is reported as a concurrent timeout. For a still-positive remaining wait the recursive call rechecks, so this gap is only the already-expired branch (including the allowed 1ns timeout, where that branch is the normal final poll). +- **Risk**: Spark treats `TimeoutException` as a task failure and can retry the partition, resending provider calls after an intentional cancel. `TaskKilledException` would preserve the recorded kill reason and avoid that retry classification. The window is the gap after line 121’s `checkTask()` during the zero-wait `get`, so it is timing-dependent rather than the common cancel path. +- **Suggested Fix**: Call `checkTask()` before throwing the timeout exception so a recorded kill or scope cancellation wins, matching the `ExecutionException` path. Keep the existing timeout exception only when that recheck does not throw. + +## Resolution Log +_Updated by the driving agent as findings are addressed._ + +### Issue 1 +- **Status**: Open +- **What changed**: pending +- **Why**: pending +- **How verified**: pending + +## Driver audit and resolution + +The original response above is preserved. All 938 assigned frozen diff lines +were read. Five searches, however, omitted the explicit product-subdirectory +paths required for this attempt. They used Scala-only patterns or the exact +`build.sbt` filename and did not reveal prior feedback, but they fail the +recorded stricter scope contract. This is not a qualifying independent result. + +### Confirmed iterator diagnostic race; corrected + +The driver independently reproduced the narrowed zero-wait case through +`sendCancellableRequestsWithContext`. A controlled task context publishes its +kill immediately after the initial interruption read, returning the value +observed before the kill. The shared monitor is held at a bounded latch, and a +blocked request is awaited with the supported one-nanosecond timeout. The +regression therefore exercises the real expired `future.get` path rather than +replacing a future, executor, or production cancellation method. + +Before the fix, on Temurin Java 8: + +```text +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyCancellableRequests -- -z "expired HTTP waits" +FAILED: 1/1; expected TaskKilledException, received TimeoutException. +``` + +`Scope.await` now calls `checkTask()` immediately before reporting an expired +result wait. The ordinary timeout exception remains unchanged when cancellation +is not observed. The regression checks the recorded kill reason and the caller's +unchanged interrupt flag; cleanup releases the monitor and worker and joins the +worker in `finally`. The guide describes this additional cancellation check. + +This confirms an iterator-level diagnostic race, not the report's claimed extra +Spark retry. Spark 3.5.0's +[`Executor.TaskRunner.run`](https://github.com/apache/spark/blob/v3.5.0/core/src/main/scala/org/apache/spark/executor/Executor.scala) +classifies a `NonFatal` failure as `TaskKilled` when the task has a recorded kill +reason. The earlier claim of a missed positive wait of up to 100 ms was also +incorrect: positive-wait timeouts recurse through `checkTask`. This narrower +zero-wait finding supplied the distinct boundary now covered by the regression. + +The complete after-fix validation and port replay results are recorded below +when finished. This source fix consumes attempt 168; none of its results +qualifies the corrected frozen patch. + +### Primary after-fix validation + +On Temurin Java 8: + +```text +core/compile +core/Test/compile +cognitive/compile +cognitive/Test/compile +PASSED. + +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyCancellableRequests +PASSED: 13/13, including the new expired-wait regression. + +core/testOnly com.microsoft.azure.synapse.ml.io.split1.VerifyBoundedRequests -- -z "minimum concurrent timeout" +PASSED: 2/2; buffered results and ordinary pending-work timeouts retain their behavior. + +cognitive/testOnly com.microsoft.azure.synapse.ml.services.ai.AIPromptCancellationSuite com.microsoft.azure.synapse.ml.services.openai.OpenAIPromptMultimodalRequestSuite com.microsoft.azure.synapse.ml.services.ai.AIPromptAttachmentSuite -- -z "cancellation" -z "CancellationException" +PASSED: 11/11 across three suites. + +core/scalastyle +core/Test/scalastyle +cognitive/scalastyle +cognitive/Test/scalastyle +PASSED: zero errors and warnings. +``` + +No selected tests were skipped. These local results are not current-head CI. + +### Completed port compatibility replay + +The two changed Scala files matched the preceding frozen source in both isolated +port checkouts before replay, and matched the corrected source after validation. +On Java 17, both Spark 4.0.1/Scala 2.13.16 and Spark 4.1.1/Scala 2.13.17 passed +the exact compile, test and style commands listed above: four compile gates, +13/13 cancellation tests, 2/2 minimum-timeout controls, 11/11 cognitive +cancellation tests, and four style gates with zero errors or warnings. +No tests were skipped. These are focused compatibility replays, not full port CI. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-claude-opus-5.5-slice1.md b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-claude-opus-5.5-slice1.md new file mode 100644 index 00000000000..14b237b39a2 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-claude-opus-5.5-slice1.md @@ -0,0 +1,60 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 169 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:39:08Z" +outcome: "clean" +reviewer_model: "claude-opus-5.5" +input_sha256: "82c91d9290e2942b75aa8d16cf9d8afa1304f4d4f7896d28fb2c1ff969db8f13" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: claude-opus-5.5 (reasoning: max) +- **Artifact**: reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-claude-opus-5.5-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] **Diff coverage, slice 1/24.** I read the assigned frozen diff (1,001 lines) with seven bounded views: 1–150, 151–300, 301–450, 451–600, 601–750, 751–900 and 901–1001. No range was truncated. All four files are new (`/dev/null` before-side), so the before-state is empty and every added line was reviewed: + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (+254, diff lines 1–260) + - `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (+10, diff lines 261–276) + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (+682, diff lines 277–964) + - `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (+31, diff lines 965–1001) +- [x] **Locking and concurrency in `CancellableRequests.scala`:** + - `cancel()` and `observeCancellation()` use the scope monitor (L80–88, L139–164). `close()` uses a separate `cleanupLock` (L166). No path takes the two locks in opposite order, so they cannot deadlock. + - The shared monitor never waits on a drain, and request aborts run on a dedicated thread (L156–160). + - A request registered in `withRequest` (L55–57) is always either aborted or rejected. Registration is followed by a volatile check of `cancelled`, while `cancel()` sets `cancelled` before taking its request snapshot (L142–145). At least one side must see the other's write. +- [x] **Cancellation and timeout paths:** + - Futures dropped by `shutdownNow` cannot hang `await`, because `checkTask()` runs on every poll of at most 100 ms (L121, L124–135). + - A rejected submission rechecks cancellation before rethrowing (L112–115). An expired wait also rechecks before reporting a timeout (L125–127). + - Task kills keep Spark's recorded reason and fall back to the generic message only when none exists (L93–94). + - `close()` clears and restores the caller's interrupt flag (L170, L202). The abort wait and worker drain share one five-second budget (L172). After a failed abort wait, the abort pass is interrupted and the workers are still drained (L185–198). + - `guarded` closes the scope only once, and cleanup failures are attached as suppressed exceptions (L214–220). +- [x] **Callers and inputs (contracts outside the slice):** + - `HTTPClients.scala` L84–90 documents the close contract. L113–127, L329–353 and L612–637 show bounded retries checking cancellation and registering requests through `withRequest`. + - `HTTPTransformer.scala` L110–118 rejects concurrency below 1 before a scope or thread pool is created. L125–155 contains the cancellable `mapPartitions` path. + - `AIPrompt.scala` L113–115 and L543 send only the default bounded handler down the cancellable path. + - `OpenAIAttachmentUtils.scala` L42–44 calls `TaskContextUtils` only with a context already checked for null. +- [x] **`TaskContextUtils` and the test helper.** + - `TaskContextUtils` delegates to Spark's `private[spark] getKillReason()`. This follows the existing pattern in the `org.apache.spark.injections` package and in `HTTPSinkV2`. + - `TaskContextTestSupport` calls the `TaskContextImpl` constructor with the correct arguments for Spark 3.5.0 (`build.sbt` L31). It passes `cpus` explicitly so no `SparkEnv` is needed. +- [x] **All 13 test cases traced.** + - Submission-time kill: creating the worker thread triggers the kill. The pool is shut down while the thread is being added, the submission is rejected, and the recorded kill reason is reported. `addWorker` creates the thread outside `mainLock`, so this cannot deadlock. + - The expired-wait race test uses stack-frame gating on `CancellableRequests$Scope.await`. + - Both shared-monitor isolation tests depend only on the monitor being single-threaded, so they are deterministic. + - Assertions are not vacuous, and the timing bounds (5 s drain against 8 s and 2 s limits) are consistent. +- [x] **Spec and build checks.** The behaviour matches `docs/Explore Algorithms/OpenAI/Unified AI prompting.md` L437–458 (timeout semantics, five-second cleanup budget, logged non-fatal abort errors) and L781–786 (reason preservation, `interruptOnCancel=false`, both rechecks). The build uses Scala 2.12.17 (`build.sbt` L34). `scala.util.Using` is already used in pre-existing main code (`OpenAIPrompt.scala`). `Test / parallelExecution := false` (L282). +- [x] **Considered and not raised:** + - The worker-drain failure uses `require` (L197–198), so it surfaces as `IllegalArgumentException`, while the abort timeout throws `IllegalStateException`. No caller treats the type differently: the `IllegalArgumentException` catches in `AIPrompt` L140 and L442 run during driver-side plan building and per-row request preparation, not around the HTTP iterator. + - Abandoned driver-side iterators that are never closed leave idle pool threads; the `HTTPClients.scala` L84 comment already documents that callers must close them. + - Inherited thread-local values on the monitor thread are explicitly out of scope in the stated design. +- [ ] **Tests not executed.** This was a read-only review; no commands were run. The pass counts in the plan were not reproduced by this reviewer. +- [x] **Independence.** I did not open, list or search `reviews/`, sibling feedback, run directories or `.validation`. The only read outside the product worktree was the assigned frozen diff. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md new file mode 100644 index 00000000000..faeb2fa08b6 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md @@ -0,0 +1,46 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 169 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:27:35Z" +outcome: "clean" +reviewer_model: "gpt-6.1-sol" +input_sha256: "643f652c5c74f4cbba87c240956ed84c0a1f8095e7673d672bd4445fd6e15e49" +--- + +--- +producer: review-code +schema_version: 1 +tier: Max +contract: direct +execution_backend: copilot +attempt: 169 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:22:20.598Z" +outcome: clean +reviewer_model: gpt-6.1-sol +input_sha256: f7ffba2e9552f1f03381a6ea1e46c94b7e0b41454e3d7f9306e627d969e9dd56 +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: gpt-6.1-sol +- **Artifact**: reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Reviewed slice 1/24 completely: the assigned frozen diff was read in bounded ranges `1-150`, `151-300`, `301-450`, `451-600`, `601-750`, `751-900`, and `901-1001`. All four files are additions; complete before/after coverage establishes an empty baseline for each. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala:1-254`: bounded submission, cancellation races, timeout handling, request registration, abort isolation, cleanup budgeting, exception suppression, interruption preservation, and classloader selection. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala:1-10`: recorded task-kill reason access and its HTTP/attachment call-site contracts. +- [x] Reviewed `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala:1-682` and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala:1-31`, including all 13 regression cases, bounded synchronization, worker cleanup, and task-context restoration. +- [x] Cross-checked integration in `HTTPClients.scala`, `HTTPTransformer.scala`, and `AIPrompt.scala`; inspected related timeout and public Spark-cancellation tests and the cancellation documentation in `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [x] Considered the caller-supplied validation summary as reported evidence only. No commands or tests were executed. This clean verdict qualifies only the assigned slice, not the complete PR. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-grok-4.7-slice1.md b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-grok-4.7-slice1.md new file mode 100644 index 00000000000..9e3eb2d4e12 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-grok-4.7-slice1.md @@ -0,0 +1,38 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 169 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:41:14Z" +outcome: "clean" +reviewer_model: "grok-4.7" +input_sha256: "8bc6c8f654fa086871be753d37c85fc9c707d6cc7bd65423e47fd973ca70f9a7" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: grok-4.7 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-169-stage-attempt-1-review-1-grok-4.7-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Assigned frozen diff lines 1–1001 read in bounded views (1–150, 151–300, 301–450, 451–600, 601–750, 751–900, 901–1001), covering all four slice files before and after. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (254 lines): task-kill versus thread-interrupt ordering, rejected-submission recheck, expired-wait recheck, request registration versus abort snapshot, shared-monitor isolation, five-second shared drain, and library classloader on monitor and abort threads only. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` against the Spark 3.5 `TaskContext.getKillReason` / `TaskContextImpl` contract used by the kill-reason path. +- [x] Reviewed `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (682 lines) for false-green coverage of classloader inheritance, submission rejection, expired-wait races, idempotent close, single abort pass, and cross-scope monitor isolation. Tests were not executed. +- [x] Checked cross-file contracts in `HTTPClients.scala` (`sendCancellableRequestsWithContext`, `withRequest`), `HTTPTransformer.scala` (per-result timeout and task-completion close), and the cancellation section of `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. No spec deviation found. + +## Driver qualification and disposition + +The original response above is preserved. This invocation does **not** qualify as independent review: checkout-wide Markdown searches violated its recorded source-only search boundary and returned prior review paths and feedback. A clean self-reported verdict cannot clear that failure. + +The other two reviewers completed clean, independently audited reads of this slice. No source defect was reported and the frozen product patch did not change. The complete review pass remains unfinished; these results are not a merge-ready sign-off. + +The next transport-only pass will require exact `view` calls only, including for repository guidance. Its audit will reject every search tool call. This is a prospective restriction, not a waiver of this attempt's failure. The three model families, full diff coverage, six-stage gate, forty-minute invocation deadlines and existing permissions remain unchanged. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-claude-opus-5.5-slice1.md b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-claude-opus-5.5-slice1.md new file mode 100644 index 00000000000..0baf7c941b7 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-claude-opus-5.5-slice1.md @@ -0,0 +1,69 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 170 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T13:14:44Z" +outcome: "clean" +reviewer_model: "claude-opus-5.5" +input_sha256: "dba01f2e02b1f3a5cab351d86eb04801a957be44a6e55ff799acc51352ce2530" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: claude-opus-5.5 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-claude-opus-5.5-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] **Slice identity.** Slice 1/24 (Cancellation lifecycle and regression coverage), source HEAD `5250eb2874612851f9f2b77dbc70dc558cd19b2a`, merge base `d3f3af6bd3e4b7d191f51c9eb29f9aeffe788193`, slice SHA-256 `f7ffba2e…`, 4 files. +- [x] **Full diff coverage.** I read all 1,001 lines of the assigned frozen diff with bounded `view` ranges: 1-150, 151-300, 301-450, 451-600, 601-750, 751-900 and 901-1001. I re-read the test portion (274-423, 424-573, 574-723, 724-873, 874-1001). No range was truncated or inaccessible. +- [x] **Before and after content.** All four files are new (`new file mode 100644`, `--- /dev/null`), so the before side is empty and the full after side is in the diff: + - `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (+254) + - `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (+10) + - `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (+682) + - `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (+31) + + The worktree copies of both production files match the diff exactly. +- [x] **Guidance read by exact path.** `AGENTS.md`, `.github/skills/synapseml-branches/SKILL.md`, `.github/skills/synapseml-branches/references/branch-spark3p5.md`, `.github/skills/code-review/SKILL.md`, and `build.sbt` (Spark 3.5.0, Scala 2.12.17). I used only the `view` tool and did not read `reviews/`, `.validation`, session state or any other reviewer's output. +- [x] **Request registration versus cancellation, `CancellableRequests.scala`.** `withRequest` (53-60) registers the request before calling `check()`. `cancel()` (139-164) sets the volatile `cancelled` flag before taking a snapshot of `requests` (a `ConcurrentHashMap` key set). So every request is either rejected before it runs or included in the abort snapshot. +- [x] **`cancel()` and `close()`.** + - `cancel()` runs once under the Scope lock and never blocks. It calls `shutdownNow`, takes the snapshot, and starts a daemon abort thread that uses the library classloader. + - `close()` (166-205) runs under a separate `cleanupLock`, so the shared monitor never waits on a drain. The abort wait and the pool drain share one 5-second budget. + - The first failure is kept and later ones are added as suppressed. The interrupt flag is cleared during the drain and restored afterwards. +- [x] **`submit`, `await` and `map`.** + - `submit` (102-117) rechecks after a `RejectedExecutionException`. Because `cancelled` is set before `shutdownNow`, a rejection caused by shutdown always reports the recorded cancellation. + - `await` (119-137) calls `checkTask()` before both the expired-timeout path and the `ExecutionException` path. An already-buffered result is still returned when the wait is zero. + - `map` (208-253) handles an already-completed task context because the listener runs first and the `closed` check stops a leaked watch. Nested `guarded` calls drain only once. +- [x] **Kill reason.** `checkTask()` (90-100) throws `TaskKilledException(getKillReason.getOrElse(fallback))` only when `task.isInterrupted()`. `TaskContextUtils` (8-10) wraps Spark's `private[spark] getKillReason()` following the existing `org.apache.spark.injections` pattern. The thread-interrupt path is unchanged. +- [x] **Contracts in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala`.** + - Lines 84-90: `sendCancellableRequestsWithContext` delegates to `CancellableRequests.map`. + - Lines 113-128: `retryPause` calls `check()` on every pause. + - Lines 320-336 and 449-454: retries check before and after each execute/buffer step, and an IOException goes to `keepTrying`. + - Lines 614-636: `boundedWithFabricAuth` wraps all attempts in one `withRequest`, so the same request is never registered twice. It calls `check()` again before returning null after retries are exhausted. +- [x] **Timeout contract in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala`.** In lines 107-108 and 136-146, `concurrentWait` yields either a finite duration or `Duration.Inf`. NaN becomes 0, +∞ becomes a finite `Long.MaxValue` nanoseconds, and −∞ is rejected by `Duration.fromNanos`. So `await` never receives `MinusInf` or `Undefined`. +- [x] **Other contracts.** `CognitiveServicesBaseNoHandler.getInternalTransformer` in `cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala` routes through the same transformer path. I also checked `RequestWithContext` in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/Clients.scala` and the `client`/`response` helpers in `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala`. +- [x] **Test soundness.** I traced all 11 test definitions (13 cases, counting the two `Seq(false, true)` parameterizations) against: + - JDK 8/17 `ThreadPoolExecutor.addWorker`, which builds the worker thread before taking `mainLock`. This makes the `InheritableThreadLocal.childValue` hook valid: `shutdownNow` runs, then the rejection, then `TaskKilledException`. + - `FutureTask` and `ScheduledFutureTask` semantics. + - The Spark 3.5 `TaskContextImpl` constructor, which takes 11 arguments plus a default `resources`. + - `markInterrupted`, `markTaskCompleted` and `setTaskContext`, which are `private[spark]`. That is why the test helper lives in the `org.apache.spark` test package. + + Every test releases its latches and semaphores in `finally` and joins its threads with bounded waits. +- [x] **Security and compatibility.** + - The abort-failure log (line 151) records only the exception class's simple name, never URIs, headers or credentials. + - The slice adds no RDD usage and changes no existing public signature. + - The new surfaces are the `private[http]` `CancellableRequests`, a public `TaskContextUtils` injection consistent with existing injections, and test-only support. +- [x] **Concerns checked and not reported, because they are not concrete defects:** + - The first test's monitor-classloader check only proves the fix when that test is the first in the JVM to create the shared monitor. Production sets the monitor's classloader explicitly either way. + - Inheritable thread-locals and the classloader pinned by the long-lived monitor; the PR explicitly disclaims clearing inheritable thread-locals. + - The abort/`setCancellable` window inside Apache HttpClient, which the library re-checks itself before sending. + - Non-finite timeouts, which production callers cannot pass. +- [ ] **Tests not executed.** This was a read-only review, so I ran no build or test commands. I took the PR's validation evidence into account but did not reproduce it. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md new file mode 100644 index 00000000000..3c7eb3e5e81 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md @@ -0,0 +1,47 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 170 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:52:46Z" +outcome: "clean" +reviewer_model: "gpt-6.1-sol" +input_sha256: "d8ac79f62cc70d298c32acdb0b722045517f92677b7db1d9448a089100affa73" +--- + +--- +producer: review-code +schema_version: 1 +tier: Max +contract: direct +execution_backend: copilot-cli +attempt: 170 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T12:48:31.499Z" +outcome: clean +reviewer_model: gpt-6.1-sol +input_sha256: f7ffba2e9552f1f03381a6ea1e46c94b7e0b41454e3d7f9306e627d969e9dd56 +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: gpt-6.1-sol +- **Artifact**: reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-gpt-6.1-sol-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Reviewed slice 1/24 completely through bounded views of the assigned frozen diff: lines 1–150, 151–300, 301–450, 451–600, 601–750, 751–900 and 901–1001. All four files are additions; complete absent-before/added-after content was covered. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` (1–254): submission and timeout cancellation races, request registration, worker shutdown, bounded cleanup, exception suppression, interruption preservation and infrastructure classloaders. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala` (1–10) and `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala` (1–31), including kill-reason access and synthetic task lifecycle behavior. +- [x] Reviewed `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` (1–682), covering all 13 regression cases and their synchronization, assertions and cleanup. +- [x] Traced necessary contracts through `HTTPClients.scala`, `HTTPTransformer.scala`, `VerifyBoundedRequests.scala`, `AIPromptCancellationSuite.scala` and the cancellation documentation in `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [x] Read `AGENTS.md`, the required branch/review skills and master branch reference, plus version declarations in `build.sbt` and `environment.yml`. +- [ ] Tests were not executed under the read-only contract. Supplied validation results are caller-provided evidence, not independently reproduced. This clean verdict applies only to slice 1. \ No newline at end of file diff --git a/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-grok-4.7-slice1.md b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-grok-4.7-slice1.md new file mode 100644 index 00000000000..2ce6c118c27 --- /dev/null +++ b/reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-grok-4.7-slice1.md @@ -0,0 +1,30 @@ +--- +producer: "review-code" +schema_version: 1 +tier: "Max" +contract: "direct" +execution_backend: "copilot-cli" +attempt: 170 +stage_attempt: 1 +round: 1 +timestamp: "2026-11-06T13:04:19Z" +outcome: "clean" +reviewer_model: "grok-4.7" +input_sha256: "1c7ca362de6497649386da96dcb915ca33df60e4855bff2825fb893f5a56c326" +--- + +## Review Summary +- **Round**: 1 +- **Theme**: Broad sweep +- **Mode**: resolved-plan +- **Model**: grok-4.7 +- **Artifact**: reviews/pr-2811/task-unknown-attempt-170-stage-attempt-1-review-1-grok-4.7-slice1.md +- **Issues Found**: 0 +- **Verdict**: CLEAN + +## Evidence Checklist +- [x] Assigned frozen diff read in full (lines 1–1001) in bounded chunks; before/after compared for all four slice files, all new. +- [x] Reviewed `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala` for kill-reason preservation, submission/timeout rechecks, abort/drain, classloader isolation, and cooperative cancellation limits. +- [x] Reviewed `core/src/main/scala/org/apache/spark/injections/TaskContextUtils.scala`, `core/src/test/scala/org/apache/spark/TaskContextTestSupport.scala`, and `core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala` against Spark 3.5.0 / Scala 2.12.17 (`build.sbt`) and the cancellation contract in `docs/Explore Algorithms/OpenAI/Unified AI prompting.md`. +- [x] Checked cross-file use in `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPClients.scala` (`withRequest`, `sendCancellableRequestsWithContext`) and `core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala` (task-local iterator consumption). No concrete correctness, security, or spec defect found in this slice. +- [ ] Tests were not executed; this read-only pass did not treat the supplied validation narrative as proof. \ No newline at end of file diff --git a/tools/ci/tests/test_aiprompt_notebook.py b/tools/ci/tests/test_aiprompt_notebook.py new file mode 100644 index 00000000000..654a28f034e --- /dev/null +++ b/tools/ci/tests/test_aiprompt_notebook.py @@ -0,0 +1,151 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +import ast +import json +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, call + +import pytest + + +NOTEBOOK = ( + Path(__file__).resolve().parents[3] + / "docs" + / "Explore Algorithms" + / "OpenAI" + / "AIPrompt.ipynb" +) + + +@pytest.mark.parametrize( + ("diagnostic", "cleanup_failure"), + [ + (None, None), + ("incomplete", None), + ("missing-answer", None), + (None, "request"), + ("incomplete", "request"), + (None, "confirmation"), + ], +) +def test_stored_response_cleanup_uses_collected_ids_before_validation( + monkeypatch, capsys, diagnostic, cleanup_failure +): + notebook = json.loads(NOTEBOOK.read_text(encoding="utf-8")) + cells = [ + "".join(cell["source"]) + for cell in notebook["cells"] + if cell["cell_type"] == "code" + ] + setup = ast.parse(next(cell for cell in cells if "def run_example(" in cell)) + helpers = [ + node + for node in setup.body + if isinstance(node, ast.FunctionDef) + and node.name in ("run_example", "validate_example") + ] + storage = next(cell for cell in cells if "created_response_ids = []" in cell) + + def stage(): + result = Mock() + for name in ( + "setPromptTemplate", + "setStore", + "setResponseIdCol", + "setPreviousResponseId", + "setPreviousResponseIdCol", + "copy", + ): + getattr(result, name).return_value = result + result.getResponseIdCol.return_value = "response_id" + return result + + stored = stage() + continued = stage() + rows = [ + { + "response_id": identifier, + "answer": "ALPHA", + "error": None, + "raw_response": json.dumps({"id": identifier}), + } + for identifier in ("created-first", "created-second") + ] + if diagnostic == "incomplete": + rows[0]["error"] = {"code": "incomplete"} + elif diagnostic == "missing-answer": + rows[0]["answer"] = None + stored.transform.return_value.collect.return_value = rows + continued.transform.return_value.collect.return_value = [ + {"answer": "ALPHA", "error": None, "raw_response": '{"output": "ALPHA"}'} + ] + response = Mock(status_code=204) + delete = Mock(return_value=response) + + class DeleteError(Exception): + pass + + if cleanup_failure == "request": + delete.side_effect = [DeleteError("Synthetic deletion failure"), response] + elif cleanup_failure == "confirmation": + unconfirmed = Mock(status_code=200) + unconfirmed.json.return_value = {"deleted": False} + delete.side_effect = [unconfirmed, response] + monkeypatch.setitem( + sys.modules, + "requests", + SimpleNamespace(delete=delete, RequestException=DeleteError), + ) + countries = Mock() + namespace = { + "json": json, + "completed_examples": [], + "prompt_for": Mock(side_effect=[stored, continued]), + "countries": countries, + "F": SimpleNamespace(lit=lambda value: value), + "openai_endpoint": "https://example.invalid", + "openai_token": "synthetic-token", + } + exec( + compile(ast.Module(body=helpers, type_ignores=[]), str(NOTEBOOK), "exec"), + namespace, + ) + + if diagnostic: + with pytest.raises(AssertionError, match="responses/store-and-response-id"): + exec(compile(storage, str(NOTEBOOK), "exec"), namespace) + assert namespace["completed_examples"] == [] + continued.transform.assert_not_called() + if cleanup_failure: + assert "Could not delete" in capsys.readouterr().err + else: + if cleanup_failure: + with pytest.raises(RuntimeError, match="Could not delete") as error: + exec(compile(storage, str(NOTEBOOK), "exec"), namespace) + assert isinstance(error.value.__cause__, (DeleteError, ValueError)) + else: + exec(compile(storage, str(NOTEBOOK), "exec"), namespace) + assert namespace["completed_examples"] == [ + "responses/store-and-response-id", + "responses/scalar-previous-id", + "responses/column-previous-id", + ] + stored.transform.assert_called_once_with(countries) + stored.transform.return_value.collect.assert_called_once_with() + assert namespace["created_response_ids"] == ["created-first", "created-second"] + expected_deleted = ( + ["created-second"] if cleanup_failure else ["created-first", "created-second"] + ) + assert namespace["deleted_response_ids"] == expected_deleted + assert delete.call_args_list == [ + call( + f"https://example.invalid/openai/v1/responses/{identifier}", + headers={"Authorization": "Bearer synthetic-token"}, + timeout=60, + ) + for identifier in ("created-first", "created-second") + ] + assert response.raise_for_status.call_count == (1 if cleanup_failure else 2) diff --git a/website/sidebars.js b/website/sidebars.js index 8ac0de367f7..a6cb4f51bee 100644 --- a/website/sidebars.js +++ b/website/sidebars.js @@ -50,6 +50,8 @@ module.exports = { items: [ "Explore Algorithms/OpenAI/Langchain", "Explore Algorithms/OpenAI/OpenAI", + "Explore Algorithms/OpenAI/Unified AI prompting", + "Explore Algorithms/OpenAI/AIPrompt", "Explore Algorithms/OpenAI/Quickstart - OpenAI Embedding", "Explore Algorithms/OpenAI/Quickstart - OpenAI Embedding and GPU based KNN", "Explore Algorithms/OpenAI/Quickstart - Custom Embeddings and Approximate KNN on GPU",