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/AIDurableInput.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIDurableInput.scala new file mode 100644 index 00000000000..39855315046 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIDurableInput.scala @@ -0,0 +1,194 @@ +// 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 org.apache.spark.sql.{DataFrame, Row, functions => F} +import org.apache.spark.sql.types.{StringType, StructType} +import spray.json._ + +private[ai] object AIDurableInput { + private val ExternalBindings = Set( + "previous_response_id", "conversation", "conversation_id", "file_id", "file_ids", + "tools", "tool_choice", "tool_resources", "functions", "function_call", "retrieve", "run_id", + "data_sources", "web_search_options", "prompt", "tool_call_id") + private val OutputSchemas = Map( + "chat_completions" -> Set(List("response_format", "json_schema", "schema")), + "responses" -> Set(List("text", "format", "schema"))) + + private def externalBindings(value: JsValue, path: List[String] = Nil, + outputSchemas: Set[List[String]] = Set.empty): Unit = value match { + case JsObject(fields) if !outputSchemas(path) => + val bindings = fields.iterator.collect { + case (name, entry) if ExternalBindings(name) && entry != JsNull => name + }.toSeq.sorted + require(bindings.isEmpty, + "Durable requests cannot depend on provider state, retrieval, tools, or external resources; " + + s"restricted fields: ${bindings.mkString(", ")}") + content(value) + fields.foreach { case (name, entry) => externalBindings(entry, path :+ name, outputSchemas) } + case JsArray(items) => items.foreach(externalBindings(_)) + case _ => + } + + private def body(value: JsValue, api: String): Unit = + externalBindings(JsObject(AIJson.obj(value).fields - "metadata"), + outputSchemas = OutputSchemas.getOrElse(api, Set.empty)) + + private def responseFormatBody(value: Any, api: String): JsObject = { + val format = AIRequestAdapter.responseFormat(value, api) + if (api == "responses") JsObject("text" -> JsObject("format" -> format)) + else JsObject("response_format" -> format) + } + + private def questions(api: String, value: JsValue): Unit = (api, value) match { + case ("systemone", JsObject(named)) => + named.values.foreach { + // Question names and criteria labels are user-defined text, not provider binding fields. + case JsObject(fields) => + externalBindings(JsObject(fields - "criteria")) + fields.get("criteria").foreach { + case JsObject(criteria) => criteria.values.foreach(externalBindings(_)) + case other => externalBindings(other) + } + case other => externalBindings(other) + } + case _ => externalBindings(value) + } + + def validate(stage: AIPrompt, schema: StructType): Unit = { + val api = stage.getApiType + Seq(stage.subscriptionKey, stage.AADToken, stage.CustomAuthHeader, + stage.customHeaders, stage.telemHeaders).foreach { p => + require(!stage.get(p).orElse(stage.getDefault(p)).exists(_.isRight), + s"Durable mode rejects ${p.name}Col; rebind scalar credentials on resume") + } + require(!stage.isDefined(stage.previousResponseId), + "Durable mode does not support previousResponseId or previousResponseIdCol") + stage.get(stage.extraBody).orElse(stage.getDefault(stage.extraBody)).foreach { + case Left(value) => body(AIJson.fromAny(value), api) + case _ => + } + stage.get(stage.responseFormat).orElse(stage.getDefault(stage.responseFormat)).foreach { + case Left(value) => body(responseFormatBody(value, api), api) + case _ => + } + stage.get(stage.questions).orElse(stage.getDefault(stage.questions)).foreach { + case Left(value) => questions(api, AIJson.fromAny(value)) + case _ => + } + stage.getColumnTypes.foreach { case (name, _) => AIColumnPlan.resolve(schema, name) } + } + + private def inline(value: JsValue, literalContent: Boolean = false): Unit = value match { + case JsString(url) if !literalContent && url.trim.nonEmpty => + // Request validation handles malformed inline media as row errors; only external dependencies abort here. + require(OpenAIAttachmentUtils.isDataUrl(url.trim), + "Durable media must be inline base64 data URLs, not remote content or paths") + case JsArray(items) => items.foreach(inline(_)) + case JsObject(fields) => fields.get("url").foreach(inline(_)) + case _ => + } + + private def content(part: JsValue): Unit = part match { + case JsObject(fields) => + Seq("image_url", "file_data", "file_url").foreach { name => + // A string file_data is literal content, not a location to resolve. + fields.get(name).foreach(inline(_, literalContent = name == "file_data")) + } + case _ => + } + + private def messages(value: JsValue): Unit = { + externalBindings(value) + value match { + case JsArray(items) => items.foreach { + case JsObject(fields) => + require(!fields.get("id").exists(_ != JsNull) && !fields.get("type").exists { + case JsString("item_reference" | "function_call" | "function_call_output") => true + case _ => false + }, "Durable messages cannot reference provider items or tool results") + fields.get("content").foreach { + case JsArray(parts) => parts.foreach(content) + case _ => + } + case _ => + } + case _ => + } + } + + private def validateDeferredJson(kind: String, text: String): Unit = { + val parsed = try Some(AIJson.parse(text)) catch { + case _: IllegalArgumentException => None + } + // Inference parses these cells again and commits malformed values as terminal row errors. + parsed.foreach { value => + if (kind.endsWith("_questions")) questions(kind.stripSuffix("_questions"), value) + else if (kind == "systemone") externalBindings(value) + else messages(value) + } + } + + private def validateBodyCell(kind: String, text: String, api: String): Unit = { + // Malformed bodies and formats become terminal row errors; resource checks must not be deferred. + val parsed = try Some(if (kind == "body") AIJson.objectString(text) + else responseFormatBody(text, kind.stripSuffix("_format"))) catch { + case _: IllegalArgumentException => None + } + parsed.foreach(body(_, api)) + } + + private def validateCell(kind: String, text: String, api: String): Unit = kind match { + case "prompt" => () // Selecting the rendered value forces evaluation before any inference requests. + case "inline" => + val normalized = text.trim + if (normalized.nonEmpty) inline(JsString(normalized)) + case "body" | "chat_completions_format" | "responses_format" => validateBodyCell(kind, text, api) + // Inference encodes message Rows directly, so parser failures cannot be deferred to row errors. + case "messages" => messages(AIJson.parse(text)) + case "systemone" | "decisions" | "systemone_questions" | "decisions_questions" => + validateDeferredJson(kind, text) + case _ => throw new IllegalArgumentException("Unsupported durable input kind") + } + + private def validateCells(kinds: Seq[String], row: Row, api: String): Unit = { + kinds.zipWithIndex.foreach { case (kind, index) => + if (!row.isNullAt(index)) validateCell(kind, row.getString(index), api) + } + } + + def validateRows(stage: AIPrompt, snapshot: DataFrame, preparedInput: Option[String] = None): Unit = { + val api = stage.getApiType + val kinds = Seq.newBuilder[String] + val columns = Seq.newBuilder[org.apache.spark.sql.Column] + def add(name: String, kind: String, json: Boolean = false): Unit = { + val resolved = AIColumnPlan.resolve(snapshot.schema, name) + kinds += kind + val column = AIColumnPlan.column(resolved) + columns += (if (json) F.to_json(column) else column) + } + stage.getColumnTypes.foreach { case (name, kind) => if (kind == "path") add(name, "inline") } + if (stage.isSet(stage.messagesCol)) add(stage.getMessagesCol, "messages", json = true) + stage.get(stage.extraBody).orElse(stage.getDefault(stage.extraBody)).flatMap(_.right.toOption) + .foreach(add(_, "body")) + stage.get(stage.responseFormat).orElse(stage.getDefault(stage.responseFormat)).flatMap(_.right.toOption) + .foreach { name => + val dataType = snapshot.schema(AIColumnPlan.resolve(snapshot.schema, name)).dataType + add(name, api + "_format", json = dataType != StringType) + } + stage.get(stage.questions).orElse(stage.getDefault(stage.questions)).flatMap(_.right.toOption) + .foreach(add(_, api + "_questions")) + if (AIRequestAdapter.NativeApis(api) && stage.getInputMode == "json") { + preparedInput.orElse(stage.get(stage.stateCol)).orElse(stage.get(stage.inputCol)) + .foreach(add(_, api)) + } else preparedInput.foreach(add(_, "prompt")) + val selectedKinds = kinds.result() + if (selectedKinds.nonEmpty) { + snapshot.select(columns.result(): _*).foreachPartition { rows: Iterator[Row] => + rows.foreach(validateCells(selectedKinds, _, api)) + } + } + } +} 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..06096e3c0a2 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPrompt.scala @@ -0,0 +1,658 @@ +// 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.{AesEncrypt, Attribute, EulerNumber, GenericRowWithSchema, + LeafExpression, Literal, NamedLambdaVariable, Pi, SubqueryExpression, UserDefinedExpression} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Generate, Project, Window} +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.Collections +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[ai] def hasConfiguredAuthentication(prompt: AIPrompt): Boolean = { + def scalar[T](param: ServiceParam[T]): Option[T] = + prompt.get(param).orElse(prompt.getDefault(param)).flatMap(_.left.toOption) + val resolved = ServiceAuthHeaders.resolve(scalar(prompt.subscriptionKey), prompt.subscriptionKeyHeaderName, + prompt.aadHeaderName, scalar(prompt.AADToken), scalar(prompt.CustomAuthHeader), + scalar(prompt.customHeaders), None, None, None) + resolved.headers.keys.exists(name => prompt.authenticationHeaderNames(name.toLowerCase(java.util.Locale.ROOT))) || + (prompt.isDefined(prompt.url) && prompt.usingImplicitFabricEndpoint(prompt.getUrl.stripSuffix("/") + "/")) + } + + private val FirstRetryMs = 100 + private val SecondRetryMs = 500 + private val ThirdRetryMs = 1000 + private val RetryIntervals = Array(FirstRetryMs, SecondRetryMs, ThirdRetryMs) + private[ai] val DefaultHttpPolicy = "synapseml-bounded-http-v1:" + RetryIntervals.mkString(",") + 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") + } + } + + def loadOutput(spark: SparkSession, outputPath: String): DataFrame = + AIPromptOutput.loadOutput(spark, outputPath) + + def getOutputStatus(spark: SparkSession, outputPath: String): AIPromptOutputStatus = + AIPromptOutput.getOutputStatus(spark, outputPath) + + def getOutputStatusJson(spark: SparkSession, outputPath: String): String = + getOutputStatus(spark, outputPath).toJson + + def releaseOutputOwner(spark: SparkSession, outputPath: String, + expectedOwnerToken: String, confirmStopped: Boolean): Unit = + AIPromptOutput.releaseOutputOwner(spark, outputPath, expectedOwnerToken, confirmStopped) +} + +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. Uncommitted requests can repeat even with outputPath configured. */ +class AIPrompt(override val uid: String) extends OpenAIPrompt(uid) + with AIPromptParams with HasAIPromptRecovery 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" + + private[ai] lazy val authenticationHeaderNames: Set[String] = + Set(subscriptionKeyHeaderName, aadHeaderName).map(_.toLowerCase(java.util.Locale.ROOT)) + + 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() + val result = stage.inferenceSchema(schema) + if (stage.isDefined(outputPath)) { + AIPromptOutput.validateSettings(stage) + AIPromptOutput.storageSchema(schema) + stage.validateDurableInput(schema, SparkSession.active) + AIPromptOutput.storageSchema(result) + } else result + } + + override def transform(dataset: Dataset[_]): DataFrame = AIPrompt.withSession(dataset.sparkSession) { + logTransform({ + val stage = effectiveCopy() + stage.inferenceSchema(dataset.schema) + if (stage.isDefined(outputPath)) AIPromptOutput.transformEffective(stage, dataset.toDF()) + else { + 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, hiddenColumns: Seq[String], + interpolateFilenames: Boolean = true): DataFrame = { + if (!isSet(promptTemplate)) input + else { + val filenames = if (!interpolateFilenames) Map.empty[String, Column] else { + 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)) + val publicInput = if (hiddenColumns.isEmpty) input else input.drop(hiddenColumns: _*) + // Bind against public columns without allowing the template to filter or expand rows. + val scoped = AIPrompt.bindTemplate(publicInput, 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 = { + infer(dataset, Seq.empty, http) + } + + private[ai] def infer(dataset: Dataset[_], hiddenColumns: Seq[String], 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.drop(hiddenColumns: _*).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) + if (hiddenColumns.nonEmpty) validateDurableTemplate(original, promptName, hiddenColumns) + val prepared = preparePrompt(original, promptName, hiddenColumns) + // 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) + } + } + + private[ai] def validateDurableInput(schema: StructType, session: => SparkSession): Unit = { + AIDurableInput.validate(this, schema) + if (isSet(promptTemplate)) { + val input = session.createDataFrame(Collections.emptyList[Row](), schema) + val outputPlan = plan(input.schema, parsed = false) + validateDurableTemplate(input, outputPlan.scratch(s"${uid}_prompt"), Seq.empty) + } + } + + private def validateDurableTemplate(input: DataFrame, promptName: String, + hiddenColumns: Seq[String]): Unit = if (isSet(promptTemplate)) { + // Validate user expressions before inserting SynapseML's stable filename UDF. + val prepared = preparePrompt(input, promptName, hiddenColumns, interpolateFilenames = false) + val analyzed = AIPrompt.withSafeTemplate(prepared.queryExecution.analyzed) + require(!analyzed.exists { + case _: Window | _: Aggregate | _: Generate => true + case _ => false + }, "Durable prompt templates require row-local expressions; materialize windows, aggregates and generators " + + "into input columns") + // Input expressions are snapshotted once; inspect only the template projection for unstable functions. + val expressions = analyzed.expressions + require(!expressions.exists(SubqueryExpression.hasSubquery), + "Durable prompt templates cannot contain SQL subqueries; materialize external values into input columns") + val unstable = expressions.exists(_.exists { + // Spark's deterministic flag does not account for AES's generated initialization vector. + case _: UserDefinedExpression | _: AesEncrypt => true + case expression if !expression.deterministic => true + case _: Attribute | _: Literal | _: NamedLambdaVariable | _: Pi | _: EulerNumber => false + case _: LeafExpression => true + case _ => false + }) + require(!unstable, "Durable prompt templates require stable built-in expressions over input columns; " + + "materialize UDFs, encryption and runtime-dependent values into input columns") + } + + private[ai] def validateDurableRows(snapshot: DataFrame): Unit = { + if (isSet(promptTemplate)) { + val outputPlan = plan(snapshot.schema, parsed = false) + val promptName = outputPlan.scratch(s"${uid}_prompt") + validateDurableTemplate(snapshot, promptName, Seq.empty) + val prepared = preparePrompt(snapshot, promptName, Seq.empty) + AIDurableInput.validateRows(this, prepared, Some(promptName)) + } else AIDurableInput.validateRows(this, snapshot) + } + + 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/AIPromptDelta.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDelta.scala new file mode 100644 index 00000000000..b2b548d2df8 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDelta.scala @@ -0,0 +1,243 @@ +// 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.hadoop.fs.Path +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ + +import java.nio.charset.StandardCharsets +import java.util.{Collections, UUID} +import scala.util.control.NonFatal + +private[ai] case class AIPromptManifest(runId: String, + libraryVersion: String, + sparkVersion: String, + fingerprint: String, + inputSchema: String, + outputSchema: String, + bindings: String, + keyColumn: String, + state: String, + inputVersion: Option[Long], + inputCount: Option[Long], + owner: Option[String]) { + def inputType: StructType = AIPromptStorageCodec.readSchema(inputSchema) + def outputType: StructType = AIPromptStorageCodec.readSchema(outputSchema) + def keyed(schema: StructType): StructType = + StructType(StructField(keyColumn, LongType) +: schema.fields.toSeq) +} + +/** SQL transactions arbitrate control transitions; releasing an owner requires stopped prior execution. */ +private[ai] object AIPromptDelta { + val FormatVersion: Int = 2 + private val ByteMask = 0xff + private val IncompleteInitialization = + "AIPrompt initialization is incomplete and may still be running in the creating application. " + + "Status, output reads and resume are unavailable. Wait for initialization to finish, or confirm the " + + "creating application has stopped before abandoning this path. Only after stopped execution is confirmed, " + + "correct any input errors and use a new outputPath or clean up this run." + val ControlSchema: StructType = StructType(Seq( + StructField("formatVersion", IntegerType), StructField("runId", StringType), + StructField("libraryVersion", StringType), StructField("sparkVersion", StringType), + StructField("fingerprint", StringType), StructField("inputSchema", StringType), + StructField("outputSchema", StringType), StructField("bindings", StringType), + StructField("keyColumn", StringType), StructField("state", StringType), + StructField("inputVersion", LongType), StructField("inputCount", LongType), + StructField("owner", StringType))) + + def validatePath(path: String): Unit = { + require(Option(path).exists(_.trim.nonEmpty), "AIPrompt outputPath must not be empty.") + require(!path.contains("${"), "AIPrompt outputPath must not contain SQL variable substitutions (${...}).") + } + + def qualify(spark: SparkSession, path: String): String = { + validatePath(path) + val qualified = org.apache.spark.ml.Serializer.makeQualifiedPath(spark, path).toString + validatePath(qualified) + qualified + } + + def requireRuntime(spark: SparkSession, path: String): Unit = { + try { + Class.forName("org.apache.spark.sql.delta.sources.DeltaDataSource", false, + Thread.currentThread().getContextClassLoader) + } catch { + case NonFatal(e) => throw new IllegalStateException( + "AIPrompt durable output requires a compatible, caller-configured Delta runtime.", e) + } + require(spark.conf.getOption("spark.sql.extensions").exists(_.contains("DeltaSparkSessionExtension")) && + spark.conf.getOption("spark.sql.catalog.spark_catalog").exists(_.contains("DeltaCatalog")), + "Configure the Delta Spark extension and catalog before creating the Spark session.") + val uri = new Path(path).toUri + val scheme = uri.getScheme + require(Set("file", "hdfs", "abfs", "abfss").contains(scheme), + "AIPrompt requires HDFS or ABFS storage, or same-host file: storage in Spark local mode. " + + "Other filesystem/LogStore configurations are not qualified.") + require(scheme != "file" || (spark.sparkContext.isLocal && Option(uri.getAuthority).forall(_.isEmpty)), + "AIPrompt file: durable storage requires Spark local mode and a same-host local directory. " + + "File URI authorities are not supported; use HDFS or ABFS storage on a non-local cluster.") + } + + def child(path: String, name: String): String = new Path(new Path(path), name).toString + + def table(path: String, name: String): String = "delta.`" + child(path, name).replace("`", "``") + "`" + + def hex(bytes: Array[Byte]): String = bytes.map(b => f"${b & ByteMask}%02x").mkString + + def literal(value: String): String = s"CAST(unhex('${hex(value.getBytes(StandardCharsets.UTF_8))}') AS STRING)" + + def column(name: String): org.apache.spark.sql.Column = AIColumnPlan.column(name) + + def version(spark: SparkSession, path: String, name: String): Long = + spark.sql(s"DESCRIBE HISTORY ${table(path, name)} LIMIT 1").select("version").head().getLong(0) + + def read(spark: SparkSession, path: String, name: String, version: Long): DataFrame = + spark.read.format("delta").option("versionAsOf", version).load(child(path, name)) + + def latest(spark: SparkSession, path: String, name: String): (Long, DataFrame) = { + val pinned = version(spark, path, name) + pinned -> read(spark, path, name, pinned) + } + + def create(spark: SparkSession, path: String, manifest: AIPromptManifest, + beforeCreate: () => Unit): Unit = { + val root = new Path(path) + val fs = root.getFileSystem(spark.sessionState.newHadoopConf()) + def uninitialized: Boolean = !fs.exists(root) || (fs.getFileStatus(root).isDirectory && + fs.listStatus(root).forall(entry => root.toUri.getScheme == "file" && entry.isFile && + entry.getPath.getName == AIPromptLocalTransaction.LockFile)) + if (!uninitialized) { + requireReady(control(spark, path)) + throw new IllegalArgumentException("AIPrompt transform requires a new or empty run directory; " + + "use loadOutput to read committed outcomes or resume to continue an existing run.") + } + val row = Row(FormatVersion, manifest.runId, manifest.libraryVersion, manifest.sparkVersion, + manifest.fingerprint, manifest.inputSchema, + manifest.outputSchema, manifest.bindings, manifest.keyColumn, manifest.state, + manifest.inputVersion.map(Long.box).orNull, manifest.inputCount.map(Long.box).orNull, + manifest.owner.orNull) + beforeCreate() + AIPromptLocalTransaction(path) { + if (root.toUri.getScheme == "file") { + require(uninitialized, + "AIPrompt initial ownership conflict: another creator has initialized this run.") + } + // The Delta create arbitrates remote writers; local rename-based LogStores also need the kernel lock. + spark.createDataFrame(Collections.singletonList(row), ControlSchema) + .write.format("delta").mode("errorifexists").save(child(path, "_control")) + val committed = control(spark, path) + require(committed.runId == manifest.runId && committed.owner == manifest.owner, + "AIPrompt initial ownership conflict.") + } + } + + def control(spark: SparkSession, path: String): AIPromptManifest = { + val root = new Path(path) + val fs = root.getFileSystem(spark.sessionState.newHadoopConf()) + require(fs.exists(root), "No AIPrompt run exists at outputPath. Check the path supplied to this operation.") + require(fs.isDirectory(root) && fs.isDirectory(new Path(root, "_control")), + "outputPath is not an AIPrompt run directory. Check the path; do not remove unrelated data. " + + "For a new run, choose a new or empty outputPath.") + val rows = try { + latest(spark, path, "_control")._2.select(ControlSchema.fieldNames.map(column): _*).take(2) + } catch { + case error: AnalysisException if error.getErrorClass == "DELTA_TABLE_ONLY_OPERATION" => + throw new IllegalArgumentException(IncompleteInitialization, error) + } + require(rows.length == 1, "Corrupt AIPrompt control table: expected exactly one row.") + val row = rows.head + require(row.getAs[Int]("formatVersion") == FormatVersion, "Unsupported AIPrompt output format version.") + require(Seq("runId", "libraryVersion", "sparkVersion", "fingerprint", "bindings", "keyColumn", + "inputSchema", "outputSchema").forall { name => + Option(row.getAs[String](name)).exists(_.trim.nonEmpty) + }, "Corrupt AIPrompt control table: required identity fields must not be null or empty.") + val manifest = AIPromptManifest(row.getAs[String]("runId"), + row.getAs[String]("libraryVersion"), row.getAs[String]("sparkVersion"), row.getAs[String]("fingerprint"), + row.getAs[String]("inputSchema"), row.getAs[String]("outputSchema"), row.getAs[String]("bindings"), + row.getAs[String]("keyColumn"), row.getAs[String]("state"), + Option(row.getAs[java.lang.Long]("inputVersion")).map(_.longValue), + Option(row.getAs[java.lang.Long]("inputCount")).map(_.longValue), + Option(row.getAs[String]("owner"))) + require(manifest.owner.forall(_.trim.nonEmpty), + "Corrupt AIPrompt control table: owner must be null or non-blank.") + require(Set("initializing", "running", "complete").contains(manifest.state), + "Corrupt AIPrompt run state.") + manifest.inputType + manifest.outputType + manifest + } + + def requireReady(manifest: AIPromptManifest): Unit = + require(manifest.state != "initializing" && manifest.inputVersion.isDefined && manifest.inputCount.isDefined, + IncompleteInitialization) + + def requireUnowned(manifest: AIPromptManifest): Unit = + require(manifest.owner.isEmpty, + "AIPrompt run already has an owner. Stopped-job administrative release is required.") + + private def update(spark: SparkSession, path: String, assignments: String, condition: String): Unit = { + AIPromptLocalTransaction(path) { + val changed = spark.sql(s"UPDATE ${table(path, "_control")} SET $assignments WHERE $condition").collect() + require(changed.length == 1 && changed.head.getLong(0) == 1L, + "AIPrompt ownership conflict: the expected owner no longer owns this run.") + } + } + + private def owned(manifest: AIPromptManifest, token: String): String = + s"runId = ${literal(manifest.runId)} AND owner = ${literal(token)}" + + def assertOwner(spark: SparkSession, path: String, manifest: AIPromptManifest, token: String): Unit = { + val current = control(spark, path) + require(current.runId == manifest.runId && current.owner.contains(token), + "AIPrompt ownership conflict: inference and writes are stopped.") + } + + def ready(spark: SparkSession, path: String, manifest: AIPromptManifest, token: String, + inputVersion: Long, count: Long): AIPromptManifest = { + update(spark, path, s"state = 'running', inputVersion = $inputVersion, inputCount = $count", + owned(manifest, token) + " AND state = 'initializing'") + control(spark, path) + } + + def acquire(spark: SparkSession, path: String, manifest: AIPromptManifest): String = { + requireReady(manifest) + require(manifest.state != "complete", "A completed AIPrompt run does not require ownership.") + requireUnowned(manifest) + val token = UUID.randomUUID().toString + update(spark, path, s"owner = ${literal(token)}, state = 'running'", + s"runId = ${literal(manifest.runId)} AND owner IS NULL AND state = ${literal(manifest.state)}") + assertOwner(spark, path, manifest, token) + token + } + + def complete(spark: SparkSession, path: String, manifest: AIPromptManifest, token: String): Unit = + update(spark, path, "state = 'complete', owner = NULL", owned(manifest, token)) + + def release(spark: SparkSession, path: String, manifest: AIPromptManifest, token: String): Unit = + update(spark, path, "owner = NULL", owned(manifest, token)) + + def validateSchema(frame: DataFrame, expected: StructType): Unit = + require(AIPromptStorageCodec.nullable(frame.schema) == AIPromptStorageCodec.physical(expected), + "Corrupt or incompatible AIPrompt Delta schema; schema evolution is not supported.") + + def input(spark: SparkSession, path: String, manifest: AIPromptManifest): DataFrame = { + requireReady(manifest) + require(version(spark, path, "_input") == manifest.inputVersion.get, + "The immutable AIPrompt input snapshot has been modified. After initialization, " + + "do not OPTIMIZE or VACUUM the private _input table; disable automatic maintenance for it.") + val frame = read(spark, path, "_input", manifest.inputVersion.get) + validateSchema(frame, manifest.keyed(manifest.inputType)) + frame + } + + def keyCount(frame: DataFrame, key: String): Long = { + val counts = frame.agg(count(lit(1)), count(column(key)), countDistinct(column(key))).head() + val total = counts.getLong(0) + require(total == counts.getLong(1) && total == counts.getLong(2), + "Corrupt AIPrompt row identity: keys must be unique and non-null.") + total + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLocalTransaction.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLocalTransaction.scala new file mode 100644 index 00000000000..59f4874ef1c --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLocalTransaction.scala @@ -0,0 +1,74 @@ +// 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.hadoop.fs.Path + +import java.nio.channels.{FileChannel, FileLock} +import java.nio.file.{Files, Paths, StandardOpenOption} +import java.util.concurrent.locks.ReentrantLock +import scala.annotation.tailrec +import scala.collection.mutable + +/** + * Local Delta LogStores serialize rename-based commits within one JVM only. The kernel lock + * serializes this feature's control transactions across local drivers, including table creation. + * It is never an ownership lease: the durable owner survives both lock release and driver death. + */ +private[ai] object AIPromptLocalTransaction { + val LockFile: String = ".aiprompt-control.lock" + private val LockPollMillis = 25L + + private case class LocalLock(lock: ReentrantLock, var users: Int) + private val Locks = mutable.Map.empty[java.nio.file.Path, LocalLock] + + @tailrec + private def acquire(channel: FileChannel): FileLock = { + // Blocking kernel acquisition can also block Thread.interrupt on Windows. + val lock = channel.tryLock() + if (lock != null) lock else { + Thread.sleep(LockPollMillis) + acquire(channel) + } + } + + def apply[T](path: String)(transaction: => T): T = try { + val uri = new Path(path).toUri + if (uri.getScheme == "file") { + val root = Paths.get(uri) + Files.createDirectories(root) + val canonical = root.toRealPath() + val local = Locks.synchronized { + val entry = Locks.getOrElseUpdate(canonical, LocalLock(new ReentrantLock(), 0)) + entry.users += 1 + entry + } + try { + local.lock.lockInterruptibly() + try { + val channel = FileChannel.open(canonical.resolve(LockFile), + StandardOpenOption.CREATE, StandardOpenOption.WRITE) + try { + val lock = acquire(channel) + try { + if (Thread.currentThread().isInterrupted) { + throw new InterruptedException("AIPrompt control transaction interrupted before execution.") + } + transaction + } finally lock.release() + } finally channel.close() + } finally local.lock.unlock() + } finally Locks.synchronized { + local.users -= 1 + if (local.users == 0) Locks.remove(canonical) + } + } else { + transaction + } + } catch { + case error: InterruptedException => + Thread.currentThread().interrupt() + throw error + } +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutput.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutput.scala new file mode 100644 index 00000000000..94999b7ce6e --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutput.scala @@ -0,0 +1,449 @@ +// 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.build.BuildInfo +import com.microsoft.azure.synapse.ml.io.http.ErrorUtils +import com.microsoft.azure.synapse.ml.services.openai.UsageUtils +import org.apache.spark.ml.param.Param +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConfInjections +import org.apache.spark.sql.types.{DataType, StringType, StructType} +import spray.json._ + +import java.nio.charset.StandardCharsets +import java.security.MessageDigest +import java.util.{Collections, Locale, UUID} +import scala.annotation.tailrec +import scala.util.Try +import scala.util.control.NonFatal + +/** Counts describe terminal outcomes, not exactly-once execution of external model requests. + * `isComplete` means every input row has a committed outcome; `allSucceeded` additionally requires no row errors. + * Both are count-based: the last batch may be committed while `state` is still `running` and an owner remains. + * Neither flag proves task quiescence or permits owner release. Stop the previous application and confirm its + * work has quiesced before releasing its exact owner token; resume then finalizes without replaying committed rows. + */ +case class AIPromptOutputStatus(formatVersion: Int, + runId: String, + state: String, + inputCount: Long, + committedCount: Long, + pendingCount: Long, + rowErrorCount: Long, + ownerToken: Option[String], + resultsVersion: Long) { + def isComplete: Boolean = pendingCount == 0L + def allSucceeded: Boolean = isComplete && rowErrorCount == 0L + + def toJson: String = JsObject( + "formatVersion" -> JsNumber(formatVersion), "runId" -> JsString(runId), "state" -> JsString(state), + "inputCount" -> JsNumber(inputCount), "committedCount" -> JsNumber(committedCount), + "pendingCount" -> JsNumber(pendingCount), "rowErrorCount" -> JsNumber(rowErrorCount), + "ownerToken" -> ownerToken.map(JsString(_)).getOrElse(JsNull), "resultsVersion" -> JsNumber(resultsVersion), + "isComplete" -> JsBoolean(isComplete), "allSucceeded" -> JsBoolean(allSucceeded)).compactPrint + + override def toString: String = toJson +} + +private[ai] object AIPromptOutput { + private val OutputParams = Set("outputCol", "errorCol", "usageCol", "rawResponseCol", "responseIdCol") + private val SavedBindings = OutputParams + "messagesCol" + private val TypedOutputs: Map[String, DataType] = Map( + "errorCol" -> ErrorUtils.ErrorSchema, "usageCol" -> UsageUtils.UsageStructType, + "rawResponseCol" -> StringType, "responseIdCol" -> StringType) + private val AuthParams = Set("subscriptionKey", "AADToken", "CustomAuthHeader") + private val StorageParams = Set("outputPath", "outputBatchSize") + private val AutomaticBatchPartitions = 200 + private val SQLSemanticKeys = Set( + "spark.sql.session.timeZone", "spark.sql.caseSensitive", "spark.sql.timestampType", + "spark.sql.mapKeyDedupPolicy", "spark.sql.storeAssignmentPolicy", + "spark.sql.columnNameOfCorruptRecord", "spark.sql.chunkBase64String.enabled", + "spark.sql.ansi.enabled", "spark.sql.legacy.sizeOfNull") + private val SQLSemanticPrefixes = Seq( + "spark.sql.ansi.", "spark.sql.legacy.", "spark.sql.parser.", "spark.sql.function.", + "spark.sql.json.", "spark.sql.csv.", "spark.sql.jsonGenerator.", "spark.sql.decimalOperations.") + private val IgnoreEvent: String => Unit = _ => () + + private def enabledOutput(prompt: AIPrompt, name: String): Boolean = name match { + case "usageCol" => prompt.isSet(prompt.usageCol) + case "rawResponseCol" => prompt.isSet(prompt.rawResponseCol) + case "responseIdCol" => prompt.responseIdEnabled + case _ => true + } + + def storageSchema(schema: StructType): StructType = { + AIPromptStorageCodec.validate(schema) + AIPromptStorageCodec.validateFieldResolution(schema) + AIPromptStorageCodec.durable(schema).asInstanceOf[StructType] + } + + private def bindings(prompt: AIPrompt): Map[String, String] = + prompt.extractParamMap().toSeq.filter(p => SavedBindings(p.param.name) && enabledOutput(prompt, p.param.name)) + .map(p => p.param.name -> p.value.asInstanceOf[String]).toMap + + private def readBindings(manifest: AIPromptManifest): Map[String, String] = { + val bindings = AIJson.objectString(manifest.bindings).fields.map { case (name, value) => + require(SavedBindings(name), "Unknown AIPrompt stored output binding.") + val column = value match { + case JsString(text) if text.nonEmpty => text + case _ => throw new IllegalArgumentException("Invalid AIPrompt stored output binding.") + } + name -> column + } + require(Set("outputCol", "errorCol", "messagesCol").subsetOf(bindings.keySet), + "AIPrompt stored output bindings are incomplete.") + val outputs = bindings.filter { case (role, _) => OutputParams(role) } + require(outputs.values.map(_.toLowerCase(Locale.ROOT)).toSeq.distinct.size == outputs.size, + "AIPrompt stored output bindings must be distinct.") + val fields = manifest.outputType.fields.map(field => field.name -> field.dataType).toMap + outputs.foreach { case (role, column) => + require(fields.contains(column), "AIPrompt stored output binding is absent from the output schema.") + require(TypedOutputs.get(role).forall(dataType => + AIPromptStorageCodec.nullable(dataType) == AIPromptStorageCodec.nullable(fields(column))), + "AIPrompt stored output binding has an incompatible type.") + } + bindings + } + + // Escape code units before UTF-8 hashing, which would otherwise replace unpaired surrogates with '?'. + private def quotedString(value: String): String = JsString(value).compactPrint.flatMap { + case character if Character.isSurrogate(character) => f"\\u${character.toInt}%04x" + case character => character.toString + } + + private def canonical(json: JsValue): String = json match { + case JsObject(fields) => fields.toSeq.sortBy(_._1) + .map { case (name, value) => quotedString(name) + ":" + canonical(value) }.mkString("{", ",", "}") + case JsArray(values) => values.map(canonical).mkString("[", ",", "]") + case JsString(text) => quotedString(text) + case other => other.compactPrint + } + + private def orderedJsonIdentity(json: JsValue): String = json match { + case JsObject(fields) => fields.iterator + .map { case (name, value) => JsString(name).compactPrint + ":" + orderedJsonIdentity(value) } + .mkString("{", ",", "}") + case JsArray(values) => values.map(orderedJsonIdentity).mkString("[", ",", "]") + case JsNumber(number) if number.bigDecimal.scale == 0 => number.toString + case JsNumber(number) if number.signum == 0 => "0e0" + case JsNumber(number) => + val digits = number.bigDecimal.unscaledValue.toString + val significant = digits.reverse.dropWhile(_ == '0').reverse + // Avoid stripTrailingZeros scale overflow and expansion of large exponents. + val exponent = digits.length - significant.length - number.bigDecimal.scale.toLong + significant + "e" + exponent + case other => other.compactPrint + } + + private def credentialValue(value: Any): Option[JsValue] = value match { + case Right(binding: String) => Some(JsString(binding)) + case _ => None + } + + private def headerValue(param: Param[Any], value: Any, authHeaders: Set[String]): Option[JsValue] = value match { + case Left(headers: scala.collection.Map[_, _]) => + val filtered = headers.toSeq.collect { + case (key: String, v: String) if !authHeaders.exists(_.equalsIgnoreCase(key)) => + key -> JsString(v) + }.toMap + if (filtered.isEmpty) None else Some(JsObject(filtered)) + case _ => Some(param.jsonEncode(value).parseJson) + } + + private def paramValue(param: Param[Any], value: Any, authHeaders: Set[String]): Option[JsValue] = param.name match { + case "handler" => Some(JsString(AIPrompt.DefaultHttpPolicy)) + case name if AuthParams(name) => credentialValue(value) + case "customHeaders" | "telemHeaders" => headerValue(param, value, authHeaders) + // Preserve provider-owned object order and integer/decimal kind, not equivalent decimal spelling. + case "responseFormat" | "extraBody" | "questions" => + val encoded = value match { + case Left(json) => orderedJsonIdentity(JsObject("left" -> AIJson.fromAny(json))) + case _ => param.jsonEncode(value) + } + Some(JsString(encoded)) + case _ => Some(param.jsonEncode(value).parseJson) + } + + private def sqlSemantics(spark: SparkSession): JsObject = { + val registered = SQLConfInjections.registeredKeys + val configured = spark.conf.getAll.keys + val keys = (registered ++ configured ++ SQLSemanticKeys).filter { key => + SQLSemanticKeys(key) || SQLSemanticPrefixes.exists(key.startsWith) + }.toSet + JsObject(keys.map(key => key -> spark.conf.getOption(key).map(JsString(_)).getOrElse(JsNull)).toMap) + } + + private[ai] def validateSettings(prompt: AIPrompt): Unit = { + if (prompt.isDefined(prompt.outputPath)) AIPromptDelta.validatePath(prompt.getOutputPath) + require(AIPrompt.usesBoundedHandler(prompt), + "AIPrompt durable output requires the default HTTP handler; custom executable handlers cannot be fingerprinted.") + batchSize(prompt) + } + + private[ai] def fingerprint(prompt: AIPrompt, input: StructType, output: StructType, + spark: SparkSession): String = { + validateSettings(prompt) + // Library-owned defaults participate too; an upgrade can change recovery identity. + val config = prompt.extractParamMap().toSeq + .filter(p => !StorageParams(p.param.name) && enabledOutput(prompt, p.param.name)).flatMap { pair => + val param = pair.param.asInstanceOf[Param[Any]] + paramValue(param, pair.value, prompt.authenticationHeaderNames).map { json => + val explicit = !OutputParams(param.name) && prompt.isSet(param) + param.name -> JsObject("value" -> json, "explicit" -> JsBoolean(explicit)) + } + }.toMap + val content = canonical(JsObject( + "format" -> JsNumber(AIPromptDelta.FormatVersion), + "libraryVersion" -> JsString(BuildInfo.version), + "sparkVersion" -> JsString(spark.version), + "stageClass" -> JsString(prompt.getClass.getName), + "authenticationConfigured" -> JsBoolean(AIPrompt.hasConfiguredAuthentication(prompt)), + "params" -> JsObject(config), + "sqlSemantics" -> sqlSemantics(spark), + "input" -> AIPromptStorageCodec.schemaJson(input).parseJson, + "output" -> AIPromptStorageCodec.schemaJson(output).parseJson)) + AIPromptDelta.hex(MessageDigest.getInstance("SHA-256").digest(content.getBytes(StandardCharsets.UTF_8))) + } + + private def batchSize(prompt: AIPrompt): Int = { + val size = prompt.getOutputBatchSize + require(prompt.outputBatchSize.isValid(size), "AIPrompt outputBatchSize must be between 1 and 100000.") + size + } + + private def requireMatchingConfiguration(prompt: AIPrompt, spark: SparkSession, + manifest: AIPromptManifest): Unit = { + require(manifest.libraryVersion == BuildInfo.version && manifest.sparkVersion == spark.version, + "AIPrompt resume build or Spark runtime does not match the saved run. " + + "Restore the original SynapseML build and Spark runtime, or use loadOutput to read committed results.") + val output = storageSchema(if (manifest.state == "complete") prompt.inferenceOutputSchema(manifest.inputType) + else prompt.inferenceSchema(manifest.inputType)) + require(fingerprint(prompt, manifest.inputType, output, spark) == manifest.fingerprint, + "AIPrompt resume configuration, public output schema or Spark SQL semantics do not match the saved run. " + + "Restore the original SynapseML build, stage class, settings and Spark runtime, " + + "including library and SQL defaults and whether authentication is configured, before resuming.") + } + + private[ai] def transformObserved(prompt: AIPrompt, dataset: DataFrame, event: String => Unit): DataFrame = + AIPrompt.withSession(dataset.sparkSession) { + val effective = prompt.effectiveCopy() + effective.inferenceSchema(dataset.schema) + transformEffective(effective, dataset, event) + } + + def transformEffective(effective: AIPrompt, dataset: DataFrame, + event: String => Unit = IgnoreEvent): DataFrame = + AIPrompt.withSession(dataset.sparkSession) { + require(!dataset.isStreaming, "AIPrompt durable output does not support streaming inputs.") + val spark = dataset.sparkSession + val path = AIPromptDelta.qualify(spark, effective.getOutputPath) + AIPromptDelta.requireRuntime(spark, path) + batchSize(effective) + val inputSchema = storageSchema(dataset.schema) + val outputSchema = storageSchema(effective.inferenceSchema(inputSchema)) + effective.validateFabricModel(inputSchema) + val token = UUID.randomUUID().toString + val key = "_aiprompt_row_" + UUID.randomUUID().toString.replace("-", "") + require(!(inputSchema.fieldNames ++ outputSchema.fieldNames).exists(_.equalsIgnoreCase(key)), + "AIPrompt internal row key conflicts with an input or output column.") + val manifest = AIPromptManifest(UUID.randomUUID().toString, BuildInfo.version, spark.version, + fingerprint(effective, inputSchema, outputSchema, spark), AIPromptStorageCodec.schemaJson(inputSchema), + AIPromptStorageCodec.schemaJson(outputSchema), + JsObject(bindings(effective).map { case (name, value) => name -> JsString(value) }).compactPrint, + key, "initializing", None, None, Some(token)) + effective.validateDurableInput(dataset.schema, spark) + val keyed = dataset.select(monotonically_increasing_id().as(key) +: dataset.columns + .map(AIPromptDelta.column): _*) + val encoded = AIPromptStorageCodec.encode(keyed, manifest.keyed(inputSchema)) + AIPromptDelta.create(spark, path, manifest, () => event("before_control_create")) + event("control_created") + encoded.write.format("delta").mode("errorifexists").save(AIPromptDelta.child(path, "_input")) + event("input_committed") + val ready = initializeResults(effective, spark, path, manifest, token, event) + execute(effective, spark, path, ready, token, event) + } + + private def initializeResults(prompt: AIPrompt, spark: SparkSession, path: String, manifest: AIPromptManifest, + token: String, event: String => Unit): AIPromptManifest = { + val (inputVersion, snapshot) = AIPromptDelta.latest(spark, path, "_input") + AIPromptDelta.validateSchema(snapshot, manifest.keyed(manifest.inputType)) + if (inputVersion > 0L) { + // The new table starts at version 0; automatic compaction may commit before this read. + val changes = spark.sql(s"DESCRIBE HISTORY ${AIPromptDelta.table(path, "_input")}") + .filter(col("version") > 0L && col("version") <= inputVersion && + !col("operation").eqNullSafe(lit("OPTIMIZE"))) + require(changes.take(1).isEmpty, "The immutable AIPrompt input snapshot has been modified.") + } + prompt.validateDurableRows( + AIPromptStorageCodec.decode(snapshot, manifest.keyed(manifest.inputType)).drop(manifest.keyColumn)) + event("input_validated") + val count = AIPromptDelta.keyCount(snapshot, manifest.keyColumn) + val physical = AIPromptStorageCodec.physical(manifest.keyed(manifest.outputType)).asInstanceOf[StructType] + spark.createDataFrame(Collections.emptyList[Row](), physical).write.format("delta").mode("errorifexists") + .save(AIPromptDelta.child(path, "results")) + event("results_created") + require(AIPromptDelta.version(spark, path, "_input") == inputVersion, + "The immutable AIPrompt input snapshot has been modified.") + val ready = AIPromptDelta.ready(spark, path, manifest, token, inputVersion, count) + event("ready") + ready + } + + def resume(prompt: AIPrompt, spark: SparkSession): DataFrame = resumeObserved(prompt, spark, IgnoreEvent) + + private[ai] def resumeObserved(prompt: AIPrompt, spark: SparkSession, event: String => Unit): DataFrame = + AIPrompt.withSession(spark) { + val effective = prompt.effectiveCopy() + require(effective.isDefined(effective.outputPath), "resume requires outputPath; set the original run directory") + val path = AIPromptDelta.qualify(spark, effective.getOutputPath) + AIPromptDelta.requireRuntime(spark, path) + batchSize(effective) + val manifest = AIPromptDelta.control(spark, path) + AIPromptDelta.requireReady(manifest) + if (manifest.state != "complete") AIPromptDelta.requireUnowned(manifest) + effective.restoreDurableOutputDefaults(readBindings(manifest)) + requireMatchingConfiguration(effective, spark, manifest) + if (manifest.state != "complete") { + val saved = AIPromptDelta.input(spark, path, manifest) + val decoded = AIPromptStorageCodec.decode(saved, manifest.keyed(manifest.inputType)).drop(manifest.keyColumn) + effective.validateDurableInput(decoded.schema, spark) + effective.validateDurableRows(decoded) + } + val committed = ledger(spark, path, manifest) + if (manifest.state == "complete") { + AIPromptStorageCodec.decode(committed.results.drop(manifest.keyColumn), manifest.outputType) + } else { + if (committed.count < manifest.inputCount.get) effective.validateFabricModel(manifest.inputType) + event("before_acquire") + val token = AIPromptDelta.acquire(spark, path, manifest) + event("acquired") + execute(effective, spark, path, manifest, token, event) + } + } + + private case class Ledger(version: Long, results: DataFrame, count: Long) + + private[ai] def batchPartitions(rowCount: Int, shufflePartitions: String): Int = { + val normalized = shufflePartitions.trim + val configured = if (normalized == "auto") AutomaticBatchPartitions else { + Try(normalized.toInt).toOption.filter(_ > 0).getOrElse { + throw new IllegalArgumentException("spark.sql.shuffle.partitions must be a positive integer or 'auto'.") + } + } + math.min(rowCount, configured) + } + + private def ledger(spark: SparkSession, path: String, manifest: AIPromptManifest, + validateInputKeys: Boolean = true): Ledger = { + val input = AIPromptDelta.input(spark, path, manifest) + if (validateInputKeys) { + require(AIPromptDelta.keyCount(input, manifest.keyColumn) == manifest.inputCount.get, + "AIPrompt input count does not match the immutable snapshot.") + } + val (version, results) = AIPromptDelta.latest(spark, path, "results") + AIPromptDelta.validateSchema(results, manifest.keyed(manifest.outputType)) + val count = AIPromptDelta.keyCount(results, manifest.keyColumn) + val key = manifest.keyColumn + require(results.select(AIPromptDelta.column(key)).join(input.select(AIPromptDelta.column(key)), + Seq(key), "left_anti").limit(1).count() == 0L, "AIPrompt results contain unknown row keys.") + require(count <= manifest.inputCount.get && (manifest.state != "complete" || count == manifest.inputCount.get), + "AIPrompt completion state contradicts committed outcomes.") + Ledger(version, results, count) + } + + private def execute(prompt: AIPrompt, spark: SparkSession, path: String, manifest: AIPromptManifest, + token: String, event: String => Unit): DataFrame = { + // Construct inside guarded batch preparation, only if this run has pending requests. + lazy val http = AIPrompt.createHttpTransformer(prompt) + val infer = (batch: DataFrame) => prompt.infer(batch, Seq(manifest.keyColumn), http) + processBatches(prompt, spark, path, manifest, token, event, ledger(spark, path, manifest), infer) + // Failed Spark jobs and uncertain commits retain ownership; they do not prove task quiescence. + AIPromptDelta.complete(spark, path, manifest, token) + event("completed") + loadOutput(spark, path) + } + + @tailrec + private def processBatches(prompt: AIPrompt, spark: SparkSession, path: String, manifest: AIPromptManifest, + token: String, event: String => Unit, committed: Ledger, + infer: DataFrame => DataFrame): Unit = { + AIPromptDelta.assertOwner(spark, path, manifest, token) + val input = AIPromptDelta.input(spark, path, manifest) + event("before_ledger_version_check") + require(AIPromptDelta.version(spark, path, "results") == committed.version, + "AIPrompt committed outcomes changed outside the current owner.") + val key = manifest.keyColumn + val keys = input.select(AIPromptDelta.column(key)) + .join(committed.results.select(AIPromptDelta.column(key)), Seq(key), "left_anti") + .take(batchSize(prompt)).map(_.getLong(0)) + if (keys.nonEmpty) { + val selectedKeys = keys.map(Long.box) + val prepared = try { + requireMatchingConfiguration(prompt, spark, manifest) + val savedBatch = input.filter(AIPromptDelta.column(key).isin(selectedKeys: _*)) + val partitions = batchPartitions(keys.length, spark.conf.get("spark.sql.shuffle.partitions")) + val batch = AIPromptStorageCodec.decode(savedBatch, manifest.keyed(manifest.inputType)) + .repartition(partitions) + val inferred = infer(batch) + val selected = inferred.select(manifest.keyed(manifest.outputType).fieldNames.map(AIPromptDelta.column): _*) + AIPromptStorageCodec.encode(selected, manifest.keyed(manifest.outputType)) + } catch { + case NonFatal(error) if !Thread.currentThread().isInterrupted => + // Previous actions finished successfully; this block only builds the next lazy plan. + try AIPromptDelta.release(spark, path, manifest, token) catch { + case releaseError: InterruptedException => + Thread.currentThread().interrupt() + error.addSuppressed(releaseError) + case NonFatal(releaseError) => error.addSuppressed(releaseError) + } + throw error + } + AIPromptDelta.assertOwner(spark, path, manifest, token) + event("before_batch_commit") + prepared.write.format("delta").mode("append").save(AIPromptDelta.child(path, "results")) + event("batch_committed") + val after = ledger(spark, path, manifest, validateInputKeys = false) + require(after.count == committed.count + keys.length && + after.results.filter(AIPromptDelta.column(key).isin(selectedKeys: _*)).count() == keys.length, + "AIPrompt committed batch cardinality does not match the selected input keys.") + processBatches(prompt, spark, path, manifest, token, event, after, infer) + } + } + + private def open(spark: SparkSession, outputPath: String): (String, AIPromptManifest) = { + val path = AIPromptDelta.qualify(spark, outputPath) + AIPromptDelta.requireRuntime(spark, path) + val manifest = AIPromptDelta.control(spark, path) + AIPromptDelta.requireReady(manifest) + readBindings(manifest) + path -> manifest + } + + def loadOutput(spark: SparkSession, outputPath: String): DataFrame = { + val (path, manifest) = open(spark, outputPath) + val committed = ledger(spark, path, manifest) + AIPromptStorageCodec.decode(committed.results.drop(manifest.keyColumn), manifest.outputType) + } + + def getOutputStatus(spark: SparkSession, outputPath: String): AIPromptOutputStatus = { + val (path, manifest) = open(spark, outputPath) + val committed = ledger(spark, path, manifest) + val errors = committed.results.filter(AIPromptDelta.column(readBindings(manifest)("errorCol")).isNotNull).count() + AIPromptOutputStatus(AIPromptDelta.FormatVersion, manifest.runId, manifest.state, + manifest.inputCount.get, committed.count, manifest.inputCount.get - committed.count, + errors, manifest.owner, committed.version) + } + + def releaseOutputOwner(spark: SparkSession, outputPath: String, + expectedOwnerToken: String, confirmStopped: Boolean): Unit = { + require(confirmStopped, "Confirm that the previous Spark application has stopped before releasing its owner.") + require(Option(expectedOwnerToken).exists(_.nonEmpty), "An exact nonempty expected owner token is required.") + val path = AIPromptDelta.qualify(spark, outputPath) + AIPromptDelta.requireRuntime(spark, path) + val manifest = AIPromptDelta.control(spark, path) + AIPromptDelta.release(spark, path, manifest, expectedOwnerToken) + } +} 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..165c1ee975b --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptParams.scala @@ -0,0 +1,106 @@ +// 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.{IntParam, Param, ParamValidators} +import spray.json.DefaultJsonProtocol._ +import spray.json.{JsObject, JsValue} + +trait AIPromptParams extends HasServiceParams with HasInputCol { + private val maximumOutputBatchSize = 100000 + + 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) + + val outputPath: Param[String] = new Param[String]( + this, "outputPath", "Fresh caller-owned Delta run directory; enables eager, resumable input and result storage.") + def setOutputPath(value: String): this.type = set(outputPath, value) + def getOutputPath: String = $(outputPath) + + val outputBatchSize: IntParam = new IntParam( + this, "outputBatchSize", "Maximum rows per durable Delta batch.", + ParamValidators.inRange(1, maximumOutputBatchSize, lowerInclusive = true, upperInclusive = true)) + def setOutputBatchSize(value: Int): this.type = set(outputBatchSize, value) + def getOutputBatchSize: Int = $(outputBatchSize) + + setDefault(inputMode -> "text", outputBatchSize -> 1000) +} 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..6baffb4e44d --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPython.scala @@ -0,0 +1,144 @@ +// 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 BridgeClass = classOf[AIPrompt].getName + 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 = + s""" + |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) + | + |def resume(self, spark): + | from pyspark.sql import DataFrame + | self._transfer_params_to_java() + | return DataFrame(self._java_obj.resume(spark._jsparkSession), spark) + | + |@staticmethod + |def loadOutput(spark, outputPath): + | from pyspark.sql import DataFrame + | bridge = spark._jvm.$BridgeClass + | return DataFrame(bridge.loadOutput(spark._jsparkSession, outputPath), spark) + | + |@staticmethod + |def getOutputStatus(spark, outputPath): + | '''Read committed outcome counts and run ownership without model calls. + | + | ``isComplete`` means no rows are pending; ``allSucceeded`` also requires + | no row errors. These count flags can be true while ``state`` is still + | ``running`` and ``ownerToken`` is present. Neither proves task quiescence. + | Stop the previous application and confirm its work has quiesced before + | releasing its exact owner token; resume finalizes without replaying + | committed rows. + | ''' + | import json + | bridge = spark._jvm.$BridgeClass + | return json.loads(bridge.getOutputStatusJson(spark._jsparkSession, outputPath)) + | + |@staticmethod + |def releaseOutputOwner(spark, outputPath, expectedOwnerToken, confirmStopped=False): + | if confirmStopped is not True: + | raise ValueError("Stop the previous Spark application and set confirmStopped=True") + | bridge = spark._jvm.$BridgeClass + | bridge.releaseOutputOwner(spark._jsparkSession, outputPath, expectedOwnerToken, True) + |""".stripMargin +} diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageCodec.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageCodec.scala new file mode 100644 index 00000000000..c6a9a855ac7 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageCodec.scala @@ -0,0 +1,189 @@ +// 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.SQLDataTypes +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.{lit, when} +import org.apache.spark.sql.types._ +import spray.json._ + +/** The on-disk schema is deliberately independent of Spark's arbitrary UDT deserializer. */ +private[ai] object AIPromptStorageCodec { + private val VectorTag = "aiprompt-vector-v1" + private val PrimitiveTypes = Seq(StringType, BinaryType, BooleanType, ByteType, ShortType, + IntegerType, LongType, FloatType, DoubleType, DateType, TimestampType, TimestampNTZType) + .map(t => t.typeName -> t).toMap + // Delta strips these source-column behaviors when reading a materialized snapshot. + private val SourceMetadataKeys = Seq( + "__autoGeneratedAlias", "__metadata_col", "__supports_qualified_star", "__qualified_access_only", + "__file_source_metadata_col", "__file_source_constant_metadata_col", "__file_source_generated_metadata_col", + "delta.generationExpression", "delta.identity.allowExplicitInsert", "delta.identity.highWaterMark", + "delta.identity.start", "delta.identity.step") + + def durable(dataType: DataType): DataType = dataType match { + case s: StructType => StructType(s.fields.map { field => + val metadata = new MetadataBuilder().withMetadata(field.metadata) + SourceMetadataKeys.foreach(metadata.remove) + field.copy(dataType = durable(field.dataType), nullable = true, metadata = metadata.build()) + }) + case ArrayType(t, _) => ArrayType(durable(t), containsNull = true) + case MapType(k, v, _) => MapType(durable(k), durable(v), valueContainsNull = true) + case other => other + } + + def nullable(dataType: DataType): DataType = AIValueCodec.nullable(dataType) + + def validate(schema: StructType): Unit = { + validateType(schema, 0) + decodeType(parseSchemaJson(encodeType(schema).compactPrint)) + } + + private def parseSchemaJson(json: String): JsValue = try { + AIJson.parse(json) + } catch { + case _: IllegalArgumentException => + throw new IllegalArgumentException( + "AIPrompt durable schema contains invalid JSON or exceeds supported parsing limits (including metadata).") + } + + + def validateFieldResolution(dataType: DataType): Unit = dataType match { + case s: StructType => + AIColumnPlan.distinct(s.fieldNames.toSeq, "durable fields") + s.fields.foreach(field => validateFieldResolution(field.dataType)) + case ArrayType(t, _) => validateFieldResolution(t) + case MapType(k, v, _) => validateFieldResolution(k); validateFieldResolution(v) + case _ => + } + + private def validateStruct(schema: StructType, depth: Int): Unit = { + schema.fields.foreach { field => + require(field.name.nonEmpty && !field.name.exists(" ,;{}()\n\t=".contains(_)), + "AIPrompt durable storage requires Delta-compatible column names (including nested fields).") + // Stored structs wrap each field type in a struct array and field object. + validateType(field.dataType, depth + 3) + } + require(schema.fieldNames.map(_.toLowerCase(java.util.Locale.ROOT)).distinct.length == schema.length, + "AIPrompt durable storage does not support case-insensitive duplicate column names.") + } + + private def validateType(dataType: DataType, depth: Int): Unit = { + require(depth <= AIJson.MaxDepth, "AIPrompt durable schema exceeds the supported depth.") + dataType match { + case s: StructType => + require(depth == 0 || s.nonEmpty, "AIPrompt durable storage does not support nested empty structs.") + validateStruct(s, depth) + case ArrayType(t, _) => validateType(t, depth + 1) + case MapType(k, v, _) => validateType(k, depth + 1); validateType(v, depth + 1) + case other => validateLeaf(other) + } + } + + private def validateLeaf(dataType: DataType): Unit = dataType match { + case t if t == SQLDataTypes.VectorType => + case _: UserDefinedType[_] => + throw new IllegalArgumentException("AIPrompt durable storage supports only the Spark ML vector UDT.") + case decimal: DecimalType => + require(decimal.scale >= 0, "AIPrompt durable storage supports only nonnegative decimal scales.") + case t if PrimitiveTypes.values.toSeq.contains(t) => + case _ => throw new IllegalArgumentException("Unsupported SQL type for AIPrompt durable storage.") + } + + def physical(dataType: DataType): DataType = dataType match { + case t if t == SQLDataTypes.VectorType => AIValueCodec.VectorStorage + case s: StructType => StructType(s.fields.map(f => f.copy(dataType = physical(f.dataType), nullable = true))) + case ArrayType(t, _) => ArrayType(physical(t), containsNull = true) + case MapType(k, v, _) => MapType(physical(k), physical(v), valueContainsNull = true) + case other => other + } + + def encode(frame: DataFrame, schema: StructType): DataFrame = convertFrame(frame, schema, encode = true) + + def decode(frame: DataFrame, schema: StructType): DataFrame = convertFrame(frame, schema, encode = false) + + private def convertFrame(frame: DataFrame, schema: StructType, encode: Boolean): DataFrame = { + val target = if (encode) physical(schema).asInstanceOf[StructType] else schema + val columns = schema.fields.map { field => + val column = AIColumnPlan.column(field.name) + val sourceField = frame.schema(field.name) + val resultField = target(field.name) + val resultType = resultField.dataType + val converted = if (AIValueCodec.containsVector(field.dataType)) { + AIValueCodec.convertColumn(column, field.dataType, field.nullable, encode) + } else if (sourceField.dataType == resultType) column else column.cast(resultType) + // Widen nullability without a row encoder; Catalyst removes the always-true branch. + val widened = if (resultField.nullable && !sourceField.nullable) when(lit(true), converted) else converted + widened.as(field.name, field.metadata) + } + AIValueCodec.reconcileSchema(frame.select(columns: _*), target) + } + + def schemaJson(schema: StructType): String = { + validate(schema) + encodeType(nullable(schema)).compactPrint + } + + private def encodeType(dataType: DataType): JsValue = dataType match { + case t if t == SQLDataTypes.VectorType => JsString(VectorTag) + case s: StructType => JsObject("struct" -> JsArray(s.fields.map { f => + JsObject("name" -> JsString(f.name), "type" -> encodeType(f.dataType), + "metadata" -> parseSchemaJson(f.metadata.json)) + }.toVector)) + case ArrayType(t, _) => JsObject("array" -> encodeType(t)) + case MapType(k, v, _) => JsObject("key" -> encodeType(k), "value" -> encodeType(v)) + case d: DecimalType => JsObject("precision" -> JsNumber(d.precision), "scale" -> JsNumber(d.scale)) + case other => JsString(other.typeName) + } + + def readSchema(json: String): StructType = { + val schema = decodeType(parseSchemaJson(json)) match { + case value: StructType => value + case _ => throw new IllegalArgumentException("AIPrompt stored schema must be a struct.") + } + validate(schema) + schema + } + + private def decodeType(json: JsValue): DataType = json match { + case JsString(VectorTag) => SQLDataTypes.VectorType + case JsString(name) => PrimitiveTypes.getOrElse(name, + throw new IllegalArgumentException("Unknown AIPrompt stored type; arbitrary UDTs are not loaded.")) + case JsObject(fields) => decodeContainer(fields) + case _ => throw new IllegalArgumentException("Invalid or unsupported AIPrompt stored schema.") + } + + private def decodeContainer(fields: Map[String, JsValue]): DataType = fields.toSeq.sortBy(_._1) match { + case Seq(("struct", JsArray(elements))) => + StructType(elements.map(decodeField)) + case Seq(("array", value)) => ArrayType(decodeType(value), true) + case Seq(("key", key), ("value", value)) => MapType(decodeType(key), decodeType(value), true) + case Seq(("precision", JsNumber(precision)), ("scale", JsNumber(scale))) => + require(precision.isValidInt && scale.isValidInt, "Invalid AIPrompt stored decimal type.") + try DecimalType(precision.toInt, scale.toInt) catch { + case _: IllegalArgumentException | _: ArithmeticException | + _: org.apache.spark.sql.AnalysisException | _: org.apache.spark.SparkException => + throw new IllegalArgumentException("Invalid AIPrompt stored decimal type.") + } + case _ => throw new IllegalArgumentException("Invalid or unsupported AIPrompt stored schema.") + } + + private def decodeField(json: JsValue): StructField = json match { + case JsObject(fields) if fields.keySet == Set("name", "type", "metadata") => + (fields("name"), fields("metadata")) match { + case (JsString(name), metadata: JsObject) => + val decoded = try Metadata.fromJson(metadata.compactPrint) catch { + case _: ClassCastException | _: IllegalArgumentException => + throw new IllegalArgumentException("Invalid AIPrompt stored field metadata.") + case error: org.apache.spark.SparkThrowable if Set( + "_LEGACY_ERROR_TEMP_2120", "_LEGACY_ERROR_TEMP_2121")(error.getErrorClass) => + throw new IllegalArgumentException("Invalid AIPrompt stored field metadata.") + } + require(AIJson.parse(decoded.json) == metadata, "Invalid AIPrompt stored field metadata.") + StructField(name, decodeType(fields("type")), nullable = true, decoded) + case _ => throw new IllegalArgumentException("Invalid AIPrompt stored schema field.") + } + case _ => throw new IllegalArgumentException("Invalid AIPrompt stored schema field.") + } +} 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..233c1bf3cbf --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodec.scala @@ -0,0 +1,150 @@ +// 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} +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[ai] 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) when(lit(true), converted) 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[ai] 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[ai] 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/HasAIPromptRecovery.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/HasAIPromptRecovery.scala new file mode 100644 index 00000000000..3b62d5f8a65 --- /dev/null +++ b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/ai/HasAIPromptRecovery.scala @@ -0,0 +1,25 @@ +// 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.{Param, Params} +import org.apache.spark.sql.{DataFrame, SparkSession} + +private[ai] trait HasAIPromptRecovery extends Params { self: AIPrompt => + def resume(spark: SparkSession): DataFrame = AIPrompt.withSession(spark) { + logVerb("resume", AIPromptOutput.resume(this, spark)) + } + + private[ai] def restoreDurableOutputDefaults(bindings: Map[String, String]): Unit = { + val roles = Set("outputCol", "errorCol", "usageCol", "rawResponseCol", "responseIdCol", "messagesCol") + require(bindings.keySet.subsetOf(roles), "Manifest contains unsupported output bindings") + bindings.foreach { case (name, saved) => + val param = getParam(name).asInstanceOf[Param[String]] + if (isSet(param)) { + require(getOrDefault(param) == saved, s"Explicit $name does not match the saved output binding") + } else if (Set("usageCol", "rawResponseCol", "responseIdCol")(name)) set(param, saved) + else setDefault(param -> saved) + } + } +} 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 68b999a17de..7d40ced9e87 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 @@ -4,16 +4,46 @@ package com.microsoft.azure.synapse.ml.services.openai import com.microsoft.azure.synapse.ml.io.binary.BinaryFileReader +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.{Path => HPath} +import org.apache.spark.{TaskContext, TaskKilledException} -import java.io.ByteArrayInputStream +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} import java.net.{URI, URLConnection} +import java.nio.channels.ClosedByInterruptException import java.nio.charset.StandardCharsets +import java.util.concurrent.CancellationException import java.util.{Base64, Locale} -import scala.util.Try +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'") + + // Preserve legacy row diagnostics without retaining storage-client cancellation details. + 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") + } + + private[openai] def readAttachment(fileName: String)(read: => Array[Byte]): Array[Byte] = try { + read + } catch { + case _: InterruptedException | _: ClosedByInterruptException => interruptedRead() + case NonFatal(_) if Thread.currentThread().isInterrupted => interruptedRead() + case NonFatal(_) if Option(TaskContext.get()).exists(_.isInterrupted()) => + throw new TaskKilledException("Attachment read canceled") + case _: CancellationException => throw new CancelledAttachmentReadException(fileName) + case error: AttachmentSizeException => throw error + case NonFatal(_) => throw new AttachmentReadException(fileName) + } private case class DataUrlAttachment( fileName: String, @@ -23,6 +53,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") } } @@ -96,7 +127,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 @@ -217,6 +248,59 @@ 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) + + private def readBoundedFile(path: HPath, limit: Double, conf: Configuration): Array[Byte] = { + val limitBytes = (limit * 1024 * 1024).toLong + if (limitBytes == Long.MaxValue) BinaryFileReader.readSingleFileBytes(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) = @@ -224,11 +308,10 @@ private[openai] object OpenAIAttachmentUtils { readDataUrl(filePathStr, sizeLimitMB) } else { val fileName = attachmentFilename(filePathStr) - val fileBytes = try { - BinaryFileReader.readSingleFileBytes(new HPath(filePathStr)) - } catch { - case NonFatal(_) => - throw new IllegalArgumentException(s"Unable to read attachment '$fileName'") + val fileBytes = readAttachment(fileName) { + val path = new HPath(filePathStr) + if (bounded && sizeLimitMB.isDefined) readBoundedFile(path, sizeLimitMB.get, conf) + else BinaryFileReader.readSingleFileBytes(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 1509cddda69..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 @@ -12,6 +12,7 @@ import com.microsoft.azure.synapse.ml.services._ import com.microsoft.azure.synapse.ml.services.aifoundry.{AIFoundryChatCompletion, HasAIFoundryTextParamsExtended} import org.apache.hadoop.conf.Configuration import org.apache.http.entity.AbstractHttpEntity +import org.apache.spark.TaskKilledException import org.apache.spark.ml.{ComplexParamsReadable, ComplexParamsWritable, Transformer} import org.apache.spark.ml.param.{BooleanParam, DoubleParam, Param, ParamMap, ParamValidators} import org.apache.spark.ml.util.{Identifiable, MLWriter} @@ -113,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)) @@ -129,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) @@ -165,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()) @@ -185,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) @@ -274,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) @@ -493,6 +505,7 @@ class OpenAIPrompt(override val uid: String) extends Transformer ) } match { case scala.util.Success(msgs) => (msgs, null) //scalastyle:ignore null + case scala.util.Failure(e: TaskKilledException) => throw e case scala.util.Failure(e) => (null, e.getMessage) //scalastyle:ignore null } } @@ -577,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 => @@ -591,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" => @@ -641,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) @@ -710,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) { @@ -723,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/aiprompt_fabric_recovery_e2e.py b/cognitive/src/test/python/synapsemltest/services/aiprompt_fabric_recovery_e2e.py new file mode 100644 index 00000000000..9947c5c81e9 --- /dev/null +++ b/cognitive/src/test/python/synapsemltest/services/aiprompt_fabric_recovery_e2e.py @@ -0,0 +1,222 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +"""Opt-in same-application recovery probe using synthetic native HTTP replies. + +This script does not qualify recovery across managed Spark application restarts. + +Supply the PR core, cognitive and cognitive test JARs to Fabric Spark, using +--expected-core-jar, --expected-cognitive-jar and --expected-tests-jar respectively. +Each expected name must exactly match the URL-decoded loaded JAR basename. +Run without Python optimization so the qualification assertions remain enabled. +The runner owns the scratch Lakehouse lifecycle. This script does not provision +resources, acquire model credentials, or claim live-provider qualification. +""" + +import argparse +import json +import threading +import uuid +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" + ), + "recovery": source( + "com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricRecovery" + ), + "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), + ("recovery", args.expected_tests_jar), + ("probe", args.expected_tests_jar), + ): + basename = unquote(urlsplit(provenance[role]).path.rsplit("/", 1)[-1]) + assert expected and basename == expected, provenance + 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 + + relative = "Files/aiprompt_e2e/" + uuid.uuid4().hex + hadoop_path = jvm.org.apache.hadoop.fs.Path(relative) + filesystem = hadoop_path.getFileSystem( + spark.sparkContext._jsc.hadoopConfiguration() + ) + qualified = str(filesystem.makeQualified(hadoop_path)) + assert qualified.startswith(("abfs:", "abfss:")), ( + "This scenario must exercise OneLake/ABFS, not local storage", + qualified, + ) + jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricRecovery.verifyLocalStorageRejected( + spark._jsparkSession + ) + implicit_apis = json.loads( + jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricProbe.verifyImplicitEndpoints( + spark._jsparkSession + ) + ) + assert implicit_apis == ["chat_completions", "responses", "embeddings"] + 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"], + }, + } + companion = jvm.com.microsoft.azure.synapse.ml.services.ai.AIPrompt + + def configured(): + return ( + 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) + .setOutputPath(qualified) + .setOutputBatchSize(2) + ) + + 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: + data = spark.createDataFrame( + [(1, "outage"), (2, "outage"), (3, None), (4, "billing")], + "id INT, text STRING", + ).repartition(2) + stage = configured() + expected = stage.transformSchema(data._jdf.schema()).json() + jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricRecovery.stopAfterFirstBatch( + stage, data._jdf + ) + partial = json.loads( + companion.getOutputStatusJson(spark._jsparkSession, qualified) + ) + assert partial["committedCount"] == 2 and partial["pendingCount"] == 2 + assert partial["ownerToken"] is not None + committed = DataFrame( + companion.loadOutput(spark._jsparkSession, qualified), spark + ).collect() + before_resume = len(sends) + assert len(committed) == 2 + companion.releaseOutputOwner( + spark._jsparkSession, qualified, partial["ownerToken"], True + ) + result = DataFrame(configured().resume(spark._jsparkSession), spark) + assert json.loads(result.schema.json()) == json.loads(expected), ( + expected, + result.schema.json(), + ) + assert result.count() == 4 + assert len(sends) == 3 + rows = result.collect() + 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 + ) + status = json.loads( + companion.getOutputStatusJson(spark._jsparkSession, qualified) + ) + reopened = DataFrame( + companion.loadOutput(spark._jsparkSession, qualified), spark + ) + assert reopened.schema == result.schema + assert reopened.count() == 4 + resumed = DataFrame(configured().resume(spark._jsparkSession), spark) + assert resumed.schema == result.schema and resumed.count() == 4 + assert len(sends) == 3, "Reading/resuming committed rows must not send HTTP" + evidence = { + "applicationId": spark.sparkContext.applicationId, + "sparkVersion": spark.version, + "classSources": provenance, + "storageScheme": qualified.split(":", 1)[0], + "rows": len(rows), + "physicalSends": len(sends), + "committedBeforeRecovery": partial["committedCount"], + "sendsBeforeRecovery": before_resume, + "matchingOwnerRecovery": True, + "status": status, + "schemaEqual": True, + "duplicateRowsRetained": True, + "credentialFreeMock": True, + "implicitFabricApis": implicit_apis, + "nonLocalFileStorageRejected": True, + } + print("SYNAPSEML_FABRIC_E2E_RESULT=" + 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/python/synapsemltest/services/test_AIPromptRecovery.py b/cognitive/src/test/python/synapsemltest/services/test_AIPromptRecovery.py new file mode 100644 index 00000000000..212d3cd7168 --- /dev/null +++ b/cognitive/src/test/python/synapsemltest/services/test_AIPromptRecovery.py @@ -0,0 +1,205 @@ +# Copyright (C) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE in project root for information. + +import json +import os +import tempfile +import unittest +from decimal import Decimal +from pathlib import Path + +from pyspark.errors import IllegalArgumentException +from pyspark.sql import DataFrame +from synapse.ml.services.ai import AIPrompt, AISystemOne +from synapse.ml.services.openai import OpenAIDecisions +from synapsemltest.services.test_AIPrompt import endpoint, spark + + +class TestAIPromptRecoveryConfirmation(unittest.TestCase): + def test_administrative_release_requires_explicit_stopped_confirmation(self): + with self.assertRaisesRegex(ValueError, "Stop the previous Spark application"): + AIPrompt.releaseOutputOwner(spark, "unused", "unused") + + +@unittest.skipUnless( + os.environ.get("AIPROMPT_RUN_PYTHON_DELTA_TESTS") == "true", + "Requires caller-configured Delta and the cognitive test JAR", +) +class TestAIPromptDurable(unittest.TestCase): + def test_generated_durable_methods_recover_without_replaying_committed_rows(self): + cases = ( + ( + AIPrompt, + { + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "answer"}, + } + ] + }, + ), + ( + AISystemOne, + {"model": "model", "answers": {"check": {"result": True}}}, + ), + ( + OpenAIDecisions, + { + "model": "model", + "answers": [ + {"type": "predicate", "name": "check", "probability": 0.9} + ], + }, + ), + ) + texts = ['{"note":"first"}', '{"note":"second"}', '{"note":"third"}'] + data = spark.createDataFrame([(text,) for text in texts], "text STRING") + numeric_text = ( + '{"trailing":2.50,"integer":1.00,"exponent":1e2,"zero":-0.00,' + '"precise":0.12345678901234567890123456789,"tiny":1e-400,"huge":1e400}' + ) + numbers = json.loads(numeric_text, parse_float=Decimal) + for stage_type, reply in cases: + with self.subTest(stage=stage_type.__name__), endpoint( + reply, exact_numbers=True + ) as ( + url, + requests, + ), tempfile.TemporaryDirectory() as directory: + root = str(Path(directory) / "run") + + def configured(): + stage = stage_type( + url=url if stage_type is AISystemOne else url + "/v1", + model="model", + outputCol="answer", + errorCol="error", + ) + stage.setExtraBody('{"numeric":' + numeric_text + "}") + if stage_type is not AISystemOne: + stage.setInputCol("text") + if stage_type is AIPrompt: + stage.setResponseSchema( + { + "type": "object", + "properties": {"value": {"type": "number"}}, + "numeric": numbers, + } + ) + elif stage_type is AISystemOne: + stage.setQuestions( + { + "check": { + "type": "noul", + "instructions": "Check", + "numeric": numbers, + } + } + ) + stage.setStateCol("text").setStateMode("json") + self.assertEqual(stage.getInputMode(), "json") + elif stage_type is OpenAIDecisions: + stage.setQuestions( + [ + { + "type": "predicate", + "name": "check", + "instructions": "Check", + "numeric": numbers, + } + ] + ) + stage.setStateMode("text") + self.assertEqual(stage.getInputMode(), "text") + return stage + + original = configured().setOutputPath(root).setOutputBatchSize(1) + original._transfer_params_to_java() + spark._jvm.com.microsoft.azure.synapse.ml.services.ai.AIPromptFabricRecovery.stopAfterFirstBatch( + original._java_obj, data._jdf + ) + partial = stage_type.getOutputStatus(spark, root) + self.assertIsInstance(partial, dict) + self.assertEqual(partial["state"], "running") + self.assertEqual( + [ + partial[key] + for key in ("inputCount", "committedCount", "pendingCount") + ], + [3, 1, 2], + ) + self.assertIs(partial["isComplete"], False) + self.assertIs(partial["allSucceeded"], False) + owner = partial["ownerToken"] + self.assertIsInstance(owner, str) + self.assertTrue(owner) + committed = stage_type.loadOutput(spark, root) + self.assertIsInstance(committed, DataFrame) + self.assertEqual(committed.count(), 1) + self.assertEqual(len(requests), 1) + with self.assertRaisesRegex( + IllegalArgumentException, "ownership conflict" + ): + stage_type.releaseOutputOwner( + spark, root, owner + "-wrong", confirmStopped=True + ) + self.assertEqual( + stage_type.getOutputStatus(spark, root)["ownerToken"], owner + ) + stage_type.releaseOutputOwner(spark, root, owner, confirmStopped=True) + self.assertIsNone(stage_type.getOutputStatus(spark, root)["ownerToken"]) + fresh = configured().setParams(outputPath=root, outputBatchSize=2) + parameters = ( + "ExtraBody", + "ResponseFormat" if stage_type is AIPrompt else "Questions", + ) + for suffix in parameters: + value = getattr(original, "get" + suffix)() + getattr(fresh, "set" + suffix)(value) + name = suffix[0].lower() + suffix[1:] + parameter = original._java_obj.getParam(name) + self.assertNotEqual( + parameter.jsonEncode( + original._java_obj.getOrDefault(parameter) + ), + parameter.jsonEncode( + fresh._java_obj.getOrDefault(fresh._java_obj.getParam(name)) + ), + ) + self.assertIsInstance( + fresh.getExtraBody()["numeric"]["trailing"], float + ) + self.assertIsInstance( + fresh.getExtraBody()["numeric"]["exponent"], float + ) + resumed = fresh.resume(spark) + self.assertIsInstance(resumed, DataFrame) + rows = resumed.orderBy("text").collect() + self.assertEqual([row.text for row in rows], sorted(texts)) + self.assertTrue( + all(row.error is None and row.answer is not None for row in rows) + ) + self.assertEqual( + stage_type.loadOutput(spark, root).orderBy("text").collect(), rows + ) + self.assertEqual(fresh.resume(spark).orderBy("text").collect(), rows) + self.assertEqual(original.resume(spark).orderBy("text").collect(), rows) + complete = stage_type.getOutputStatus(spark, root) + self.assertEqual(complete["runId"], partial["runId"]) + self.assertEqual(complete["state"], "complete") + self.assertEqual(complete["committedCount"], 3) + self.assertEqual(complete["pendingCount"], 0) + self.assertEqual(complete["rowErrorCount"], 0) + self.assertIsNone(complete["ownerToken"]) + self.assertIs(complete["isComplete"], True) + self.assertIs(complete["allSucceeded"], True) + self.assertGreater( + complete["resultsVersion"], partial["resultsVersion"] + ) + self.assertEqual(len(requests), 3) + self.assertTrue(all(body["numeric"] == numbers for _, body in requests)) + + +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..f94fa5b6961 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptCancellationSuite.scala @@ -0,0 +1,120 @@ +// 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.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 } + assert(killed.exists(_.reason == "HTTP request task canceled")) + } 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/AIPromptDurableBindingsSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableBindingsSuite.scala new file mode 100644 index 00000000000..d4d13c4dbfb --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableBindingsSuite.scala @@ -0,0 +1,711 @@ +// 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.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import spray.json._ + +import java.io.File +import java.util.Collections +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptDurableBindingsSuite extends AIPromptOutputTestBase { + private val chatReply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + + test("durable custom URL roots ignore unrelated global endpoint changes during recovery") { + delta() + val previous = GlobalParams.getGlobalParam(URLKey) + try { + GlobalParams.setGlobalParam(URLKey, "https://first.example.invalid") + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = new AIPrompt().setCustomUrlRoot(server.url).setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root).setOutputBatchSize(2) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("one", "two", "three").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 2 && status.pendingCount == 1 && server.requests.get() == 2) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + GlobalParams.setGlobalParam(URLKey, "https://second.example.invalid") + assert(stage.resume(session).count() == 3 && server.requests.get() == 3) + GlobalParams.resetGlobalParam(URLKey) + assert(stage.resume(session).count() == 3 && server.requests.get() == 3) + assert(AIPrompt.getOutputStatus(session, root).ownerToken.isEmpty) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + + Seq("chat_completions", "responses", "systemone", "decisions", "embeddings").foreach { api => + test(s"durable $api template evaluation failures stop initialization before any HTTP requests") { + delta() + val session = spark.newSession() + session.conf.set("spark.sql.legacy.timeParserPolicy", "EXCEPTION") + import session.implicits._ + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val root = path() + val stage = new AIPrompt().setApiType(api) + .setUrl(if (api == "systemone") server.url else server.url + "/v1").setModel("model") + .setPromptTemplate("{cast(to_date(text, 'MM/dd/yyyy') as string)}") + .setOutputPath(root).setOutputBatchSize(1) + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + else if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + val input = Seq("01/02/2020", "1/2/2020").toDF("text") + val error = intercept[org.apache.spark.SparkException](stage.transform(input)) + assert(Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null) + .exists(_.getClass.getSimpleName == "SparkUpgradeException")) + assert(server.requests.isEmpty) + val qualified = AIPromptDelta.qualify(session, root) + assert(AIPromptDelta.control(session, qualified).state == "initializing") + assert(!new File(root, "results").exists()) + val resumeError = intercept[IllegalArgumentException](stage.resume(session)) + assert(resumeError.getMessage.contains("initialization is incomplete") && server.requests.isEmpty) + } + } + } + + test("durable fingerprints distinguish explicit and default bindings with identical telemetry") { + val schema = StructType(Seq(StructField("text", StringType))) + val output = StructType(Seq(StructField("answer", StringType))) + val defaults = new AIPrompt("same") + def defaultTemperature(): AIPrompt = new AIPrompt("same") { + setDefault(temperature -> Left(0.5), + telemHeaders -> defaults.getDefault(defaults.telemHeaders).get, + handler -> defaults.getDefault(defaults.handler).get) + } + val implicitTemperature = defaultTemperature() + val explicitTemperature = defaultTemperature().setTemperature(0.5) + def values(stage: AIPrompt): Map[String, Any] = + stage.extractParamMap().toSeq.map(pair => pair.param.name -> pair.value).toMap + assert(implicitTemperature.getClass == explicitTemperature.getClass) + assert(values(implicitTemperature) == values(explicitTemperature)) + assert(AIPromptOutput.fingerprint(implicitTemperature, schema, output, spark) == + AIPromptOutput.fingerprint(defaultTemperature(), schema, output, spark)) + assert(AIPromptOutput.fingerprint(implicitTemperature, schema, output, spark) != + AIPromptOutput.fingerprint(explicitTemperature, schema, output, spark)) + } + + Seq(true).foreach { durable => + test(s"attachment wildcard templates use only public columns with durable=$durable") { + if (durable) delta() + val session = spark + val inline = "data:image/png;base64,AQID" + val data = session.range(3).select(col("id").cast(StringType).as("ticket"), + lit(inline).as("asset"), lit("user-owned").as("asset_filename")) + val expected = data.toJSON.collect().toVector.sorted + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate("{asset};{to_json(struct(*))}").setColumnTypes(Map("asset" -> "path")) + .setOutputCol("answer").setErrorCol("error") + if (durable) { + stage.setOutputPath(root).setOutputBatchSize(1) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 2 && server.requests.size() == 1) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + } + val result = if (durable) stage.resume(session) else 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) + if (durable) assert(stage.resume(session).count() == 3 && server.requests.size() == 3) + } + } + } + + Seq(false, true).foreach { perRow => + test(s"durable top-level metadata preserves inert binding and media names with perRow=$perRow") { + delta() + val metadata = Map("conversation" -> "ticket-42", "prompt" -> "v3", "tools" -> "label", + "image_url" -> "https://example.invalid/inert", "file_data" -> "ordinary text") + val extension = Map("metadata" -> metadata) + val data = spark.range(2).withColumn("text", lit("hello")) + .withColumn("body", lit(AIJson.fromAny(extension).compactPrint)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), path(), server.url + "/v1") + if (perRow) stage.setExtraBodyCol("body") else stage.setExtraBody(extension) + val result = stage.transform(data) + assert(result.collect().forall(_.getAs[Row]("error") == null)) + assert(server.requests.size() == 2) + server.requests.asScala.foreach { case (_, body) => + assert(AIJson.objectString(body).fields("metadata") == AIJson.fromAny(metadata)) + } + assert(stage.resume(spark).count() == 2 && server.requests.size() == 2) + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).allSucceeded) + } + } + } + + test("durable shuffle parallelism supports automatic sizing and caps explicit partition counts") { + Seq("auto", " auto ").foreach(value => assert(AIPromptOutput.batchPartitions(3, value) == 3)) + Seq(199, 200, 100000).foreach { rows => + assert(AIPromptOutput.batchPartitions(rows, "auto") == math.min(rows, 200)) + } + Seq("2", " 2", "2 ", "\t2\n").foreach(value => assert(AIPromptOutput.batchPartitions(3, value) == 2)) + assert(AIPromptOutput.batchPartitions(3, "20") == 3) + Seq("0", "-1", "unknown").foreach { setting => + val error = intercept[IllegalArgumentException](AIPromptOutput.batchPartitions(3, setting)) + assert(error.getMessage.contains("spark.sql.shuffle.partitions")) + } + } + + test("durable batches use configured shuffle parallelism rather than snapshot file packing") { + delta() + val session = spark.newSession() + session.conf.set("spark.sql.shuffle.partitions", " 2 ") + assert(session.conf.get("spark.sql.shuffle.partitions") == " 2 ") + session.conf.set("spark.sql.adaptive.enabled", "true") + val arrived = new CountDownLatch(2) + val overlapped = new AtomicBoolean(true) + Using.resource(new AIRecordingServer((_, _) => { + arrived.countDown() + if (!arrived.await(10, TimeUnit.SECONDS)) overlapped.set(false) + 200 -> chatReply + })) { server => + val data = session.createDataFrame(Seq(Row("first"), Row("last")).asJava, + new StructType().add("text", StringType)).coalesce(1) + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setConcurrency(1) + .setOutputPath(root).setOutputBatchSize(2) + val rows = stage.transform(data).collect() + assert(rows.length == 2 && rows.forall(_.getAs[String]("answer") == "answer")) + assert(overlapped.get(), "Two durable rows should reach the provider concurrently at concurrency=1") + assert(server.requests.size() == 2) + session.conf.set("spark.sql.shuffle.partitions", " 1 ") + assert(stage.resume(session).count() == 2 && server.requests.size() == 2) + assert(AIPrompt.getOutputStatus(session, root).allSucceeded) + } + } + + test("active durable batches reject changed SQL semantics before sending and release quiescent ownership") { + delta() + val session = spark.newSession() + session.conf.set("spark.sql.ansi.enabled", "false") + session.conf.set("spark.sql.legacy.sizeOfNull", "true") + val data = session.range(3).withColumn("tags", expr("cast(null as array)")) + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/chat/completions") + .setPromptTemplate("{coalesce(size(tags), 0)}").setOutputBatchSize(1) + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + val mismatch = intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") session.conf.set("spark.sql.legacy.sizeOfNull", "false")) + } + assert(mismatch.getMessage.contains("Spark SQL semantics")) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && status.ownerToken.isEmpty) + assert(contents == Vector("-1")) + intercept[IllegalArgumentException](stage.resume(session)) + assert(contents == Vector("-1")) + session.conf.set("spark.sql.legacy.sizeOfNull", "true") + session.conf.set("spark.sql.shuffle.partitions", "1") + assert(stage.resume(session).count() == 3 && contents == Vector("-1", "-1", "-1")) + assert(AIPrompt.getOutputStatus(session, root).isComplete) + assert(stage.resume(session).count() == 3 && contents.size == 3) + } + } + + test("durable JSON parsing identity rejects corrupt-record setting changes between batches and on resume") { + delta() + val session = spark.newSession() + val setting = "spark.sql.columnNameOfCorruptRecord" + session.conf.set(setting, "_corrupt_record") + val input = """{"a":"bad"}""" + val data = session.range(3).withColumn("text", lit(input)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setPromptTemplate("{to_json(from_json(text, 'a INT, _corrupt_record STRING'))}").setOutputBatchSize(1) + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") session.conf.set(setting, "other_corrupt_record")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && status.ownerToken.isEmpty) + assert(contents.size == 1 && AIJson.objectString(contents.head).fields("_corrupt_record") == JsString(input)) + val original = contents.head + intercept[IllegalArgumentException](stage.resume(session)) + assert(contents.size == 1) + session.conf.set(setting, "_corrupt_record") + assert(stage.resume(session).count() == 3 && contents == Vector.fill(3)(original)) + assert(stage.resume(session).count() == 3 && contents.size == 3) + } + } + + test("durable JSON partial-result semantics reject changed defaults between batches and on resume") { + delta() + val session = spark.newSession() + val setting = "spark.sql.json.enablePartialResults" + session.conf.unset(setting) + val data = session.range(3).withColumn("text", lit("""{"a":{"ok":1,"bad":"x"}}""")) + val expression = "to_json(from_json(text, 'a STRUCT'))" + val original = data.selectExpr(expression).first().getString(0) + session.conf.set(setting, "false") + val changed = data.selectExpr(expression).first().getString(0) + assert(AIJson.objectString(original) == AIJson.objectString("""{"a":{"ok":1}}""")) + assert(AIJson.objectString(changed) == JsObject.empty) + session.conf.unset(setting) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setPromptTemplate("{" + expression + "}").setOutputBatchSize(1) + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + val mismatch = intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") session.conf.set(setting, "false")) + } + assert(mismatch.getMessage.contains("Spark SQL semantics")) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && status.ownerToken.isEmpty) + assert(contents == Vector(original)) + intercept[IllegalArgumentException](stage.resume(session)) + assert(contents == Vector(original)) + session.conf.unset(setting) + assert(stage.resume(session).count() == 3 && contents == Vector.fill(3)(original)) + assert(AIPrompt.getOutputStatus(session, root).allSucceeded) + assert(stage.resume(session).count() == 3 && contents.size == 3) + } + } + + test("Decisions JSON null message types normalize consistently through direct and facade durable stages") { + delta() + val session = spark + import session.implicits._ + Seq(new OpenAIDecisions(), new AIPrompt().setApiType("decisions")).foreach { stage => + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val variants = Seq("absent" -> None, "null" -> Some(JsNull), "message" -> Some(JsString("message")), + "unknown" -> Some(JsString("other")), "number" -> Some(JsNumber(42))) + val data = variants.map { case (name, tag) => + val fields = Map("role" -> JsString("user"), "content" -> JsString("Evidence")) ++ tag.map("type" -> _) + name -> JsArray(JsObject(fields)).compactPrint + }.toDF("tag", "text") + stage.setUrl(server.url + "/v1").setModel("model").setInputCol("text").setInputMode("json") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setOutputCol("answer").setErrorCol("error") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + val rows = stage.transform(data).collect().map(row => row.getAs[String]("tag") -> row).toMap + Seq("absent", "null", "message").foreach(name => assert(rows(name).getAs[Row]("error") == null)) + Seq("unknown", "number").foreach(name => + assert(rows(name).getAs[Row]("error").getString(0).contains("invalid_request"))) + val expected = if (durable) 6 else 3 + assert(server.requests.size() == expected) + server.requests.asScala.foreach { case (_, body) => + val message = AIJson.objectString(body).fields("input").asInstanceOf[JsArray].elements.head.asJsObject + assert(message.fields.get("type").forall(_ == JsString("message"))) + } + if (durable) { + assert(stage.resume(session).count() == 5 && server.requests.size() == expected) + assert(AIPrompt.getOutputStatus(session, stage.getOutputPath).rowErrorCount == 2) + } + } + } + } + } + + test("durable non-row-local templates are rejected before storage but materialized windows are retained") { + delta() + val session = spark.newSession() + val data = session.range(3).withColumn("text", col("id").cast(StringType)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq("{cast(row_number() over (order by text) as string)}: {text}", + "{cast(max(text) as string)}", "{cast(explode(array(text)) as string)}").foreach { template => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setPromptTemplate(template) + intercept[IllegalArgumentException](stage.transform(data)) + assert(!new File(root).exists() && server.requests.isEmpty) + } + val materialized = data.withColumn("rank", expr("row_number() over (order by text)")) + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), path(), server.url + "/v1") + .setPromptTemplate("{rank}: {text}") + assert(stage.transform(materialized).count() == 3 && server.requests.size() == 3) + assert(stage.resume(session).count() == 3 && server.requests.size() == 3) + } + } + + Seq("systemone", "decisions").foreach { api => + test(s"durable $api recovery requires the original stage class") { + delta() + val questions = if (api == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions + val reply = if (api == "systemone") NativeRequestFixtures.SystemResponse + else NativeRequestFixtures.DecisionResponse + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + def direct(): AIPrompt = if (api == "systemone") new AISystemOne() else new OpenAIDecisions() + def configure(stage: AIPrompt): AIPrompt = stage.setApiType(api).setUrl(server.url).setModel("model") + .setInputCol("text").setQuestions(questions).setOutputCol("answer").setErrorCol("error") + .setOutputPath(root) + val data = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", StringType)) + assert(configure(direct()).transform(data).count() == 1) + val mismatch = intercept[IllegalArgumentException](configure(new AIPrompt()).resume(spark)) + assert(mismatch.getMessage.contains("stage class")) + assert(configure(direct()).resume(spark).count() == 1) + assert(server.requests.size() == 1) + } + } + } + + Seq("responses", "decisions").foreach { api => + test(s"$api message type tags are validated before HTTP in lazy and durable execution") { + delta() + 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) + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + rows.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")) + } + } + assert(server.requests.size() == (if (durable) 4 else 2)) + } + assert(stage.resume(spark).count() == 4 && server.requests.size() == 4) + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).rowErrorCount == 2) + 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)))) + stage.clear(stage.outputPath) + assert(stage.transform(invalidType).head().getAs[Row]("error").getString(0).contains("invalid_request")) + assert(server.requests.size() == 4) + } + } + } + + test("implicit embedding deployment excludes unused model defaults from planning and durable identity") { + 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) + val fingerprint = AIPromptOutput.fingerprint(stage.effectiveCopy(), input.schema, expected, spark) + Seq[Either[String, String]](Right("absent-model"), Left("changed-unused-model")).foreach { value => + GlobalParams.setGlobalParam(AIFoundryModel, value) + assert(stage.transformSchema(input.schema) == expected) + assert(AIPromptOutput.fingerprint(stage.effectiveCopy(), input.schema, expected, spark) == fingerprint) + 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("durable null API versions commit row errors and are not replayed on recovery") { + delta() + val data = spark.createDataFrame( + Seq(Row("first", "selected-version"), Row("missing", Option.empty[String].orNull), + Row("last", "selected-version")).asJava, + new StructType().add("text", StringType).add("version", StringType)).coalesce(1) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url).setModel("model").setInputCol("text") + .setApiVersionCol("version").setOutputCol("answer").setErrorCol("error") + .setOutputPath(root).setOutputBatchSize(2) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("simulated lost acknowledgement")) + } + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == 2 && partial.rowErrorCount == 1 && server.requests.size() == 1) + AIPrompt.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + val rows = stage.resume(spark).collect().map(row => row.getAs[String]("text") -> row).toMap + assert(rows.size == 3 && 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"))) + val complete = AIPrompt.getOutputStatus(spark, root) + assert(complete.isComplete && complete.rowErrorCount == 1 && complete.ownerToken.isEmpty) + } + } + + test("durable embedding recovery ignores a changed unused implicit model default") { + delta() + val previousModel = GlobalParams.getGlobalParam(AIFoundryModel) + val previousEmbedding = GlobalParams.getGlobalParam(OpenAIEmbeddingDeploymentNameKey) + val session = spark + import session.implicits._ + try { + GlobalParams.setGlobalParam(AIFoundryModel, Left("unused-first")) + GlobalParams.setGlobalParam(OpenAIEmbeddingDeploymentNameKey, Left("embedding-model")) + withServer { server => + val root = path() + val stage = new AIPrompt().setApiType("embeddings") + .setUrl(server.url.replace("chat/completions", "embeddings")).setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root).setOutputBatchSize(1) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("first", "last").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("simulated lost acknowledgement")) + } + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == 1 && server.requests.get() == 1) + AIPrompt.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + GlobalParams.setGlobalParam(AIFoundryModel, Left("unused-changed")) + assert(stage.resume(spark).count() == 2 && server.requests.get() == 2) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + GlobalParams.setGlobalParam(AIFoundryModel, Right("absent-model")) + assert(stage.resume(spark).count() == 2 && server.requests.get() == 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("durable Responses accepts null optional resource fields in map and struct content") { + delta() + val session = spark.newSession() + session.conf.set("spark.sql.jsonGenerator.ignoreNullFields", "false") + val empty = Option.empty[String].orNull + val inline = "data:text/plain;base64,aGVsbG8=" + val partSchema = new StructType().add("type", StringType).add("filename", StringType) + .add("file_data", StringType).add("file_id", StringType) + val parts: Seq[(DataType, Any)] = Seq( + MapType(StringType, StringType) -> + Map("type" -> "input_file", "filename" -> "note.txt", "file_data" -> inline, "file_id" -> empty), + partSchema -> Row("input_file", "note.txt", inline, empty)) + val reply = """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + parts.foreach { case (partType, part) => + val messageType = new StructType().add("role", StringType).add("id", StringType) + .add("content", ArrayType(partType)) + val data = session.createDataFrame(Collections.singletonList(Row(Seq(Row("user", empty, Seq(part))))), + new StructType().add("messages", ArrayType(messageType))) + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + assert(stage.transform(data).head().getAs[String]("answer") == "answer") + val root = path() + val row = stage.setOutputPath(root).transform(data).head() + assert(row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null) + assert(AIPrompt.getOutputStatus(session, root).allSucceeded) + assert(stage.resume(session).count() == 1) + } + assert(server.requests.size() == 4) + server.requests.asScala.foreach { case (_, body) => + val message = AIJson.objectString(body).fields("input").asInstanceOf[JsArray].elements.head.asJsObject + assert(!message.fields.contains("id")) + val part = message.fields("content").asInstanceOf[JsArray].elements.head.asJsObject + assert(!part.fields.contains("file_id") && part.fields("file_data") == JsString(inline)) + } + } + } + + test("durable resume ignores auth-named scalar telemetry but keeps ordinary telemetry") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val first = prompt(root, server).setTelemHeaders( + Map("x-trace" -> "same", "Authorization" -> "synthetic-first")) + assert(first.transform(Seq("hello").toDF("text")).count() == 1) + val second = first.copy(ParamMap.empty).setTelemHeaders( + Map("x-trace" -> "same", "AUTHORIZATION" -> "synthetic-rotated", "Api-Key" -> "synthetic-key")) + assert(second.resume(spark).count() == 1) + second.setTelemHeaders(Map("x-trace" -> "changed")) + intercept[IllegalArgumentException](second.resume(spark)) + assert(server.requests.get() == 1) + } + } + + test("durable telemetry column bindings fail before snapshot creation or HTTP") { + delta() + import spark.implicits._ + withServer { server => + Seq("Authorization", "api-key", "x-trace").foreach { header => + val data = Seq("hello").toDF("text") + .withColumn("headers", map(lit(header), lit("synthetic-value"))) + val root = path() + val stage = prompt(root, server) + stage.setVectorParam(stage.telemHeaders, "headers") + intercept[IllegalArgumentException](stage.transform(data)) + assert(!new File(root).exists()) + } + assert(server.requests.get() == 0) + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"$api JSON NOT NULL fields have nullable lazy and durable output schemas") { + delta() + 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)) + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + 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)) + if (durable) { + val resumed = stage.resume(spark) + assert(resumed.schema == result.schema && resumed.head() == row) + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).allSucceeded) + } + } + assert(server.requests.size() == 2) + } + } + } + + test("durable large inline messages remain valid without hiding remote media") { + delta() + val inline = "data:image/png;base64," + ("A" * 20000004) + val parts = Seq(Map("type" -> "input_image", "image_url" -> inline), + Map("type" -> "input_image", "image_url" -> "https://example.invalid/image.png")) + val message = new StructType().add("role", StringType) + .add("content", ArrayType(MapType(StringType, StringType))) + val data = spark.createDataFrame(Collections.singletonList(Row(Seq(Row("user", parts)))), + new StructType().add("messages", ArrayType(message))) + val reply = """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = new AIPrompt().setApiType("responses").setUrl(server.url + "/v1").setModel("model") + .setMessagesCol("messages").setOutputCol("answer").setErrorCol("error") + val inlineData = spark.createDataFrame(Collections.singletonList(Row(Seq(Row("user", parts.take(1))))), + data.schema) + val lazyRow = stage.transform(inlineData).head() + val saved = stage.setOutputPath(path()).transform(inlineData) + assert(saved.schema == stage.transformSchema(inlineData.schema) && saved.head() == lazyRow) + assert(lazyRow.getAs[String]("answer") == "answer" && lazyRow.getAs[Row]("error") == null) + assert(stage.resume(spark).count() == 1 && server.requests.size() == 2) + assert(server.requests.asScala.forall(_._2.contains(inline))) + val failure = intercept[org.apache.spark.SparkException](stage.setOutputPath(root).transform(data)) + assert(failure.getMessage.contains("Durable media must be inline base64")) + assert(server.requests.size() == 2 && !new File(root, "results").exists()) + } + } + + test("durable nested empty structs fail before storage while a zero-column root remains supported") { + delta() + val empty = new StructType() + val values = Seq[(DataType, Any)](empty -> Row.empty, ArrayType(empty) -> Seq(Row.empty), + MapType(StringType, empty) -> Map("key" -> Row.empty), + new StructType().add("nested", empty) -> Row(Row.empty)) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + values.foreach { case (dataType, value) => + val data = spark.createDataFrame(Collections.singletonList(Row(value)), + new StructType().add("extra", dataType)) + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setPromptTemplate("hello") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root) + val failure = intercept[IllegalArgumentException](stage.transform(data)) + assert(failure.getMessage.contains("empty structs")) + assert(!new File(root).exists() && server.requests.isEmpty) + } + val data = spark.createDataFrame(Collections.singletonList(Row.empty), empty) + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setPromptTemplate("hello") + .setOutputCol("answer").setErrorCol("error").setOutputPath(path()) + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(empty) && result.head().getAs[String]("answer") == "answer") + assert(stage.resume(spark).count() == 1 && server.requests.size() == 1) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableRequestSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableRequestSuite.scala new file mode 100644 index 00000000000..26fb3748c55 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableRequestSuite.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 org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import spray.json.DefaultJsonProtocol._ +import spray.json._ + +import java.io.File +import java.nio.charset.StandardCharsets +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptDurableRequestSuite extends AIPromptOutputTestBase { + test("durable prompt template Unicode rejects malformed literals before creating output") { + delta() + import spark.implicits._ + val data = Seq("row").toDF("text") + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + val root = path() + val stage = new AISystemOne().setUrl(server.url).setModel("model") + .setQuestions(NativeRequestFixtures.SystemQuestions).setInputMode("json") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root) + 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(!new File(root).exists() && server.requests.isEmpty) + } + val unicode = "caf\u00e9 \u96ea \uD83D\uDE00 \uFFFD" + val rows = stage.setPromptTemplate(JsString(unicode).compactPrint).transform(data).collect() + assert(rows.length == 1 && rows.head.getAs[Row]("error") == null) + assert(AIJson.objectString(server.requests.peek()._2).fields("state") == JsString(unicode)) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + assert(stage.resume(spark).collect().toSet == rows.toSet) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == rows.toSet && server.requests.size() == 1) + } + } + + test("durable fingerprints distinguish malformed UTF-16 across interrupted recovery") { + delta() + import spark.implicits._ + val data = Seq("one", "two").toDF("text") + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setSystemPrompt("\uD800").setOutputCol("answer").setErrorCol("error") + .setOutputPath(root).setOutputBatchSize(1) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after invalid request")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 1 && status.pendingCount == 1 && status.rowErrorCount == 1) + assert(server.requests.isEmpty) + val committed = AIPrompt.loadOutput(spark, root).head() + assert(committed.getAs[Row]("error").getString(0).contains("invalid_request")) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + Seq("?", "\uDC00", "\\" + "ud800", "\uFFFD", "\uD83D\uDE00").foreach { changed => + val error = intercept[IllegalArgumentException] { + stage.copy(ParamMap.empty).setSystemPrompt(changed).resume(spark) + } + assert(error.getMessage.contains("configuration") && server.requests.isEmpty) + val unchanged = AIPrompt.getOutputStatus(spark, root) + assert(unchanged.ownerToken.isEmpty && unchanged.committedCount == 1 && unchanged.pendingCount == 1) + } + val rows = stage.resume(spark).collect() + assert(rows.length == 2 && rows.contains(committed)) + assert(rows.forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + val complete = AIPrompt.getOutputStatus(spark, root) + assert(complete.isComplete && complete.rowErrorCount == 2 && !complete.allSucceeded) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == rows.toSet && server.requests.isEmpty) + } + } + + test("durable structured JSON surrogate errors remain terminal across resume and reload") { + delta() + import spark.implicits._ + val escape = "\\" + "u" + Seq("chat_completions", "responses").foreach { api => + Using.resource(new AIRecordingServer((_, body) => { + val document = if (body.contains("invalid-value")) s"""{"value":"${escape}d800"}""" + else if (body.contains("invalid-key")) s"""{"${escape}dc00":"private-value","?":"other"}""" + else s"""{"value":"${escape}d83d${escape}de00"}""" + 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 root = path() + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setPostProcessingOptions(Map("jsonSchema" -> "MAP")).setOutputPath(root) + val rows = stage.transform(Seq("invalid-value", "invalid-key", "valid").toDF("text")).collect() + rows.filter(_.getAs[String]("text") != "valid").foreach { row => + assert(row.getAs[Row]("error") != null) + assert(row.getAs[Row]("error").getString(0).contains("parse_error")) + assert(row.getAs[String]("raw").contains("\\\\u")) + } + val valid = rows.find(_.getAs[String]("text") == "valid").get + assert(valid.getAs[Row]("error") == null && + valid.getAs[Map[String, String]]("answer")("value") == "\uD83D\uDE00") + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.isComplete && !status.allSucceeded && status.rowErrorCount == 2) + assert(stage.resume(spark).collect().toSet == rows.toSet) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == rows.toSet && server.requests.size() == 3) + } + } + } + + test("durable JSON surrogate validation stores parse errors rather than corrupted successful answers") { + delta() + import spark.implicits._ + Using.resource(new AIRecordingServer((_, body) => { + val escape = "\\" + "u" + val content = "\"" + escape + (if (body.contains("invalid")) "d800" else s"d83d${escape}de00") + "\"" + 200 -> s"""{"choices":[{"index":0,"message":{"role":"assistant","content":$content}}]}""" + })) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw").setOutputPath(root) + val result = stage.transform(Seq("invalid", "valid").toDF("text")).collect() + val failed = result.find(_.getAs[String]("text") == "invalid").get + assert(failed.getAs[String]("answer") == null) + assert(failed.getAs[Row]("error").getString(0).contains("parse_error")) + assert(failed.getAs[String]("raw").contains("\\ud800")) + val valid = result.find(_.getAs[String]("text") == "valid").get + assert(valid.getAs[Row]("error") == null && valid.getAs[String]("answer") == "\uD83D\uDE00") + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.isComplete && !status.allSucceeded && status.rowErrorCount == 1) + assert(stage.resume(spark).collect().toSet == result.toSet) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == result.toSet && server.requests.size() == 2) + } + } + + test("durable native UTF-8 failures remain terminal across interrupted recovery") { + delta() + import spark.implicits._ + val unicode = "caf\u00e9 \u96ea \uFFFD" + val valid = s"""{"model":"model","answers":{"q":${JsString(unicode).compactPrint}}}""" + Using.resource(new AIRecordingServer((count, _) => + 200 -> (if (count == 1) """{"model":"model","answers":{"q":"private-value~"}}""" else valid), + encodeResponse = text => text.getBytes(StandardCharsets.UTF_8) + .map(byte => if (byte == '~'.toByte) 0xff.toByte else byte))) { server => + val root = path() + val stage = new AISystemOne().setUrl(server.url).setModel("model").setStateCol("text") + .setQuestions("""{"q":{"type":"noul","instructions":"Return an answer."}}""") + .setSubscriptionKey("synthetic-key").setOutputCol("answer").setErrorCol("error") + .setRawResponseCol("raw").setUsageCol("usage").setResponseIdCol("id") + .setOutputPath(root).setOutputBatchSize(1) + val data = Seq("one", "two", "three").toDF("text") + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed UTF-8 failure")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && status.rowErrorCount == 1) + assert(!status.allSucceeded && server.requests.size() == 1) + val failed = AIPrompt.loadOutput(spark, root).head() + assert(Seq("answer", "raw", "usage", "id").forall(field => failed.isNullAt(failed.fieldIndex(field)))) + assert(failed.getAs[Row]("error").getString(0).parseJson.asJsObject.fields("code") == JsString("parse_error")) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val resumed = stage.resume(spark) + assert(resumed.schema == stage.transformSchema(data.schema)) + val rows = resumed.collect() + assert(rows.length == 3 && rows.contains(failed)) + rows.filter(_.getAs[Row]("error") == null).foreach { row => + assert(row.getAs[String]("raw") == valid) + assert(row.getAs[Row]("answer").getMap[String, String](1)("q") == JsString(unicode).compactPrint) + } + val complete = AIPrompt.getOutputStatus(spark, root) + assert(complete.isComplete && !complete.allSucceeded && complete.rowErrorCount == 1) + assert(complete.ownerToken.isEmpty && complete.pendingCount == 0) + assert(stage.resume(spark).collect().toSet == rows.toSet) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == rows.toSet && server.requests.size() == 3) + } + } + + test("durable Decisions preserves exact score bounds in committed and resumed results") { + delta() + val session = spark + import session.implicits._ + val scores = Seq("0", "2", "-1e-400", "2.0000000000000000001") + 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 => + val root = path() + val stage = new AIPrompt().setApiType("decisions").setUrl(server.url + "/v1").setModel("model") + .setInputCol("input").setQuestions(NativeRequestFixtures.DecisionQuestions) + .setOutputCol("answer").setErrorCol("error").setRawResponseCol("raw") + .setOutputPath(root).setOutputBatchSize(4) + val result = stage.transform(scores.toDF("input")) + val rows = result.collect() + assert(rows.length == 4 && server.requests.size() == 4) + rows.foreach { row => + val score = row.getAs[String]("input") + assert(row.getAs[String]("raw") == response(score)) + 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")) + } + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.rowErrorCount == 2 && !status.allSucceeded) + assert(stage.resume(spark).collect().toSet == rows.toSet) + assert(AIPrompt.loadOutput(spark, root).collect().toSet == rows.toSet) + assert(server.requests.size() == 4) + } + } + + test("durable Chat empty refusal commits successful answers across interrupted recovery") { + delta() + import spark.implicits._ + val data = Seq("one", "two", "three").toDF("text") + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer","refusal":""}, + |"finish_reason":"stop"}]}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/chat/completions") + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed answer")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 2 && status.pendingCount == 1 && status.rowErrorCount == 0) + assert(server.requests.size() == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val result = stage.resume(spark) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + assert(rows.length == 3 && rows.forall(row => + row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer" && + row.getAs[String]("raw") == reply)) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + assert(stage.resume(spark).count() == 3 && server.requests.size() == 3) + } + } + + test("durable templates require materialized encryption while lazy random IV encryption remains supported") { + delta() + import spark.implicits._ + val data = Seq.fill(2)("hello").toDF("text").withColumn("iv", lit("abcdefghijkl")) + val encrypted = "base64(aes_encrypt(text, '0000111122223333'))" + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + Seq(encrypted, + "base64(aes_encrypt(text, '0000111122223333', 'CBC', 'DEFAULT', ''))", + "base64(aes_encrypt(text, '0000111122223333', 'GCM', 'DEFAULT', iv))", + "base64(aes_encrypt(text, '0000111122223333', 'GCM', 'DEFAULT', 'abcdefghijkl'))", + "base64(aes_encrypt(text, '0000111122223333', 'ECB'))").foreach { expression => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate("{" + expression + "}") + val schemaError = intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + val transformError = intercept[IllegalArgumentException](stage.transform(data)) + Seq(schemaError, transformError).foreach { error => + assert(error.getMessage.contains("stable built-in expressions") && error.getCause == null) + } + assert(!new File(root).exists() && server.requests.isEmpty) + } + val lazyStage = AIPromptStorageTestConfig.configure(new AIPrompt(), path(), + server.url + "/v1/chat/completions").setPromptTemplate("{" + encrypted + "}") + lazyStage.clear(lazyStage.outputPath) + val result = lazyStage.transform(data) + assert(result.schema == lazyStage.transformSchema(data.schema)) + assert(result.collect().forall(_.getAs[Row]("error") == null)) + val ciphertexts = server.requests.asScala.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + }.toVector + assert(ciphertexts.size == 2 && ciphertexts.distinct.size == 2) + } + } + + test("durable recovery retains materialized random IV ciphertext without resending committed rows") { + delta() + import spark.implicits._ + val data = Seq.fill(3)("hello").toDF("text") + .withColumn("ciphertext", expr("base64(aes_encrypt(text, '0000111122223333'))")) + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate("{ciphertext}") + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed ciphertext")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 2 && status.pendingCount == 1 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val result = stage.resume(spark) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + assert(rows.length == 3 && rows.forall(_.getAs[Row]("error") == null)) + val ciphertexts = server.requests.asScala.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + }.toVector + assert(ciphertexts.size == 3 && ciphertexts.distinct.size == 3) + assert(ciphertexts.sorted == rows.map(_.getAs[String]("ciphertext")).sorted.toVector) + assert(stage.resume(spark).count() == 3 && server.requests.size() == 3) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + } + } + + Seq(false, true).foreach { perRow => + test(s"durable Responses output-only include survives interrupted recovery with perRow=$perRow") { + delta() + import spark.implicits._ + val requested = Seq("message.output_text.logprobs", "reasoning.encrypted_content") + val extension = JsObject("include" -> JsArray(requested.map(JsString(_)).toVector)) + val data = Seq("one", "two", "three").toDF("text").withColumn("body", lit(extension.compactPrint)) + val reply = + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer", + |"logprobs":[{"token":"answer","logprob":-0.1}]}]}, + |{"type":"reasoning","encrypted_content":"synthetic-encrypted-reasoning"}], + |"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/responses") + .setApiType("responses") + if (perRow) stage.setExtraBodyCol("body") else stage.setExtraBody(Map("include" -> requested)) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed include output")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 2 && status.pendingCount == 1 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val result = stage.resume(spark) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect() + assert(rows.length == 3 && rows.forall(row => + row.getAs[Row]("error") == null && row.getAs[String]("answer") == "answer" && + row.getAs[String]("raw") == reply)) + assert(server.requests.size() == 3) + server.requests.asScala.foreach { case (url, body) => + assert(url == "/v1/responses") + assert(AIJson.objectString(body).fields("include") == extension.fields("include")) + } + assert(stage.resume(spark).count() == 3 && server.requests.size() == 3) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableSchemaSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableSchemaSuite.scala new file mode 100644 index 00000000000..fbe7761341d --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptDurableSchemaSuite.scala @@ -0,0 +1,660 @@ +// 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.HandlingUtils +import com.microsoft.azure.synapse.ml.services.openai.OpenAIDecisions +import org.apache.spark.ml.Pipeline +import org.apache.spark.ml.linalg.{SQLDataTypes, Vectors} +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.sql.{Encoders, Row, SparkSession} +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.{ArrayType, CalendarIntervalType, IntegerType, + MapType, NullType, StringType, StructType} +import spray.json.{JsArray, JsNull, JsObject, JsString} + +import java.io.File +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptDurableSchemaSuite extends AIPromptOutputTestBase { + private val inputSchema = new StructType().add("text", StringType) + .add("credential", StringType).add("headers", MapType(StringType, StringType)) + + Seq(false, true).foreach { execution => + test(s"durable preflight rejects custom default handlers with execution=$execution") { + if (execution) delta() + withServer { server => + val root = path() + val stage = new TestableAIAlternateDefaultHandlerPrompt("alternate") + .setUrl(server.url).setModel("model").setInputCol("text").setOutputPath(root) + val data = spark.range(1).select(lit("hello").as("text")) + assert(!stage.isSet(stage.handler)) + val failure = intercept[IllegalArgumentException] { + if (execution) stage.transform(data) else stage.transformSchema(data.schema) + } + assert(failure.getMessage.contains("default HTTP handler")) + intercept[IllegalArgumentException] { + AIPromptOutput.fingerprint(stage, data.schema, data.schema, spark) + } + assert(!new File(root).exists() && server.requests.get() == 0) + } + } + } + + Seq(true).foreach { durable => + test(s"non-row-local templates are rejected before inference or storage with durable=$durable") { + 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 root = path() + val stage = prompt(root, server).setPromptTemplate(template) + if (!durable) stage.clear(stage.outputPath) + val schemaError = intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + val transformError = intercept[IllegalArgumentException](stage.transform(data)) + Seq(schemaError, transformError).foreach { error => + assert(error.getMessage.contains("row-local") && Option(error.getCause).isEmpty) + } + assert(!new File(root).exists() && server.requests.get() == 0) + } + } + } + } + + 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")) + ).foreach { case (expression, columns) => + Seq("complete", "resume").foreach { mode => + test(s"wildcard template $expression validates only public columns in $mode mode") { + if (mode != "lazy") delta() + val data = spark.range(3).selectExpr(columns: _*) + val expected = data.selectExpr(expression).collect().map(_.getString(0)).toVector.sorted + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setPromptTemplate(s"{$expression}").setOutputCol("answer").setErrorCol("error") + if (mode != "lazy") stage.setOutputPath(root).setOutputBatchSize(1) + val expectedSchema = stage.transformSchema(data.schema) + if (mode == "resume") { + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && server.requests.size() == 1) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + } + val output = if (mode == "resume") stage.resume(spark) else stage.transform(data) + val rows = output.collect() + assert(output.schema == expectedSchema && 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) + val contents = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + }.sorted + assert(contents == expected && server.requests.size() == 3) + if (mode != "lazy") { + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + assert(stage.resume(spark).collect().toSet == rows.toSet && server.requests.size() == 3) + } + } + } + } + } + + test("lazy and durable templates preserve precomputed window columns and null input rows") { + delta() + 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 => + Seq(true).foreach { durable => + val stage = prompt(path(), server) + .setPromptTemplate("{coalesce(cast(ts as string), 'missing')}:{cast(ordinal as string)}") + if (!durable) stage.clear(stage.outputPath) + 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() == 4) + } + } + + test("fused JVM message arrays preserve lazy and durable outcomes") { + delta() + 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 + 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").setMessagesCol("messages") + .setOutputCol("answer").setErrorCol("error") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + 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(server.requests.size() == 1) + val body = AIJson.objectString(server.requests.remove()._2) + assert(body.fields("messages").asInstanceOf[JsArray].elements.head.asJsObject + .fields("content") == JsString("hello")) + if (durable) { + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && status.rowErrorCount == 1) + assert(stage.resume(spark).count() == 2 && server.requests.isEmpty) + } + } + } + } + + test("fused JVM storage arrays preserve vectors nulls and nested values") { + val schema = new StructType().add("vectors", ArrayType(SQLDataTypes.VectorType)) + .add("nested", MapType(StringType, ArrayType(StringType))).add("numbers", ArrayType(IntegerType)) + val input = spark.range(1).map { _ => + Row(Array(Vectors.sparse(3, Array(1), Array(0.5)), Option.empty[org.apache.spark.ml.linalg.Vector].orNull), + Map("key" -> Array("value", Option.empty[String].orNull)), Array(1, 2)) + }(Encoders.row(schema)) + val expected = input.repartition(1).collect().toSeq + val output = AIPromptStorageCodec.decode(AIPromptStorageCodec.encode(input, schema), schema) + assert(output.schema == schema && output.collect().toSeq == expected) + } + + Seq("chat_completions", "responses").foreach { api => + test(s"$api failed projections retain partial observations and terminal durable provider errors") { + delta() + 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") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + 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() == (if (durable) 8 else 4)) + if (durable) { + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && !status.allSucceeded && status.rowErrorCount == 4) + assert(stage.resume(spark).collect().toSet == rows.toSet && server.requests.size() == 8) + } + } + } + } + } + + test("nullable embedding object tags preserve lazy and durable vectors without replay") { + delta() + 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") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + 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() == (if (durable) 2 else 1)) + if (durable) { + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).allSucceeded) + assert(stage.resume(spark).head() == row && server.requests.size() == 2) + } + } + } + } + + test("Responses in-progress messages retain partial observations and terminal durable errors") { + delta() + 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") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()) + 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() == (if (durable) 4 else 2)) + if (durable) { + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && !status.allSucceeded && status.rowErrorCount == 2) + assert(stage.resume(spark).collect().toSet == rows.toSet && server.requests.size() == 4) + } + } + } + } + + test("durable Responses preserves literal file_data and never replays committed provider errors") { + delta() + 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 + } + val reply = """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + Using.resource(new AIRecordingServer((_, body) => { + if (Set(raw, inline)(fileData(body))) 200 -> reply + 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 lazyRows = stage.transform(input).collect().map(row => row.getAs[String]("text") -> row).toMap + assert(lazyRows.size == 4 && lazyRows.values.count(_.getAs[Row]("error") != null) == 2) + assert(server.requests.asScala.map(entry => fileData(entry._2)).toSeq.sorted == contents.sorted) + server.requests.clear() + val root = path() + stage.setOutputPath(root).setOutputBatchSize(2) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, input, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == 2 && partial.pendingCount == 2 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + val output = stage.resume(spark) + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect().map(row => row.getAs[String]("text") -> row).toMap + assert(rows.keySet == lazyRows.keySet) + rows.foreach { case (text, row) => + assert(row.getAs[String]("answer") == lazyRows(text).getAs[String]("answer")) + assert(row.getAs[Row]("error") == lazyRows(text).getAs[Row]("error")) + } + val complete = AIPrompt.getOutputStatus(spark, root) + assert(complete.isComplete && complete.rowErrorCount == 2 && complete.ownerToken.isEmpty) + assert(stage.resume(spark).count() == 4 && server.requests.size() == 4) + assert(server.requests.asScala.map(entry => fileData(entry._2)).toSeq.sorted == contents.sorted) + } + } + + Seq(false, true).foreach { perRow => + test(s"durable body extensions preserve literal file_data with perRow=$perRow") { + delta() + import spark.implicits._ + val extension = Map("extension" -> Map("file_data" -> "JVBERi0xLjQK")) + val input = Seq("one", "two").toDF("text").withColumn("body", lit(AIJson.fromAny(extension).compactPrint)) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + if (perRow) stage.setExtraBodyCol("body") else stage.setExtraBody(extension) + val output = stage.transform(input) + assert(output.collect().forall(_.getAs[Row]("error") == null)) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + assert(stage.resume(spark).count() == 2 && server.requests.size() == 2) + server.requests.asScala.foreach { case (_, body) => + assert(AIJson.objectString(body).fields("extension") == AIJson.fromAny(extension("extension"))) + } + } + } + } + + test("durable container JSON conversion failures remain terminal row errors on resume") { + delta() + import spark.implicits._ + Seq("chat_completions", "responses").foreach { api => + Using.resource(new AIRecordingServer((_, body) => { + val text = if (body.contains("incompatible")) """[{"value":"wrong"}]""" + else if (body.contains("bad-singleton")) """{"value":"wrong"}""" + else if (body.contains("singleton")) """{"value":1}""" + else """[{"value":null}]""" + 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 + } + 200 -> reply + })) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setApiType(api).setPostProcessingOptions(Map("jsonSchema" -> "ARRAY>")) + val data = Seq("incompatible", "valid", "singleton", "bad-singleton").toDF("text") + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val rows = output.collect().map(row => row.getAs[String]("text") -> row).toMap + Seq("incompatible", "bad-singleton").foreach { name => + assert(rows(name).getAs[Row]("error").getString(0).contains("parse_error")) + } + Seq("valid", "singleton").foreach(name => assert(rows(name).getAs[Row]("error") == null)) + assert(rows("valid").getAs[Seq[Row]]("answer").head.isNullAt(0)) + assert(rows("singleton").getAs[Seq[Row]]("answer") == Seq(Row(1))) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.rowErrorCount == 2 && !status.allSucceeded && status.pendingCount == 0) + val resumed = stage.resume(spark).collect().map(row => row.getAs[String]("text") -> row).toMap + assert(resumed == rows && server.requests.size() == 4) + } + } + } + + test("Decisions null image options normalize through lazy and durable direct and facade stages") { + delta() + 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("data:image/png;base64,AQID"), + "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") + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(path()).setOutputBatchSize(3) + val output = stage.transform(data) + assert(output.schema == stage.transformSchema(data.schema)) + val rows = output.collect() + assert(rows.length == 3 && rows.forall(_.getAs[Row]("error") == null)) + assert(rows.forall(_.getAs[Row]("answer") != null)) + val sends = if (durable) 6 else 3 + assert(server.requests.size() == sends) + 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) + } + if (durable) { + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.allSucceeded && status.ownerToken.isEmpty) + assert(stage.resume(spark).count() == 3 && server.requests.size() == sends) + } + } + } + } + } + + test("durable recovery rotates custom authentication hooks but fingerprints ordinary header values") { + delta() + val data = spark.createDataFrame(Seq(Row("first"), Row("last")).asJava, + new StructType().add("text", StringType)) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Seq("x-model-key", "x-model-token").foreach { authName => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val first = new TestableAIHeaderPrompt("durable-custom-auth").setUrl(server.url + "/v1") + .setModel("model").setInputCol("text").setOutputCol("answer").setErrorCol("error") + .setOutputPath(root).setOutputBatchSize(1) + .setCustomHeaders(Map(authName -> "synthetic-first", "api-key" -> "ordinary")) + .setTelemHeaders(Map(authName -> "ignored-first", "Authorization" -> "ordinary")) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(first, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 1 && status.pendingCount == 1 && server.requests.size() == 1) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val second = first.copy(ParamMap.empty) + .setCustomHeaders(Map(authName.toUpperCase(java.util.Locale.ROOT) -> "synthetic-rotated", + "api-key" -> "ordinary")) + .setTelemHeaders(Map(authName -> "ignored-rotated", "Authorization" -> "ordinary")) + val result = second.resume(spark) + assert(result.collect().forall(_.getAs[Row]("error") == null)) + assert(result.count() == 2 && server.requests.size() == 2) + assert(server.headers.asScala.last(authName) == "synthetic-rotated") + assert(second.resume(spark).count() == 2 && server.requests.size() == 2) + second.setCustomHeaders(Map(authName -> "synthetic-rotated", "api-key" -> "changed")) + intercept[IllegalArgumentException](second.resume(spark)) + second.setCustomHeaders(Map(authName -> "synthetic-rotated", "api-key" -> "ordinary")) + .setTelemHeaders(Map("Authorization" -> "changed")) + intercept[IllegalArgumentException](second.resume(spark)) + assert(server.requests.size() == 2 && AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + } + } + } + + test("durable scalar header preflight rejects invalid configuration before storage and allows correction") { + delta() + val data = spark.createDataFrame(Seq(Row("hello")).asJava, new StructType().add("text", StringType)) + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Seq("ap\u0131-key" -> "api-key", "author\u0131zation" -> "Authorization").foreach { case (alias, canonical) => + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setOutputPath(path()) + .setCustomHeaders(Map(alias -> "synthetic-first")) + val invalid: Seq[AIPrompt => AIPrompt] = Seq( + _.setCustomHeaders(Map(alias -> "synthetic-private")), + _.setTelemHeaders(Map(alias -> "synthetic-private")), + _.setCustomHeaders(Map("X-Caller" -> "synthetic-private\n")), + _.setTelemHeaders(Map("X-Caller" -> "synthetic-private\n")), + _.setSubscriptionKey("synthetic-private\n"), _.setAADToken("synthetic-private\n"), + _.setCustomAuthHeader("synthetic-private\n")) + stage.clear(stage.customHeaders) + invalid.foreach { configure => + val candidate = configure(stage.copy(ParamMap.empty)) + Seq( + intercept[IllegalArgumentException](candidate.transformSchema(data.schema)), + intercept[IllegalArgumentException](candidate.transform(data)) + ).foreach { error => + assert(error.getMessage.contains("Invalid HTTP header") && !error.getMessage.contains("synthetic-private")) + } + assert(!new File(stage.getOutputPath).exists() && server.requests.isEmpty) + } + val result = stage.setCustomHeaders(Map(canonical -> "synthetic-corrected")).transform(data).head() + assert(result.getAs[Row]("error") == null && server.requests.size() == 1) + assert(stage.resume(spark).head() == result && server.requests.size() == 1) + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).allSucceeded) + } + } + } + + test("durable schema rejects nonresumable settings without jobs, storage or requests") { + 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 + spark.sparkContext.setJobGroup(group, "Durable schema-only configuration validation") + try { + settings.zipWithIndex.foreach { case (configure, index) => + withClue(s"configuration $index: ") { + val root = path() + val stage = configure(prompt(root, server).setApiType("responses") + .setUrl(server.url.stripSuffix("/chat/completions"))) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + lazyStage.transformSchema(inputSchema) + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + intercept[IllegalArgumentException] { + new Pipeline().setStages(Array(stage)).transformSchema(inputSchema) + } + assert(!new File(root).exists()) + } + } + assert(spark.sparkContext.statusTracker.getJobIdsForGroup(group).isEmpty) + assert(server.requests.get() == 0) + } finally spark.sparkContext.clearJobGroup() + } + } + + test("durable schema rejects external native questions through direct and facade stages") { + withServer { server => + val cases = Seq( + (new AIPrompt().setApiType("systemone"), new AISystemOne(), + """{"urgent":{"type":"noul","instructions":"Is it urgent?","file_id":"provider-file"}}"""), + (new AIPrompt().setApiType("decisions"), new OpenAIDecisions(), + """[{"type":"predicate","instructions":"Is it urgent?","file_id":"provider-file"}]""")) + cases.foreach { case (facade, direct, questions) => + Seq(facade, direct).foreach { stage => + val root = path() + stage.setInputCol("text").setModel("model").setQuestions(questions) + .setUrl(if (stage.getApiType == "systemone") server.url.stripSuffix("/v1/chat/completions") + else server.url.stripSuffix("/chat/completions")) + stage.transformSchema(inputSchema) + stage.setOutputPath(root) + val error = intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + assert(error.getMessage.contains("external resources") && !new File(root).exists()) + } + } + assert(server.requests.get() == 0) + } + } + + test("durable schema analyses stable templates without evaluating input or requiring Delta") { + val session = spark.newSession() + session.conf.set("spark.sql.ansi.enabled", "true") + session.udf.register("aiprompt_schema_udf", (value: String) => value) + AIPrompt.withSession(session) { + withServer { server => + val group = java.util.UUID.randomUUID().toString + session.sparkContext.setJobGroup(group, "Durable schema-only template analysis") + try { + Seq("{current_timestamp()}", "{aiprompt_schema_udf(text)}", "{(select 'external')}") + .foreach { template => + val root = path() + val stage = prompt(root, server).setPromptTemplate(template) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + lazyStage.transformSchema(inputSchema) + intercept[IllegalArgumentException](stage.transformSchema(inputSchema)) + assert(!new File(root).exists()) + } + val root = path() + val valid = prompt(root, server).setPromptTemplate("{upper(text)}:{1 / 0}") + val lazyStage = valid.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + assert(valid.transformSchema(inputSchema) == AIPromptOutput.storageSchema( + lazyStage.transformSchema(inputSchema))) + assert(!new File(root).exists() && server.requests.get() == 0) + assert(session.sparkContext.statusTracker.getJobIdsForGroup(group).isEmpty) + } finally session.sparkContext.clearJobGroup() + } + } + } + + test("durable schema validates input types even when an output role replaces the field") { + withServer { server => + val roles: Seq[AIPrompt => AIPrompt] = Seq( + _.setOutputCol("replaced"), _.setErrorCol("replaced"), _.setUsageCol("replaced"), + _.setRawResponseCol("replaced"), _.setResponseIdCol("replaced")) + val group = java.util.UUID.randomUUID().toString + spark.sparkContext.setJobGroup(group, "Durable input type preflight") + try { + Seq(CalendarIntervalType, NullType, new StructType().add("empty", new StructType())) + .foreach { dataType => + roles.zipWithIndex.foreach { case (configure, index) => + withClue(s"${dataType.typeName}, role $index: ") { + val root = path() + val schema = inputSchema.add("replaced", dataType) + val stage = configure(prompt(root, server)) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + AIPromptOutput.storageSchema(lazyStage.transformSchema(schema)) + Seq(stage, new Pipeline().setStages(Array(stage))).foreach { transformer => + val error = intercept[IllegalArgumentException](transformer.transformSchema(schema)) + assert(error.getMessage.contains("AIPrompt durable storage")) + } + assert(!new File(root).exists()) + } + } + } + assert(spark.sparkContext.statusTracker.getJobIdsForGroup(group).isEmpty) + assert(server.requests.get() == 0) + } finally spark.sparkContext.clearJobGroup() + } + } +} + +class AIPromptSessionTeardownSuite extends AIPromptOutputTestBase { + test("durable fixture teardown clears stopped session aliases before other stages are constructed") { + val alias = spark.newSession() + SparkSession.setActiveSession(alias) + SparkSession.setDefaultSession(alias) + AIPromptDeltaTestSession.stopSparkSession() + assert(alias.sparkContext.isStopped) + assert(!SparkSession.getActiveSession.exists(_.sparkContext.isStopped)) + assert(!SparkSession.getDefaultSession.exists(_.sparkContext.isStopped)) + new AIPrompt().setInputCol("text").setUrl("https://example.invalid/v1").setModel("model") + .transformSchema(new StructType().add("text", StringType)) + } +} 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/AIPromptFabricRecovery.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricRecovery.scala new file mode 100644 index 00000000000..ca6c6a30fb5 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptFabricRecovery.scala @@ -0,0 +1,51 @@ +// 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.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.types.StructType + +import java.io.File +import java.util.{Collections, UUID} + +/** Test-jar bridge for managed-runtime storage guards and interruption after a real storage commit. */ +object AIPromptFabricRecovery { + def verifyLocalStorageRejected(spark: SparkSession): Unit = { + require(!spark.sparkContext.isLocal, "This regression requires a non-local Spark master") + val directory = new File(System.getProperty("java.io.tmpdir"), "aiprompt-reject-local-" + UUID.randomUUID()) + val path = directory.toURI.toString + require(!directory.exists(), "The local rejection probe requires a fresh path") + def rejected(action: => Any): Unit = { + try { + action + throw new IllegalStateException("Non-local Spark accepted file: durable storage") + } catch { + case error: IllegalArgumentException => + require(error.getMessage.contains("Spark local mode"), "Expected the local storage runtime guard") + } + require(!directory.exists(), "The runtime guard must reject before creating storage") + } + rejected(AIPromptDelta.requireRuntime(spark, path)) + val input = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", "string")) + val stage = new AIPrompt().setUrl("http://127.0.0.1:1/v1").setModel("synthetic-model") + .setInputCol("text").setOutputPath(path) + rejected(stage.transform(input)) + rejected(stage.resume(spark)) + rejected(AIPrompt.loadOutput(spark, path)) + rejected(AIPrompt.getOutputStatus(spark, path)) + rejected(AIPrompt.releaseOutputOwner(spark, path, "synthetic-owner", confirmStopped = true)) + } + + def stopAfterFirstBatch(prompt: AIPrompt, input: DataFrame): Unit = { + val stopped = new IllegalStateException("AIPrompt managed-runtime post-commit interruption") + try { + AIPromptOutput.transformObserved(prompt, input, event => + if (event == "batch_committed") throw stopped) + throw new IllegalStateException("Expected a nonempty input and a committed batch") + } catch { + case error: IllegalStateException if error eq stopped => () + } + } +} 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/AIPromptInitializationSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptInitializationSuite.scala new file mode 100644 index 00000000000..7853b7780b4 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptInitializationSuite.scala @@ -0,0 +1,745 @@ +// 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.{StringType, StructType} +import spray.json._ + +import java.io.File +import scala.collection.JavaConverters._ +import scala.util.{Try, Using} + +class AIPromptInitializationSuite extends AIPromptOutputTestBase { + 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 + + test("missing run paths report no run without preventing a fresh transform") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + val calls = Seq[() => Any]( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => AIPrompt.releaseOutputOwner(spark, root, "stopped-owner", confirmStopped = true)) + calls.foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("No AIPrompt run exists")) + assert(!error.getMessage.contains("incomplete") && !error.getMessage.contains(root)) + assert(!new File(root).exists() && server.requests.get() == 0) + } + assert(stage.transform(Seq("hello").toDF("text")).count() == 1) + assert(AIPrompt.getOutputStatus(spark, root).isComplete && server.requests.get() == 1) + } + } + + test("unrelated output directories are not described as incomplete runs or modified") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val input = Seq("unrelated data").toDF("text") + input.coalesce(1).write.parquet(root) + val files = new File(root).list().toSet + val stage = prompt(root, server) + val calls = Seq[() => Any]( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => AIPrompt.releaseOutputOwner(spark, root, "stopped-owner", confirmStopped = true), + () => stage.transform(input)) + calls.foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("not an AIPrompt run directory")) + assert(!error.getMessage.contains("incomplete") && !error.getMessage.contains("cleaning up")) + assert(!error.getMessage.contains(root)) + assert(new File(root).list().toSet == files && server.requests.get() == 0) + } + assert(spark.read.parquet(root).as[String].collect().toSeq == Seq("unrelated data")) + } + } + + test("uncommitted control data reports incomplete initialization without changing storage") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val control = new File(root, "_control") + val input = Seq("hello").toDF("text") + input.coalesce(1).write.parquet(control.getPath) + val files = control.list().toSet + val stage = prompt(root, server) + val calls = Seq[() => Any]( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => AIPrompt.releaseOutputOwner(spark, root, "stopped-owner", confirmStopped = true), + () => stage.transform(input)) + calls.foreach { call => + val result = Try(call()) + assert(result.isFailure) + val error = result.failed.get + assert(error.isInstanceOf[IllegalArgumentException], error.toString) + assert(error.getMessage.contains("initialization is incomplete") && error.getMessage.contains("new outputPath")) + assert(error.getMessage.contains("may still be running")) + assert(error.getMessage.contains("Wait for initialization to finish")) + assert(error.getMessage.contains("confirm the creating application has stopped")) + assert(!error.getMessage.contains(root)) + assert(control.list().toSet == files && server.requests.get() == 0) + } + } + } + + test("active initialization requires waiting or confirmed stopped execution before abandoning the run") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(spark, root) + val stage = prompt(root, server) + val input = Seq("hello").toDF("text") + var observed = Set.empty[String] + val output = AIPromptOutput.transformObserved(stage, input, event => + if (Set("input_committed", "input_validated").contains(event)) { + val before = AIPromptDelta.control(spark, qualified) + assert(before.state == "initializing" && before.owner.isDefined) + val calls = Seq[() => Any]( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => stage.copy(ParamMap.empty).transform(input)) + calls.foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("may still be running")) + assert(error.getMessage.contains("Wait for initialization to finish")) + assert(error.getMessage.contains("confirm the creating application has stopped")) + assert(!error.getMessage.contains(root)) + } + assert(AIPromptDelta.control(spark, qualified) == before && server.requests.get() == 0) + observed += event + }) + assert(observed == Set("input_committed", "input_validated")) + assert(output.count() == 1 && server.requests.get() == 1) + assert(AIPrompt.getOutputStatus(spark, root).isComplete) + } + } + + test("SQL variable substitution output paths reject before durable storage or inference") { + delta() + Seq(true, false).foreach { enabled => + Seq(false, true).foreach { defined => + val session = spark.newSession() + session.conf.set("spark.sql.variable.substitute", enabled.toString) + if (defined) session.conf.set("aiprompt_test_segment", "substituted") + import session.implicits._ + withServer { server => + val parent = path() + val root = new File(parent, "${aiprompt_test_segment}").toString + val stage = prompt(root, server) + val input = Seq("hello").toDF("text") + val calls = Seq[() => Any]( + () => stage.transform(input), + () => stage.transformSchema(input.schema), + () => stage.resume(session), + () => AIPrompt.getOutputStatus(session, root), + () => AIPrompt.loadOutput(session, root), + () => AIPrompt.releaseOutputOwner(session, root, "stopped-owner", confirmStopped = true)) + calls.foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("outputPath must not contain SQL variable substitutions")) + assert(!error.getMessage.contains(parent)) + } + assert(!new File(parent).exists() && server.requests.get() == 0) + } + } + } + } + + private def mediaInput: DataFrame = { + val session = spark + import session.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 = { + delta() + 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 root = path() + val stage = configure(new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root).setOutputBatchSize(4)) + def verify(rows: Array[Row]): Unit = { + 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) + } + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + verify(lazyStage.transform(input).collect()) + assert(server.requests.size() == sends) + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + verify(output.collect()) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.isComplete && status.rowErrorCount == 3 && status.committedCount == 4) + verify(stage.resume(spark).collect()) + assert(server.requests.size() == 2 * sends && status.ownerToken.isEmpty) + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"durable malformed inline $api paths retain lazy terminal row errors") { + verifyMedia(api, mediaInput, Seq("missing_input", "invalid_request", "invalid_request"), sends = 1) { + _.setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")) + } + } + + test(s"durable malformed inline $api messages retain local and provider terminal errors") { + verifyMedia(api, messageInput(api), Seq("invalid_request", "http_error", "http_error"), sends = 3) { + _.setMessagesCol("messages") + } + } + } + + test("durable malformed inline Decisions JSON remains terminal without sending invalid rows") { + 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) + } + } + + Seq(true).foreach { durable => + test(s"Decisions image URL casing and MIME parameters preserve payloads with durable=$durable") { + if (durable) delta() + 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")))) + Seq(false, true).foreach { json => + val input = if (json) messages.withColumn("input", to_json(col("messages"))).drop("messages") else messages + Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions()).foreach { stage => + 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") + if (durable) stage.setOutputPath(path()).setOutputBatchSize(accepted.size + rejected.size) + def verify(rows: Array[Row]): Unit = { + 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, asset) + assert(byAsset(asset).getAs[Row]("answer") != null) + } + rejected.foreach { asset => + assert(byAsset(asset).getAs[Row]("error").getString(0).contains("invalid_request"), asset) + assert(byAsset(asset).getAs[Row]("answer") == null) + } + } + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + verify(output.collect()) + assert(server.requests.size() == accepted.size) + if (durable) { + verify(stage.resume(spark).collect()) + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && status.rowErrorCount == rejected.size) + assert(server.requests.size() == accepted.size) + } + val sent = Iterator.continually(server.requests.poll()).takeWhile(_ != null).map { case (_, body) => + body.parseJson.asJsObject.fields("input").asInstanceOf[JsArray].elements.head.asJsObject + .fields("content").asInstanceOf[JsArray].elements.head.asJsObject.fields("image_url") + }.toSet + assert(sent == accepted.map(JsString(_)).toSet) + } + } + } + } + } + + test("durable malformed inline body media retains provider terminal errors") { + val input = mediaInput.withColumn("body", + to_json(struct(struct(struct(col("asset").as("url")).as("image_url")).as("extension")))) + verifyMedia("chat_completions", input, Seq.fill(3)("http_error"), sends = 4) { + _.setInputCol("text").setExtraBodyCol("body") + } + } + + private def underOutputSchema(value: JsValue): Seq[JsObject] = Seq( + JsObject("response_format" -> JsObject("json_schema" -> JsObject("schema" -> value))), + JsObject("response_format" -> JsObject("schema" -> value)), + JsObject("text" -> JsObject("format" -> JsObject("schema" -> value)))) + + private def underArrayOutputSchema(value: JsValue): Seq[JsObject] = Seq( + JsObject("response_format" -> JsArray(JsObject("json_schema" -> JsObject("schema" -> value)))), + JsObject("response_format" -> JsObject("json_schema" -> JsArray(JsObject("schema" -> value)))), + JsObject("text" -> JsArray(JsObject("format" -> JsObject("schema" -> value)))), + JsObject("text" -> JsObject("format" -> JsArray(JsObject("schema" -> value))))) + + private def bodyStage(api: String, url: String): AIPrompt = { + val stage = new AIPrompt().setApiType(api).setUrl(if (api == "systemone") url else url + "/v1") + .setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setOutputPath(path()) + if (api == "systemone") stage.setQuestions(NativeRequestFixtures.SystemQuestions) + else if (api == "decisions") stage.setQuestions(NativeRequestFixtures.DecisionQuestions) + stage + } + + Seq("inputCol", "promptTemplate", "questionsCol").foreach { binding => + test(s"durable native malformed shapes reject resources before sibling inference with $binding") { + delta() + import spark.implicits._ + val systemone = binding == "questionsCol" + val reply = if (systemone) NativeRequestFixtures.SystemResponse else NativeRequestFixtures.DecisionResponse + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val stage = new AIPrompt().setApiType(if (systemone) "systemone" else "decisions") + .setUrl(if (systemone) server.url else server.url + "/v1").setModel("model") + .setOutputCol("answer").setErrorCol("error") + if (systemone) stage.setInputCol("text").setQuestionsCol("payload") + else { + stage.setInputMode("json").setQuestions(NativeRequestFixtures.DecisionQuestions) + binding match { + case "inputCol" => stage.setInputCol("payload") + case "promptTemplate" => stage.setPromptTemplate("{payload}") + } + } + val resources = Seq("""{"file_id":"provider-file"}""", + """{"image_url":"https://example.invalid/media"}""") + resources.foreach { resource => + val invalids = if (systemone) Seq(s"""{"urgent":[$resource]}""", + s"""{"urgent":{"type":"choice","instructions":"Classify","criteria":{"label":[$resource]}}}""", + s"""{"urgent":{"type":"score","instructions":"Score","criteria":[$resource]}}""") + else Seq(resource) + val safe = if (systemone) NativeRequestFixtures.SystemQuestions else "\"hello\"" + invalids.foreach { invalid => + val root = path() + val data = Seq(("hello", safe), ("hello", invalid)).toDF("text", "payload") + val failure = intercept[org.apache.spark.SparkException](stage.setOutputPath(root).transform(data)) + assert(Iterator.iterate[Throwable](failure)(_.getCause).takeWhile(_ != null) + .exists(_.isInstanceOf[IllegalArgumentException])) + assert(server.requests.isEmpty) + assert(AIPromptDelta.control(spark, root).state == "initializing") + } + } + } + } + } + + test("durable native binding diagnostics identify restricted names without state values") { + import spark.implicits._ + val stage = new AISystemOne().setStateCol("state").setStateMode("json") + Seq("conversation", "prompt").foreach { name => + val value = JsObject(name -> JsString("private-state-value")).compactPrint + val error = intercept[org.apache.spark.SparkException](stage.validateDurableRows(Seq(value).toDF("state"))) + val cause = Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null).toVector.last + assert(cause.getMessage.contains(s"restricted fields: $name")) + assert(!cause.getMessage.contains("private-state-value")) + } + stage.validateDurableRows(Seq("""{"conversation":null,"prompt":null,"text":"private-state-value"}""").toDF("state")) + } + + test("durable output schema exemptions apply only to request bodies without a storage dependency") { + import spark.implicits._ + val resources = Seq(JsObject("file_id" -> JsString("provider-file")), + JsObject("image_url" -> JsString("https://example.invalid/media"))) + val stage = new AISystemOne().setStateCol("payload").setStateMode("json") + resources.flatMap(underOutputSchema).foreach { value => + val error = intercept[org.apache.spark.SparkException] { + stage.validateDurableRows(Seq(value.compactPrint).toDF("payload")) + } + assert(Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null) + .exists(_.isInstanceOf[IllegalArgumentException])) + } + val schema = JsObject("type" -> JsString("object"), "properties" -> JsObject( + "file_id" -> JsObject("type" -> JsString("string")), + "image_url" -> JsObject("type" -> JsString("string")))) + Seq("chat_completions" -> 0, "responses" -> 2).foreach { case (api, index) => + val value = underOutputSchema(schema)(index) + val data = Seq(("hello", value.compactPrint)).toDF("payload", "body") + val scalar = new AIPrompt().setApiType(api).setInputCol("payload").setExtraBody(value.compactPrint) + AIDurableInput.validate(scalar, data.schema) + new AIPrompt().setApiType(api).setInputCol("payload").setExtraBodyCol("body").validateDurableRows(data) + } + } + + test("durable output schema exemptions reject wrong APIs and array carriers before storage") { + val schema = new StructType().add("text", StringType) + val resources = Seq(JsObject("file_id" -> JsString("provider-file")), + JsObject("image_url" -> JsString("https://example.invalid/media"))) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + Seq("chat_completions", "responses", "systemone", "decisions", "embeddings").foreach { api => + val stage = bodyStage(api, server.url) + resources.foreach { resource => + val bodies = underOutputSchema(resource).zipWithIndex.collect { + case (body, index) if !(api == "chat_completions" && index == 0 || + api == "responses" && index == 2) => body + } ++ underArrayOutputSchema(resource) + bodies.foreach { body => + withClue(s"$api: ${body.compactPrint}: ") { + val error = intercept[IllegalArgumentException] { + stage.setExtraBody(body.compactPrint).transformSchema(schema) + } + assert(error.getMessage.contains("Durable")) + assert(!new File(stage.getOutputPath).exists() && server.requests.isEmpty) + } + } + } + } + } + } + + test("durable output schema exemptions validate body columns before sibling inference") { + delta() + import spark.implicits._ + Seq("chat_completions", "responses", "systemone", "decisions", "embeddings").foreach { api => + val reply = if (api == "responses") responsesReply else chatReply + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val crossApi = underOutputSchema(JsObject("file_id" -> JsString("provider-file")))( + if (api == "responses") 0 else 2) + val arrayBody = underArrayOutputSchema( + JsObject("image_url" -> JsString("https://example.invalid/media")))( + if (api == "responses") 3 else 1) + Seq(crossApi, arrayBody).foreach { body => + val stage = bodyStage(api, server.url).setExtraBodyCol("body") + val data = Seq(("plain", "{}"), ("invalid", body.compactPrint)).toDF("text", "body") + val error = withClue(s"$api: ${body.compactPrint}: ") { + intercept[org.apache.spark.SparkException](stage.transform(data)) + } + assert(Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null) + .exists(cause => cause.isInstanceOf[IllegalArgumentException] && cause.getMessage.contains("Durable"))) + assert(server.requests.isEmpty && AIPromptDelta.control(spark, stage.getOutputPath).state == "initializing") + } + if (api == "chat_completions" || api == "responses") { + val schema = JsObject("type" -> JsString("object"), "properties" -> JsObject( + "file_id" -> JsObject("type" -> JsString("string")), + "image_url" -> JsObject("type" -> JsString("string")))) + val body = underOutputSchema(schema)(if (api == "responses") 2 else 0) + val stage = bodyStage(api, server.url).setExtraBodyCol("body") + val data = Seq(("plain", "{}"), ("schema", body.compactPrint)).toDF("text", "body") + val output = stage.transform(data) + val rows = output.collect() + assert(rows.length == 2 && rows.forall(_.getAs[Row]("error") == null)) + assert(stage.resume(spark).collect().toSet == rows.toSet && server.requests.size() == 2) + val key = if (api == "responses") "text" else "response_format" + assert(server.requests.asScala.exists { case (_, sent) => + AIJson.objectString(sent).fields.get(key).contains(body.fields(key)) + }) + } + } + } + } + + private def arrayMedia(url: String): Seq[Map[String, Any]] = Seq( + Map("image_url" -> Seq(url)), Map("file_url" -> Seq(Seq(url))), + Map("image_url" -> Map("url" -> Seq(url), "detail" -> "auto")), + Map("file_data" -> Map("url" -> Seq(url))), + Map("file_data" -> Seq(Map("url" -> Seq(url)), url)), + Map("file_data" -> Seq(url)), Map("file_data" -> Seq(Seq(url)))) + + test("durable media arrays reject external locations without a storage dependency") { + import spark.implicits._ + arrayMedia("https://example.invalid/media").foreach { media => + val value = AIJson.fromAny(media).compactPrint + val stage = new AISystemOne().setStateCol("state").setStateMode("json") + intercept[org.apache.spark.SparkException](stage.validateDurableRows(Seq(value).toDF("state"))) + val scalar = new AIPrompt().setInputCol("state").setExtraBody(Map("extension" -> media)) + intercept[IllegalArgumentException](AIDurableInput.validate(scalar, Seq("hello").toDF("state").schema)) + } + val inline = arrayMedia(validImage).map(AIJson.fromAny(_).compactPrint) + new AISystemOne().setStateCol("state").setStateMode("json") + .validateDurableRows((inline :+ """{"file_data":"https://example.invalid/literal"}""").toDF("state")) + } + + test("durable media arrays in native state and body columns abort before inference") { + delta() + import spark.implicits._ + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + Seq(false, true).foreach { bodyColumn => + val stage = new AISystemOne().setUrl(server.url).setModel("model") + .setQuestions(NativeRequestFixtures.SystemQuestions).setOutputCol("answer").setErrorCol("error") + if (bodyColumn) stage.setStateCol("text").setExtraBodyCol("payload") + else stage.setStateCol("payload").setStateMode("json") + val media = arrayMedia("https://example.invalid/media") + (Seq(media(if (bodyColumn) 3 else 0)) ++ media.takeRight(2)).foreach { value => + val root = path() + val payload = AIJson.fromAny(value) + val data = Seq(("hello", "{}"), ("hello", payload.compactPrint)).toDF("text", "payload") + intercept[org.apache.spark.SparkException](stage.setOutputPath(root).transform(data)) + assert(server.requests.isEmpty && AIPromptDelta.control(spark, root).state == "initializing") + } + } + } + } + + Seq("stateCol", "inputCol", "promptTemplate").zipWithIndex.foreach { case (binding, index) => + test(s"durable System One JSON $binding rejects resources and commits malformed state without replay") { + delta() + import spark.implicits._ + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + Seq(new AISystemOne(), new AIPrompt().setApiType("systemone")).foreach { configured => + configured.setUrl(server.url).setModel("model").setInputMode("json") + .setQuestions(NativeRequestFixtures.SystemQuestions).setOutputCol("answer").setErrorCol("error") + binding match { + case "stateCol" => configured.setStateCol("payload") + case "inputCol" => configured.setInputCol("payload") + case "promptTemplate" => configured.setPromptTemplate("{payload}") + } + val before = server.requests.size() + val resources = Seq(JsObject("file_id" -> JsString("provider-file")), + JsObject("image_url" -> JsString("https://example.invalid/media"))) + val states = Seq("""{"nested":[{"file_id":"provider-file"}]}""") ++ + resources.map(value => underOutputSchema(value)(index).compactPrint) + states.foreach { state => + intercept[org.apache.spark.SparkException] { + configured.setOutputPath(path()).transform(Seq(state).toDF("payload")) + } + assert(server.requests.size() == before) + } + val root = path() + val stage = configured.copy(ParamMap.empty).setOutputPath(root) + val safe = """{"text":"https://example.invalid/inert","image_url":"data:image/png;base64,AQID"}""" + val data = Seq(safe, "not json").toDF("payload") + val result = stage.transform(data) + assert(result.schema == stage.transformSchema(data.schema)) + val rows = result.collect().map(row => row.getAs[String]("payload") -> row).toMap + assert(rows(safe).getAs[Row]("error") == null) + assert(rows("not json").getAs[Row]("error").getString(0).contains("invalid_request")) + assert(server.requests.size() == before + 1) + assert(server.requests.asScala.last._2.parseJson.asJsObject.fields("state") == safe.parseJson) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.state == "complete" && status.rowErrorCount == 1 && status.pendingCount == 0) + assert(stage.resume(spark).count() == 2 && server.requests.size() == before + 1) + } + } + } + } + + test("undefined durable lifecycle states fail status reads output reads and resume without inference") { + delta() + val session = spark + import session.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + stage.transform(Seq("original").toDF("text")) + val sends = server.requests.get() + val qualified = AIPromptDelta.qualify(session, root) + val control = AIPromptDelta.table(qualified, "_control") + Seq("ready", "unknown").foreach { state => + session.sql(s"UPDATE $control SET state = ${AIPromptDelta.literal(state)}").collect() + val status = intercept[IllegalArgumentException](AIPrompt.getOutputStatus(session, root)) + val output = intercept[IllegalArgumentException](AIPrompt.loadOutput(session, root)) + val resume = intercept[IllegalArgumentException](stage.resume(session)) + Seq(status, output, resume).foreach(error => assert(error.getMessage.contains("Corrupt AIPrompt run state"))) + assert(server.requests.get() == sends) + } + session.sql(s"UPDATE $control SET state = 'complete'").collect() + assert(AIPrompt.getOutputStatus(session, root).isComplete) + assert(AIPrompt.loadOutput(session, root).count() == 1) + assert(stage.resume(session).count() == 1 && server.requests.get() == sends) + } + } + + test("owned resume refuses before saved template validation and retains validation after release") { + delta() + val session = spark.newSession() + import session.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server).setPromptTemplate("{lower(text)}") + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("HELLO").toDF("text"), event => + if (event == "ready") throw new IllegalStateException("synthetic stopped owner")) + } + val qualified = AIPromptDelta.qualify(session, root) + val before = AIPromptDelta.control(session, qualified) + assert(before.state == "running" && before.owner.isDefined && server.requests.get() == 0) + session.udf.register("lower", (text: String) => Seq("changed:" + text)) + val owned = intercept[IllegalArgumentException](stage.resume(session)) + assert(owned.getMessage.contains("already has an owner")) + assert(AIPromptDelta.control(session, qualified) == before && server.requests.get() == 0) + AIPrompt.releaseOutputOwner(session, root, before.owner.get, confirmStopped = true) + val invalid = intercept[IllegalArgumentException](stage.resume(session)) + assert(invalid.getMessage.contains("Prompt templates must produce strings")) + assert(AIPromptDelta.control(session, qualified).owner.isEmpty && server.requests.get() == 0) + assert(stage.resume(session.newSession()).count() == 1 && server.requests.get() == 1) + assert(AIPrompt.getOutputStatus(session, root).isComplete) + } + } + + test("completed resume reads committed rows without resolving saved template functions again") { + delta() + val session = spark.newSession() + import session.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server).setPromptTemplate("{lower(text)}") + val original = stage.transform(Seq("HELLO", "WORLD").toDF("text")) + val rows = original.collect().toSet + val before = AIPrompt.getOutputStatus(session, root) + assert(before.isComplete && before.ownerToken.isEmpty && server.requests.get() == 2) + session.udf.register("lower", (text: String) => Seq("changed:" + text)) + val resumed = stage.resume(session) + assert(resumed.schema == original.schema && resumed.collect().toSet == rows) + assert(AIPrompt.getOutputStatus(session, root) == before && server.requests.get() == 2) + val mismatch = intercept[IllegalArgumentException] { + stage.copy(ParamMap.empty).setPromptTemplate("changed {lower(text)}").resume(session) + } + assert(mismatch.getMessage.contains("do not match the saved run")) + assert(server.requests.get() == 2) + } + } + + test("durable initialization rejects schema evolution before validating the committed input") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(spark, root) + val stage = prompt(root, server) + val error = intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, Seq("original").toDF("text"), event => + if (event == "input_committed") { + val manifest = AIPromptDelta.control(spark, qualified) + val input = AIPromptDelta.latest(spark, qualified, "_input")._2 + input.withColumn(manifest.keyColumn, AIPromptDelta.column(manifest.keyColumn) + lit(1L)) + .withColumn("foreign", lit("not supplied")).write.format("delta").mode("append") + .option("mergeSchema", "true").save(AIPromptDelta.child(qualified, "_input")) + }) + } + assert(error.getMessage.contains("schema evolution is not supported")) + val manifest = AIPromptDelta.control(spark, qualified) + assert(manifest.state == "initializing" && manifest.inputVersion.isEmpty && manifest.owner.isDefined) + assert(!new File(root, "results").exists() && server.requests.get() == 0) + intercept[IllegalArgumentException](stage.resume(spark)) + assert(server.requests.get() == 0) + } + } + + test("durable initialization allows compaction of the committed input without changing its rows") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(spark, root) + val stage = prompt(root, server).setOutputBatchSize(4) + val input = Seq("one", "two", "three", "four").toDF("text").repartition(2) + var compacted = Option.empty[Long] + val output = AIPromptOutput.transformObserved(stage, input, event => + if (event == "input_committed") { + assert(AIPromptDelta.version(spark, qualified, "_input") == 0L) + spark.sql(s"OPTIMIZE ${AIPromptDelta.table(qualified, "_input")}").collect() + compacted = Some(AIPromptDelta.version(spark, qualified, "_input")) + assert(compacted.get > 0L) + }) + assert(output.select("text").as[String].collect().toSet == Set("one", "two", "three", "four")) + assert(AIPromptDelta.control(spark, qualified).inputVersion == compacted) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded && server.requests.get() == 4) + assert(stage.resume(spark).count() == 4 && server.requests.get() == 4) + } + } + + Seq(false, true).foreach { complete => + test(s"durable input maintenance after initialization is rejected with guidance for complete=$complete") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(spark, root) + val stage = prompt(root, server).setOutputBatchSize(2) + val input = Seq("one", "two", "three", "four").toDF("text").repartition(2) + if (complete) { + assert(stage.transform(input).count() == 4) + } else { + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, input, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 2 && status.pendingCount == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + } + val before = AIPromptDelta.control(spark, qualified) + val resultsVersion = AIPromptDelta.version(spark, qualified, "results") + val sends = server.requests.get() + spark.sql(s"OPTIMIZE ${AIPromptDelta.table(qualified, "_input")}").collect() + assert(AIPromptDelta.version(spark, qualified, "_input") > before.inputVersion.get) + assert(AIPromptDelta.read(spark, qualified, "_input", before.inputVersion.get).count() == 4) + val operations = Seq[() => Any]( + () => AIPrompt.loadOutput(spark, root).collect(), + () => AIPrompt.getOutputStatus(spark, root), + () => stage.resume(spark).collect()) + operations.foreach { operation => + val error = intercept[IllegalArgumentException](operation()) + assert(error.getMessage.contains("immutable AIPrompt input snapshot")) + assert(error.getMessage.contains("do not OPTIMIZE or VACUUM")) + } + assert(AIPromptDelta.control(spark, qualified) == before) + assert(AIPromptDelta.version(spark, qualified, "results") == resultsVersion) + assert(server.requests.get() == sends && sends == (if (complete) 4 else 2)) + } + } + } + + Seq("input_committed", "input_validated", "results_created").foreach { boundary => + test(s"durable initialization retains the validated snapshot when input changes after $boundary") { + delta() + val session = spark + import session.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(session, root) + val stage = prompt(root, server) + val failure = Try(AIPromptOutput.transformObserved(stage, Seq("original").toDF("text"), event => + if (event == boundary) { + val manifest = AIPromptDelta.control(session, qualified) + val input = AIPromptDelta.latest(session, qualified, "_input")._2 + input.withColumn(manifest.keyColumn, AIPromptDelta.column(manifest.keyColumn) + lit(1L)) + .withColumn("text", lit("not validated")).write.format("delta").mode("append") + .save(AIPromptDelta.child(qualified, "_input")) + })).failed.toOption + assert(server.requests.get() == 0, "Unvalidated input must not reach inference") + assert(failure.exists(error => error.isInstanceOf[IllegalArgumentException] && + error.getMessage.contains("immutable AIPrompt input snapshot"))) + val manifest = AIPromptDelta.control(session, qualified) + assert(manifest.state == "initializing" && manifest.inputVersion.isEmpty && manifest.owner.isDefined) + intercept[IllegalArgumentException](stage.resume(session)) + assert(server.requests.get() == 0) + } + } + } +} 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..48c2084ce70 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptLoggingSuite.scala @@ -0,0 +1,552 @@ +// 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.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)(_.readAllBytes()) + 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)(_.readAllBytes()) + 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/AIPromptOutputCancellationSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputCancellationSuite.scala new file mode 100644 index 00000000000..9d10ffdc093 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputCancellationSuite.scala @@ -0,0 +1,161 @@ +// 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.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} +import spray.json._ + +import java.io.File +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +class AIPromptOutputCancellationSuite extends AIPromptOutputTestBase { + Seq("one", "ROW_ERROR").foreach { text => + test(s"last-batch outcome flags do not imply owner release for $text") { + delta() + val session = spark + import session.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + val stopped = new IllegalStateException("stop after last batch commit") + val failure = intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq(text).toDF("text"), event => + if (event == "batch_committed") throw stopped) + } + assert(failure eq stopped) + val status = AIPrompt.getOutputStatus(session, root) + val succeeds = text != "ROW_ERROR" + assert(status.inputCount == 1 && status.committedCount == 1 && status.pendingCount == 0) + assert(status.rowErrorCount == (if (succeeds) 0L else 1L)) + assert(status.isComplete && status.allSucceeded == succeeds) + assert(status.state == "running" && status.ownerToken.isDefined) + val json = AIPrompt.getOutputStatusJson(session, root).parseJson.asJsObject.fields + assert(json("isComplete") == JsBoolean(true) && json("allSucceeded") == JsBoolean(succeeds)) + assert(json("state") == JsString("running") && json("ownerToken") == JsString(status.ownerToken.get)) + val committed = AIPrompt.loadOutput(session, root).collect().toSeq + assert(committed.size == 1 && server.requests.get() == 1) + intercept[IllegalArgumentException](stage.resume(session)) + assert(AIPrompt.getOutputStatus(session, root) == status && server.requests.get() == 1) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + assert(stage.resume(session).collect().toSeq == committed) + val complete = AIPrompt.getOutputStatus(session, root) + assert(complete == status.copy(state = "complete", ownerToken = None)) + assert(complete.isComplete && complete.allSucceeded == succeeds && server.requests.get() == 1) + } + } + } + + test("durable preparation errors survive interrupted ownership release") { + delta() + val session = spark.newSession() + import session.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server).setPromptTemplate("{lower(text)}") + val data = Seq("ONE", "TWO", "THREE").toDF("text") + val startHolder = new CountDownLatch(1) + val holding = new CountDownLatch(1) + val release = new CountDownLatch(1) + val driver = new AtomicReference[Thread]() + val failure = new AtomicReference[Exception]() + val interrupted = new AtomicBoolean() + val pool = Executors.newFixedThreadPool(2) + try { + val holder = pool.submit(new Callable[Unit] { + override def call(): Unit = { + assert(startHolder.await(120, TimeUnit.SECONDS)) + AIPromptLocalTransaction(AIPromptDelta.qualify(session, root)) { + holding.countDown() + assert(release.await(120, TimeUnit.SECONDS)) + } + } + }) + val run = pool.submit(new Callable[Unit] { + override def call(): Unit = { + driver.set(Thread.currentThread()) + try { + AIPromptOutput.transformObserved(stage, data, event => if (event == "batch_committed") { + session.udf.register("lower", (text: String) => "changed:" + text) + startHolder.countDown() + assert(holding.await(10, TimeUnit.SECONDS)) + }) + } catch { + case error: Exception => failure.set(error) + } finally interrupted.set(Thread.currentThread().isInterrupted) + } + }) + assert(holding.await(120, TimeUnit.SECONDS)) + eventually(timeout(Span(60, Seconds))) { + assert(driver.get().getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.locks.ReentrantLock" && + frame.getMethodName == "lockInterruptibly")) + } + driver.get().interrupt() + run.get(10, TimeUnit.SECONDS) + assert(failure.get().isInstanceOf[IllegalArgumentException]) + assert(failure.get().getMessage.contains("stable built-in expressions")) + assert(failure.get().getSuppressed.exists(_.isInstanceOf[InterruptedException])) + assert(interrupted.get()) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 2 && status.pendingCount == 1 && status.ownerToken.isDefined) + assert(server.requests.get() == 2) + release.countDown() + holder.get(10, TimeUnit.SECONDS) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + assert(AIPrompt.getOutputStatus(session, root).ownerToken.isEmpty) + } finally { + startHolder.countDown() + release.countDown() + pool.shutdownNow() + assert(pool.awaitTermination(30, TimeUnit.SECONDS)) + } + } + } + + test("local control interruption preserves the caller flag and does not leak the lock") { + val root = new File(path()).toPath + val holding = new CountDownLatch(1) + val release = new CountDownLatch(1) + val caller = new AtomicReference[Thread]() + val entered = new AtomicBoolean() + val pool = Executors.newFixedThreadPool(2) + try { + val holder = pool.submit(new Callable[Unit] { + override def call(): Unit = AIPromptLocalTransaction(root.toUri.toString) { + holding.countDown() + assert(release.await(30, TimeUnit.SECONDS)) + } + }) + assert(holding.await(10, TimeUnit.SECONDS)) + val waiting = pool.submit(new Callable[Boolean] { + override def call(): Boolean = { + caller.set(Thread.currentThread()) + try { + AIPromptLocalTransaction(root.toUri.toString) { entered.set(true) } + false + } catch { + case _: InterruptedException => Thread.currentThread().isInterrupted + } + } + }) + eventually(timeout(Span(10, Seconds))) { + assert(Option(caller.get()).exists(_.getStackTrace.exists(frame => + frame.getClassName == "java.util.concurrent.locks.ReentrantLock" && + frame.getMethodName == "lockInterruptibly"))) + } + caller.get().interrupt() + assert(waiting.get(10, TimeUnit.SECONDS)) + assert(!entered.get()) + release.countDown() + holder.get(10, TimeUnit.SECONDS) + assert(AIPromptLocalTransaction(root.toUri.toString) { 3 } == 3) + } finally { + release.countDown() + pool.shutdownNow() + assert(pool.awaitTermination(10, TimeUnit.SECONDS)) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputGuardsSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputGuardsSuite.scala new file mode 100644 index 00000000000..a78321be9f0 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputGuardsSuite.scala @@ -0,0 +1,175 @@ +// 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.DataFrame +import org.apache.spark.sql.functions.lit + +import java.nio.file.{Files, Paths} +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} +import scala.util.Try + +class AIPromptOutputGuardsSuite extends AIPromptOutputTestBase { + private def input: DataFrame = { + val session = spark + import session.implicits._ + Seq("first", "second").toDF("text") + } + + for (resume <- Seq(false, true)) { + test(s"concurrent local owners arbitrate before inference with resume=$resume") { + delta() + withServer { server => + val root = path() + val data = input + def stage: AIPrompt = prompt(root, server).setOutputBatchSize(1) + if (resume) { + AIPromptFabricRecovery.stopAfterFirstBatch(stage, data) + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == 1 && partial.pendingCount == 1 && server.requests.get() == 1) + AIPrompt.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + } + val arrived = new CountDownLatch(2) + val proceed = new CountDownLatch(1) + val pool = Executors.newFixedThreadPool(2) + val futures = Seq.fill(2)(pool.submit(new Callable[Try[Long]] { + override def call(): Try[Long] = Try { + val event: String => Unit = boundary => { + if (boundary == (if (resume) "before_acquire" else "before_control_create")) { + arrived.countDown() + assert(proceed.await(180, TimeUnit.SECONDS), "Timed out waiting for the competing owner") + } + } + val output = if (resume) AIPromptOutput.resumeObserved(stage, spark, event) + else AIPromptOutput.transformObserved(stage, data, event) + output.count() + } + })) + try { + assert(arrived.await(180, TimeUnit.SECONDS), "Both owners must reach the same transaction boundary") + proceed.countDown() + val outcomes = futures.map(_.get(180, TimeUnit.SECONDS)) + assert(outcomes.count(_.isSuccess) == 1, outcomes.toString) + assert(outcomes.filter(_.isSuccess).map(_.get) == Seq(2L)) + val error = outcomes.find(_.isFailure).get.failed.get + val diagnostic = if (resume) "AIPrompt ownership conflict" else "AIPrompt initial ownership conflict" + assert(error.isInstanceOf[IllegalArgumentException] && error.getMessage.contains(diagnostic)) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.state == "complete" && status.ownerToken.isEmpty && status.allSucceeded) + assert(server.requests.get() == 2, "The losing owner must not send or replay any model request") + } finally { + proceed.countDown() + futures.foreach(_.cancel(true)) + pool.shutdownNow() + assert(pool.awaitTermination(60, TimeUnit.SECONDS), "Competing owners did not stop") + } + } + } + } + + for (mutation <- Seq("delete", "foreign-key", "version")) { + test(s"active committed ledger rejects $mutation before another request and retains its owner") { + delta() + withServer { server => + val root = path() + val stage = prompt(root, server).setOutputBatchSize(1) + val boundary = if (mutation == "version") "before_ledger_version_check" else "batch_committed" + var owner = Option.empty[String] + var changed = false + val error = intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, input, event => { + if (event == boundary && !changed) { + changed = true + val manifest = AIPromptDelta.control(spark, root) + owner = manifest.owner + val table = AIPromptDelta.table(root, "results") + mutation match { + case "delete" => spark.sql(s"DELETE FROM $table") + case "foreign-key" => + val location = AIPromptDelta.child(root, "results") + spark.read.format("delta").load(location) + .withColumn(manifest.keyColumn, lit(Long.MaxValue)) + .write.format("delta").mode("append").save(location) + case "version" => spark.sql(s"ALTER TABLE $table SET TBLPROPERTIES ('aiprompt.test' = 'changed')") + } + } + }) + } + val expected = mutation match { + case "delete" => "AIPrompt committed batch cardinality does not match the selected input keys" + case "foreign-key" => "AIPrompt results contain unknown row keys" + case "version" => "AIPrompt committed outcomes changed outside the current owner" + } + assert(changed && error.getMessage.contains(expected)) + assert(owner.isDefined && AIPromptDelta.control(spark, root).owner == owner) + assert(server.requests.get() == (if (mutation == "version") 0 else 1)) + if (mutation == "foreign-key") { + val manifest = AIPromptDelta.control(spark, root) + spark.sql(s"DELETE FROM ${AIPromptDelta.table(root, "results")} " + + s"WHERE `${manifest.keyColumn}` = ${Long.MaxValue}") + } + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == (if (mutation == "foreign-key") 1 else 0)) + assert(AIPrompt.loadOutput(spark, root).count() == partial.committedCount) + assert(server.requests.get() == (if (mutation == "version") 0 else 1)) + AIPrompt.releaseOutputOwner(spark, root, owner.get, confirmStopped = true) + assert(stage.resume(spark).count() == 2) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.state == "complete" && status.ownerToken.isEmpty && status.allSucceeded) + assert(server.requests.get() == (if (mutation == "delete") 3 else 2), + "Recovery must send only the rows missing from the repaired committed ledger") + } + } + } + + test("durable runtime rejects unqualified schemes and a missing catalog without accessing storage") { + delta() + Seq("s3a", "s3", "gs", "dbfs").foreach { scheme => + val error = intercept[IllegalArgumentException] { + AIPromptDelta.requireRuntime(spark, s"$scheme://synthetic-bucket/run") + } + assert(error.getMessage.contains("Other filesystem/LogStore configurations are not qualified")) + } + val missingCatalog = spark.newSession() + missingCatalog.conf.unset("spark.sql.catalog.spark_catalog") + val root = path() + val error = intercept[IllegalArgumentException] { + AIPromptDelta.requireRuntime(missingCatalog, Paths.get(root).toUri.toString) + } + assert(error.getMessage.contains("Configure the Delta Spark extension and catalog")) + assert(!Files.exists(Paths.get(root))) + } + + test("durable System One JSON state checks recursive resources without a storage dependency") { + import spark.implicits._ + val stage = new AISystemOne().setStateCol("state").setStateMode("json") + val external = Seq( + """{"nested":[{"file_id":"private-state-value"}]}""" -> "restricted fields: file_id", + """{"image_url":"https://example.invalid/private-state-value"}""" -> "Durable media must be inline", + """{"previous_response_id":"private-state-value"}""" -> "restricted fields: previous_response_id", + """{"tools":[{"type":"web_search","query":"private-state-value"}]}""" -> "restricted fields: tools") + external.foreach { case (state, diagnostic) => + val error = intercept[org.apache.spark.SparkException] { + stage.validateDurableRows(Seq(state).toDF("state")) + } + val causes = Iterator.iterate[Throwable](error)(_.getCause).takeWhile(_ != null).toVector + assert(causes.exists(cause => cause.isInstanceOf[IllegalArgumentException] && + Option(cause.getMessage).exists(_.contains(diagnostic)))) + assert(causes.forall(cause => !Option(cause.getMessage).exists(_.contains("private-state-value")))) + } + stage.validateDurableRows(Seq("not json", "null", """{"file_id":null}""", + """{"text":"https://example.invalid/inert","image_url":"data:image/png;base64,AQID"}""").toDF("state")) + stage.setStateMode("text").validateDurableRows(external.map(_._1).toDF("state")) + } + + test("public durable transform rejects a streaming input before storage or HTTP") { + withServer { server => + val root = path() + val data = spark.readStream.format("rate").load().selectExpr("CAST(value AS STRING) AS text") + val error = intercept[IllegalArgumentException](prompt(root, server).transform(data)) + assert(error.getMessage.contains("AIPrompt durable output does not support streaming inputs")) + assert(!Files.exists(Paths.get(root)) && server.requests.get() == 0) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputPreflightSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputPreflightSuite.scala new file mode 100644 index 00000000000..a8c024d6a26 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputPreflightSuite.scala @@ -0,0 +1,795 @@ +// 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.spark.ml.param.ParamMap +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.functions._ +import spray.json._ + +import java.io.File +import java.nio.file.Files +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit, TimeoutException} +import java.util.concurrent.atomic.AtomicBoolean +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.util.Using + +class AIPromptOutputPreflightSuite extends AIPromptOutputTestBase { + test("rendered template types are validated before durable initialization") { + import spark.implicits._ + withServer { server => + Seq("{array(text)}", "{encode(text, 'UTF-8')}").foreach { template => + val root = path() + val stage = prompt(root, server).setPromptTemplate(template) + val data = Seq("hello").toDF("text") + val schemaError = intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + val transformError = intercept[IllegalArgumentException](stage.transform(data)) + Seq(schemaError, transformError).foreach { error => + assert(error.getMessage.contains("Prompt templates must produce strings") && error.getCause == null) + } + assert(!new File(root).exists() && server.requests.get() == 0) + } + } + } + + test("durable malformed response schema names are terminal row errors and are not resent on resume") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val data = Seq( + ("invalid", """{"properties":{},"name":null}"""), + ("valid", """{"type":"text"}""")).toDF("text", "format") + val stage = prompt(root, server).setResponseFormatCol("format") + 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("invalid").getAs[Row]("error").getString(0).contains("invalid_request")) + assert(rows("valid").getAs[Row]("error") == null && server.requests.get() == 1) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.state == "complete" && status.pendingCount == 0 && status.ownerToken.isEmpty) + assert(stage.resume(spark).count() == 2 && server.requests.get() == 1) + } + } + + test("durable direct and facade Decisions commit blank inputs as terminal errors without resending") { + delta() + import spark.implicits._ + Seq(new AIPrompt().setApiType("decisions"), new OpenAIDecisions()).foreach { stage => + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val root = path() + val data = Seq(" \t\n", " evidence ").toDF("text") + stage.setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setQuestions(NativeRequestFixtures.DecisionQuestions).setErrorCol("error") + .setOutputPath(root).setOutputBatchSize(2) + 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(" \t\n").getAs[Row]("error").getString(0).contains("invalid_request")) + assert(rows(" evidence ").getAs[Row]("error") == null && server.requests.size() == 1) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.state == "complete" && status.pendingCount == 0 && status.ownerToken.isEmpty) + assert(stage.resume(spark).count() == 2 && server.requests.size() == 1) + } + } + } + + test("global API versions do not prevent v1 durable recovery or replay committed requests") { + delta() + import spark.implicits._ + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + try { + GlobalParams.resetGlobalParam(OpenAIApiVersion) + withServer { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("one", "two", "three").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after first commit")) + } + val before = AIPrompt.getOutputStatus(spark, root) + assert(before.committedCount == 2 && before.pendingCount == 1 && server.requests.get() == 2) + AIPrompt.releaseOutputOwner(spark, root, before.ownerToken.get, confirmStopped = true) + GlobalParams.setGlobalParam(OpenAIApiVersion, Left("irrelevant-version")) + val resumed = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url).resume(spark) + assert(resumed.count() == 3 && server.requests.get() == 3) + GlobalParams.setGlobalParam(OpenAIApiVersion, Right("absentVersion")) + assert(AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url).resume(spark).count() == 3) + assert(server.requests.get() == 3 && AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + + test("durable templates reject mutable functions and runtime values but snapshot materialized input expressions") { + delta() + val session = spark.newSession() + import session.implicits._ + session.udf.register("aiprompt_render_ticket", (text: String) => "old:" + text) + val data = Seq("hello", "hello", "hello").toDF("text") + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + Seq("{aiprompt_render_ticket(text)}", "{concat('prefix:', aiprompt_render_ticket(text))}", + "{current_timestamp()}", "{current_user()}", "{rand()}").foreach { template => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate(template) + val error = intercept[IllegalArgumentException](stage.transform(data)) + assert(error.getMessage.contains("stable built-in expressions") && Option(error.getCause).isEmpty) + assert(!new File(root).exists() && server.requests.isEmpty) + } + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate("{aiprompt_render_ticket(text)}") + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + assert(lazyStage.transform(data.limit(1)).count() == 1 && contents == Vector("old:hello")) + server.requests.clear() + stage.setPromptTemplate( + "{element_at(transform(array(rendered), x -> x), 1)}:{cast(round(pi()) as int)}:{cast(round(e()) as int)}") + val materialized = data.withColumn("rendered", expr("aiprompt_render_ticket(text)")) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, materialized, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && contents == Vector.fill(2)("old:hello:3:3")) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + session.udf.register("aiprompt_render_ticket", (text: String) => "new:" + text) + assert(stage.resume(session).count() == 3 && contents == Vector.fill(3)("old:hello:3:3")) + } + } + + Seq(false, true).foreach { recovering => + val operation = if (recovering) "resume" else "transform" + test(s"durable $operation releases quiescent ownership after a between-batch preparation failure") { + delta() + val session = spark.newSession() + import session.implicits._ + val initiallyCommitted = if (recovering) 2 else 0 + val total = initiallyCommitted + 3 + val rebound = new AtomicBoolean() + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((count, _) => { + if (count > initiallyCommitted && rebound.compareAndSet(false, true)) { + session.udf.register("lower", (text: String) => "changed:" + text) + } + 200 -> reply + })) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate("{lower(text)}") + val data = Seq.fill(total)("HELLO").toDF("text") + if (recovering) { + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("unknown commit acknowledgement")) + } + val interrupted = AIPrompt.getOutputStatus(session, root) + assert(interrupted.committedCount == 2 && interrupted.ownerToken.isDefined) + AIPrompt.releaseOutputOwner(session, root, interrupted.ownerToken.get, confirmStopped = true) + } + val error = intercept[IllegalArgumentException] { + if (recovering) stage.resume(session) else stage.transform(data) + } + assert(error.getMessage.contains("stable built-in expressions")) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && status.ownerToken.isEmpty && server.requests.size() == total - 1) + intercept[IllegalArgumentException](stage.resume(session)) + assert(AIPrompt.getOutputStatus(session, root).ownerToken.isEmpty && server.requests.size() == total - 1) + val resumed = stage.resume(session.newSession()) + assert(resumed.count() == total && resumed.schema == stage.transformSchema(data.schema)) + assert(server.requests.size() == total && AIPrompt.getOutputStatus(session, root).allSucceeded) + } + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"$api attachments follow input column order through copied and loaded partial durable resumes") { + delta() + val session = spark + import session.implicits._ + val left = "data:image/png;base64,AQID" + val right = "data:image/png;base64,BAUG" + val data = Seq.fill(3)((left, right)).toDF("left", "right").select("right", "left") + val reply = if (api == "chat_completions") { + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + } else { + """{"status":"completed","output":[{"type":"message","role":"assistant", + |"content":[{"type":"output_text","text":"answer"}]}]}""".stripMargin + } + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setApiType(api).setOutputBatchSize(1) + .setPromptTemplate("Compare {left} and {`RiGhT`}; literal {{LEFT}}") + .setColumnTypes(ListMap("LEFT" -> "path", "RIGHT" -> "path")) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + assert(lazyStage.transform(data.limit(1)).count() == 1) + def stopAfterBatch(event: String): Unit = + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch") + intercept[IllegalStateException](AIPromptOutput.transformObserved(stage, data, stopAfterBatch)) + val first = AIPrompt.getOutputStatus(session, root) + assert(first.state == "running" && first.pendingCount == 2 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(session, root, first.ownerToken.get, confirmStopped = true) + val reordered = stage.copy(ParamMap.empty).setColumnTypes(ListMap("RIGHT" -> "path", "LEFT" -> "path")) + intercept[IllegalStateException](AIPromptOutput.resumeObserved(reordered, session, stopAfterBatch)) + val second = AIPrompt.getOutputStatus(session, root) + assert(second.state == "running" && second.pendingCount == 1 && server.requests.size() == 3) + AIPrompt.releaseOutputOwner(session, root, second.ownerToken.get, confirmStopped = true) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-ordered-attachments").resolve("model").toString + reordered.write.save(saved) + val loaded = AIPrompt.load(saved) + assert(loaded.resume(session).count() == 3 && server.requests.size() == 4) + assert(loaded.resume(session).count() == 3 && server.requests.size() == 4) + assert(AIPrompt.getOutputStatus(session, root).state == "complete") + 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("durable native questions retain order through persistence and reject order-only changes before ownership") { + delta() + val session = spark + import session.implicits._ + val names = Seq("e_one", "d_two", "c_three", "b_four", "a_five", "z_six") + def questions(order: Seq[String], criteriaOrder: Seq[String]): JsObject = + JsObject(ListMap(order.map(name => name -> JsObject("type" -> JsString("choice"), + "instructions" -> JsString(name), + "criteria" -> JsObject(ListMap(criteriaOrder.map(label => label -> JsString(label)): _*)))): _*)) + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.SystemResponse)) { server => + val root = path() + val stage = new AIPrompt().setApiType("systemone").setUrl(server.url).setModel("model") + .setInputCol("text").setQuestions(questions(names, names)).setOutputPath(root).setOutputBatchSize(1) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("hello", "hello").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && server.requests.size() == 1) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + Seq(questions(names.reverse, names), questions(names, names.reverse)).foreach { changed => + intercept[IllegalArgumentException](stage.copy(ParamMap.empty).setQuestions(changed).resume(session)) + assert(server.requests.size() == 1 && AIPrompt.getOutputStatus(session, root).ownerToken.isEmpty) + } + val saved = Files.createTempDirectory(tmpDir, "aiprompt-ordered-questions").resolve("model").toString + stage.copy(ParamMap.empty).write.save(saved) + assert(AIPrompt.load(saved).resume(session).count() == 2 && server.requests.size() == 2) + server.requests.asScala.foreach { case (_, body) => + val sent = AIJson.objectString(body).fields("questions").asJsObject + assert(sent.fields.keys.toSeq == names) + sent.fields.values.foreach(question => + assert(question.asJsObject.fields("criteria").asJsObject.fields.keys.toSeq == names)) + } + } + } + + test("durable templates reject subqueries before writes while snapshotted input expressions remain supported") { + delta() + val session = spark + import session.implicits._ + val view = "aiprompt_template_context" + Seq("old").toDF("prefix").createOrReplaceTempView(view) + val data = Seq("hello", "hello", "hello").toDF("text") + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + try { + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + Seq(s"{(SELECT prefix FROM $view)} {text}", + s"{concat((SELECT prefix FROM $view), ' ', text)}", + s"{(SELECT max(prefix) FROM $view)} {text}").foreach { template => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate(template) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + server.requests.clear() + assert(lazyStage.transform(data.limit(1)).count() == 1 && contents == Vector("old hello")) + val error = intercept[IllegalArgumentException](stage.transform(data)) + assert(error.getMessage.contains("cannot contain SQL subqueries") && Option(error.getCause).isEmpty) + assert(!new File(root).exists() && server.requests.size() == 1) + } + server.requests.clear() + val root = path() + val materialized = data.withColumn("prefix", expr(s"(SELECT prefix FROM $view)")) + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, + server.url + "/v1/chat/completions").setPromptTemplate("{prefix} {text}") + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, materialized, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && contents == Vector("old hello", "old hello")) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + Seq("new").toDF("prefix").createOrReplaceTempView(view) + assert(stage.resume(session).count() == 3 && contents == Vector("old hello", "old hello", "old hello")) + } + } finally { + session.catalog.dropTempView(view) + } + } + + test("durable resume preserves ordered schemas in typed options and body extensions") { + delta() + val session = spark + import session.implicits._ + Seq(false, true).foreach { bodyExtension => + withServer { server => + val root = path() + val names = Seq("first", "second") + def configure(stage: AIPrompt, order: Seq[String]): AIPrompt = { + val schema = Map[String, Any]("type" -> "object", + "properties" -> ListMap(order.map(_ -> Map("type" -> "string")): _*), + "required" -> names, "additionalProperties" -> false) + if (bodyExtension) stage.setExtraBody(Map("response_format" -> Map("type" -> "json_schema", + "json_schema" -> Map("name" -> "ordered", "strict" -> true, "schema" -> schema)))) + else stage.setResponseSchema(schema) + } + val stage = configure(prompt(root, server).setOutputBatchSize(1), names) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("hello", "hello").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && server.requests.get() == 1) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + intercept[IllegalArgumentException](configure(stage.copy(ParamMap.empty), names.reverse).resume(session)) + assert(server.requests.get() == 1 && AIPrompt.getOutputStatus(session, root).ownerToken.isEmpty) + val saved = Files.createTempDirectory(tmpDir, "aiprompt-resume-ordered").resolve("model").toString + stage.copy(ParamMap.empty).write.save(saved) + assert(AIPrompt.load(saved).resume(session).count() == 2 && server.requests.get() == 2) + } + } + } + + test("resume rejects changed SQL expression semantics but permits execution tuning") { + delta() + val first = spark.newSession() + first.conf.set("spark.sql.ansi.enabled", "false") + first.conf.set("spark.sql.legacy.sizeOfNull", "true") + val resumed = first.newSession() + resumed.conf.set("spark.sql.ansi.enabled", "true") + resumed.conf.set("spark.sql.legacy.sizeOfNull", "true") + val data = first.range(3).withColumn("tags", expr("cast(null as array)")) + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/chat/completions") + .setPromptTemplate("{coalesce(size(tags), 0)}") + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(first, root) + assert(status.pendingCount == 1 && contents == Vector("-1", "-1")) + AIPrompt.releaseOutputOwner(first, root, status.ownerToken.get, confirmStopped = true) + val mismatch = intercept[IllegalArgumentException](stage.resume(resumed)) + Seq("original SynapseML build", "Spark runtime", "library and SQL defaults").foreach { guidance => + assert(mismatch.getMessage.contains(guidance)) + } + resumed.conf.set("spark.sql.ansi.enabled", "false") + resumed.conf.set("spark.sql.legacy.sizeOfNull", "false") + intercept[IllegalArgumentException](stage.resume(resumed)) + assert(contents.size == 2 && AIPrompt.getOutputStatus(first, root).ownerToken.isEmpty) + resumed.conf.set("spark.sql.legacy.sizeOfNull", "true") + resumed.conf.set("spark.sql.shuffle.partitions", "1") + assert(stage.resume(resumed).count() == 3 && contents == Vector("-1", "-1", "-1")) + assert(stage.resume(resumed).count() == 3 && contents.size == 3) + } + } + + test("durable path normalization preserves missing-input errors and self-contained attachments") { + delta() + val session = spark + import session.implicits._ + val data = Seq(None, Some(""), Some(" \t "), Some("data:text/plain;base64,aGVsbG8="), + Some(" \tDATA:text/plain;BASE64,aGVsbG8= \t")).toDF("asset") + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}, + |"finish_reason":"stop"}]}""".stripMargin + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setPromptTemplate("{asset}").setColumnTypes(Map("asset" -> "path")) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + def verify(rows: Array[Row]): Unit = { + 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")) + } + } + def verifyPrompts(): Unit = server.requests.asScala.foreach { case (_, body) => + val message = AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last.asJsObject + val text = message.fields("content").asInstanceOf[JsString].value + assert(text.takeWhile(_ != '\n') == "attachment.txt") + } + verify(lazyStage.transform(data).collect()) + assert(server.requests.size() == 2) + verifyPrompts() + val output = stage.transform(data) + verify(output.collect()) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.state == "complete" && status.rowErrorCount == 3 && status.committedCount == 5) + assert(stage.resume(session).count() == 5 && server.requests.size() == 4) + verifyPrompts() + val remote = stage.copy(ParamMap.empty).setOutputPath(path()) + intercept[org.apache.spark.SparkException] { + remote.transform(Seq(Some(" \thttps://example.invalid/media \t")).toDF("asset")) + } + assert(server.requests.size() == 4) + } + } + + test("local failed control initialization does not poison a fresh public run") { + delta() + val session = spark + import session.implicits._ + withServer { server => + val root = path() + val qualified = AIPromptDelta.qualify(spark, root) + intercept[IllegalStateException] { + AIPromptLocalTransaction(qualified) { throw new IllegalStateException("before control write") } + } + assert(Files.exists(new File(root, AIPromptLocalTransaction.LockFile).toPath)) + assert(!new File(root, "_control").exists()) + assert(prompt(root, server).transform(Seq("hello").toDF("text")).count() == 1) + assert(server.requests.get() == 1 && AIPrompt.getOutputStatus(spark, root).state == "complete") + } + } + + test("local control transactions serialize same-driver callers and release after failures") { + val root = new File(path()).toPath + val entered = new CountDownLatch(1) + val secondStarted = new CountDownLatch(1) + val release = new CountDownLatch(1) + val pool = Executors.newFixedThreadPool(2) + try { + val first = pool.submit(new Callable[Int] { + override def call(): Int = AIPromptLocalTransaction(root.toUri.toString) { + entered.countDown() + assert(release.await(10, TimeUnit.SECONDS)) + 1 + } + }) + assert(entered.await(10, TimeUnit.SECONDS)) + val second = pool.submit(new Callable[Int] { + override def call(): Int = { + secondStarted.countDown() + AIPromptLocalTransaction(root.resolve(".").toUri.toString) { 2 } + } + }) + assert(secondStarted.await(10, TimeUnit.SECONDS)) + intercept[TimeoutException](second.get(1, TimeUnit.SECONDS)) + release.countDown() + assert(first.get(10, TimeUnit.SECONDS) == 1 && second.get(10, TimeUnit.SECONDS) == 2) + intercept[IllegalStateException] { + AIPromptLocalTransaction(root.toUri.toString) { throw new IllegalStateException("transaction failed") } + } + assert(AIPromptLocalTransaction(root.toUri.toString) { 3 } == 3) + } finally { + release.countDown() + pool.shutdownNow() + assert(pool.awaitTermination(10, TimeUnit.SECONDS)) + } + } + + test("durable wildcard templates preserve public inputs through partial recovery") { + delta() + val session = spark + import session.implicits._ + val data = Seq(("duplicate", "public"), ("duplicate", "public"), ("tail", "public")) + .toDF("text", "_aiprompt_row_public") + val reply = + """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/chat/completions") + .setPromptTemplate("{to_json(struct(*))}") + val lazyStage = stage.copy(org.apache.spark.ml.param.ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + val lazyOutput = lazyStage.transform(data) + assert(lazyOutput.collect().length == 3) + def contents: Seq[String] = server.requests.asScala.toSeq.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + }.sorted + val expected = contents + assert(expected.count(_.contains("duplicate")) == 2) + expected.foreach(text => assert(AIJson.objectString(text).fields.keySet == data.columns.toSet)) + server.requests.clear() + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.pendingCount == 1 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val result = stage.resume(spark) + assert(result.schema == AIPromptOutput.storageSchema(lazyOutput.schema)) + assert(result.count() == 3 && result.filter(col("text") === "duplicate").count() == 2) + assert(contents == expected && server.requests.size() == 3) + assert(stage.resume(spark).count() == 3 && server.requests.size() == 3) + } + } + + test("durable native questions reject external dependencies before inference") { + delta() + val session = spark + import session.implicits._ + val data = Seq("hello").toDF("text") + withServer { server => + Seq("systemone", "decisions").foreach { api => + def stage(root: String): AIPrompt = new AIPrompt().setApiType(api) + .setUrl(if (api == "systemone") server.url.stripSuffix("/v1/chat/completions") + else server.url.stripSuffix("/chat/completions")) + .setModel("model").setInputCol("text").setOutputPath(root) + Seq("file_id" -> JsString("file-example"), "tools" -> JsArray(JsObject("type" -> JsString("web_search")))) + .foreach { case (binding, value) => + val question = JsObject("type" -> JsString(if (api == "systemone") "noul" else "predicate"), + "instructions" -> JsString("Is this urgent?"), binding -> value) + val questions = if (api == "systemone") JsObject("urgent" -> question) else JsArray(question) + val scalarPath = path() + intercept[IllegalArgumentException](stage(scalarPath).setQuestions(questions.compactPrint).transform(data)) + assert(!new File(scalarPath).exists()) + intercept[org.apache.spark.SparkException] { + stage(path()).setQuestionsCol("questions") + .transform(data.withColumn("questions", lit(questions.compactPrint))) + } + } + val safe = if (api == "systemone") { + """{"file_id":{"type":"choice","instructions":"Classify this text", + |"criteria":{"tools":"Tool-related text","conversation":"Conversation-related text", + |"image_url":"https://example.invalid/inert-description"}}}""".stripMargin + } else { + """[{"type":"predicate","name":"file_id","instructions":"Does the text mention tools?"}]""" + } + val scalar = stage(path()).setQuestions(safe) + scalar.validateDurableInput(data.schema, session) + val perRow = stage(path()).setQuestionsCol("questions") + perRow.validateDurableRows(data.withColumn("questions", lit(safe))) + } + assert(server.requests.get() == 0) + } + } + + test("durable JSON postprocessing commits only actual parsing failures as row errors") { + delta() + val session = spark + import session.implicits._ + Seq("chat_completions", "responses").foreach { api => + Using.resource(new AIRecordingServer((_, body) => { + val text = if (body.contains("malformed")) "not json" + else if (body.contains("incompatible")) """{"value":"wrong"}""" + else if (body.contains("fragment")) """{"value":1},"extra":2""" + else if (body.contains("trailing")) """{"value":1} {"value":2}""" + else if (body.contains("singlequotes")) "{'value':1.0}" + else if (body.contains("nonnumeric")) """{"value":NaN}""" + else if (body.contains("duplicatekeys")) """{"value":1.0,"value":2.0}""" + else """{"value":null}""" + 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 + } + 200 -> reply + })) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setApiType(api).setPostProcessingOptions(Map("jsonSchema" -> "value DOUBLE")) + val invalid = Set("malformed", "incompatible", "fragment", "trailing") + val rows = stage.transform((invalid.toSeq ++ Seq("singlequotes", "nonnumeric", "duplicatekeys", "valid")) + .toDF("text")).collect() + assert(rows.length == 8 && rows.count(_.getAs[Row]("error") != null) == 4) + assert(rows.filter(row => invalid(row.getAs[String]("text"))) + .forall(_.getAs[Row]("error").getString(0).contains("parse_error"))) + assert(rows.find(_.getAs[String]("text") == "valid").get.getAs[Row]("answer").isNullAt(0)) + assert(rows.find(_.getAs[String]("text") == "singlequotes").get.getAs[Row]("answer").getDouble(0) == 1.0) + assert(rows.find(_.getAs[String]("text") == "nonnumeric").get.getAs[Row]("answer").getDouble(0).isNaN) + assert(rows.find(_.getAs[String]("text") == "duplicatekeys").get.getAs[Row]("answer").getDouble(0) == 2.0) + assert(rows.forall(_.getAs[String]("raw").contains("assistant"))) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.rowErrorCount == 4 && !status.allSucceeded && status.pendingCount == 0) + assert(stage.resume(spark).count() == 8 && server.requests.size() == 8) + } + } + } + + test("sub-nanosecond concurrent timeouts fail before durable storage or HTTP") { + delta() + import spark.implicits._ + withServer { server => + Seq(Double.MinPositiveValue, 0.5e-9, 0.999e-9).foreach { timeout => + val root = path() + val error = intercept[IllegalArgumentException] { + prompt(root, server).setConcurrentTimeout(timeout).transform(Seq("hello").toDF("text")) + } + assert(error.getMessage.contains("concurrentTimeout")) + assert(!new File(root).exists() && server.requests.get() == 0) + } + } + } + + test("invalid scalar questions and sub-millisecond timeouts fail before durable storage") { + delta() + import spark.implicits._ + Using.resource(new AIRecordingServer((_, _) => 200 -> "{}")) { server => + val stages = Seq( + new AIPrompt().setApiType("systemone").setQuestions("null"), + new AIPrompt().setApiType("decisions").setQuestions("null"), + new AIPrompt().setTimeout(0.0005)) + stages.foreach { stage => + val root = path() + stage.setUrl(if (stage.getApiType == "systemone") server.url else server.url + "/v1") + .setModel("model").setInputCol("text").setOutputPath(root) + intercept[IllegalArgumentException](stage.transform(Seq("hello").toDF("text"))) + assert(!new File(root).exists()) + } + assert(server.requests.isEmpty) + } + } + + test("durable decisions templates reject provider bindings on the committed snapshot before HTTP") { + delta() + import spark.implicits._ + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + def stage(root: String): AIPrompt = new AIPrompt().setApiType("decisions") + .setUrl(server.url + "/v1").setModel("native-model").setPromptTemplate("{payload}") + .setInputMode("json").setQuestions(NativeRequestFixtures.DecisionQuestions) + .setOutputPath(root).setOutputCol("answer").setErrorCol("error") + val external = + """[{"role":"user","content":[{"type":"input_text","text":"hello","file_id":"provider-file"}]}]""" + intercept[org.apache.spark.SparkException] { + stage(path()).transform(Seq(external).toDF("payload")) + } + assert(server.requests.isEmpty) + val root = path() + val valid = stage(root) + val rows = valid.transform(Seq("\"ordinary https://example.invalid text\"", "not json") + .toDF("payload")).collect() + assert(rows.count(_.getAs[Row]("error") == null) == 1) + assert(rows.count(row => Option(row.getAs[Row]("error")) + .exists(_.getString(0).contains("invalid_request"))) == 1) + assert(AIPrompt.getOutputStatus(spark, root).rowErrorCount == 1) + assert(valid.resume(spark).count() == 2 && server.requests.size() == 1) + } + } + + test("durable long Retry-After terminates with the original HTTP error and never replays a committed outcome") { + delta() + val session = spark + import session.implicits._ + val reply = """{"error":{"code":"rate_limit","message":"retry later"}}""" + Using.resource(new AIRecordingServer((_, _) => 429 -> reply, retryAfter = "120")) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1/chat/completions") + val row = stage.transform(Seq("one").toDF("text")).head() + assert(row.getAs[Row]("error").getString(0).contains("http_error")) + assert(row.getAs[String]("raw") == reply) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.rowErrorCount == 1 && !status.allSucceeded && status.pendingCount == 0) + assert(stage.resume(spark).count() == 1 && server.requests.size() == 1) + } + } + + test("durable map-backed chat content commits terminal errors without sending or replaying") { + delta() + withServer { server => + val root = path() + val stage = prompt(root, server) + stage.clear(stage.promptTemplate) + stage.setMessagesCol("messages") + 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.count(_.getAs[Row]("error") == null) == 1) + assert(AIPrompt.getOutputStatus(spark, root).rowErrorCount == 5) + assert(stage.resume(spark).count() == 6 && server.requests.get() == 1) + } + } + + test("durable body extensions reject known remote media but permit inert URL text and schemas") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val input = Seq("one").toDF("text") + val remote = "https://example.invalid/mutable" + Seq(Map("image_url" -> remote), Map("image_url" -> Map("url" -> remote)), + Map("file_url" -> remote), + Map("file_url" -> Map("url" -> remote)), Map("file_data" -> Map("url" -> remote))).foreach { media => + val root = path() + val stage = prompt(root, server).setExtraBody(Map("extension" -> media)) + intercept[IllegalArgumentException](stage.transform(input)) + assert(!new File(root).exists()) + } + Seq("file_url", "file_data").foreach { field => + val stage = prompt(path(), server).setExtraBodyCol("body") + intercept[org.apache.spark.SparkException] { + stage.transform(input.withColumn("body", lit(s"""{"extension":{"$field":{"url":"$remote"}}}"""))) + } + } + assert(server.requests.get() == 0) + val inline = "data:image/png;base64,AQID" + val schema = Map("type" -> "object", "properties" -> Map("file_url" -> Map("type" -> "string"))) + val safe = prompt(path(), server).setExtraBody(Map( + "extension" -> Map("image_url" -> inline, "file_url" -> Map("url" -> inline), + "file_data" -> Map("url" -> inline), "text" -> remote), + "response_format" -> Map("type" -> "json_schema", "json_schema" -> + Map("name" -> "result", "schema" -> schema)))) + assert(safe.transform(input).count() == 1 && server.requests.get() == 1) + } + } + + test("resume fingerprints the executing session rather than another active session") { + delta() + withServer { server => + val session = spark.newSession() + import session.implicits._ + val root = path() + val stage = prompt(root, server) + val previous = SparkSession.getActiveSession + val previousZone = spark.conf.get("spark.sql.session.timeZone") + session.conf.set("spark.sql.session.timeZone", "UTC") + spark.conf.set("spark.sql.session.timeZone", "UTC") + SparkSession.setActiveSession(spark) + try { + val input = Seq(1, 2, 3).toDF("id").withColumn("text", lit("2026-01-01 12:00:00").cast("timestamp")) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, input, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.pendingCount == 1 && server.requests.get() == 2) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + session.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + intercept[IllegalArgumentException](stage.resume(session)) + assert(server.requests.get() == 2) + session.conf.set("spark.sql.session.timeZone", "UTC") + spark.conf.set("spark.sql.session.timeZone", "Asia/Tokyo") + assert(stage.resume(session).count() == 3 && server.requests.get() == 3) + assert(SparkSession.getActiveSession.contains(spark)) + } finally { + spark.conf.set("spark.sql.session.timeZone", previousZone) + previous match { + case Some(active) => SparkSession.setActiveSession(active) + case None => SparkSession.clearActiveSession() + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputSuite.scala new file mode 100644 index 00000000000..b9d745fcbf8 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputSuite.scala @@ -0,0 +1,777 @@ +// 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.SparkSessionManagement +import com.microsoft.azure.synapse.ml.io.http.URLKey +import com.microsoft.azure.synapse.ml.param.GlobalParams +import com.sun.net.httpserver.{HttpExchange, HttpHandler, HttpServer} +import org.apache.spark.SparkConf +import org.apache.spark.ml.linalg.{SQLDataTypes, SparseVector, Vector, Vectors} +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import org.apache.hadoop.fs.{LocalFileSystem, Path} +import spray.json._ + +import java.io.File +import java.net.{InetSocketAddress, URI} +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.nio.file.attribute.PosixFilePermissions +import java.util.Collections +import java.util.concurrent.atomic.AtomicInteger +import scala.util.Try +import scala.util.Using + +private[ai] object AIPromptSessionFileSystem { + val Checks: AtomicInteger = new AtomicInteger() +} + +class AIPromptSessionFileSystem extends LocalFileSystem { + override def getWorkingDirectory: Path = Option(getConf) + .flatMap(conf => Option(conf.get("aiprompt.test.workingDir"))).map(new Path(_)) + .getOrElse(super.getWorkingDirectory) + + override def exists(path: Path): Boolean = { + AIPromptSessionFileSystem.Checks.incrementAndGet() + super.exists(path) + } +} + +private[ai] object AIPromptDeltaTestSession extends SparkSessionManagement { + def hasDelta: Boolean = Try(Class.forName("org.apache.spark.sql.delta.sources.DeltaDataSource")).isSuccess + + override def stopSparkSession(): Unit = { + super.stopSparkSession() + SparkSession.getActiveSession.filter(_.sparkContext.isStopped).foreach(_ => SparkSession.clearActiveSession()) + SparkSession.getDefaultSession.filter(_.sparkContext.isStopped).foreach(_ => SparkSession.clearDefaultSession()) + } + + override def sparkConfiguration: SparkConf = { + val conf = super.sparkConfiguration + .set("spark.sql.shuffle.partitions", "2") + .set("spark.default.parallelism", "2") + .set("spark.databricks.delta.snapshotPartitions", "1") + if (hasDelta) conf.set("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") + .set("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog") + conf + } +} + +private[ai] class AIPromptRecoveryRecordingServer extends AutoCloseable { + val requests: AtomicInteger = new AtomicInteger() + private val server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + server.createContext("/", new HttpHandler { + override def handle(exchange: HttpExchange): Unit = { + requests.incrementAndGet() + val stream = exchange.getRequestBody + val source = scala.io.Source.fromInputStream(stream, "UTF-8") + val body = try source.mkString finally source.close() + val failed = body.contains("ROW_ERROR") + val response = if (failed) { + """{"error":{"code":"invalid_request","message":"synthetic row error"}}""" + } else if (exchange.getRequestURI.getPath.endsWith("/embeddings")) { + """{"object":"list","model":"local-model","data":[{"object":"embedding", + |"index":0,"embedding":[1.25,0.0,-2.5]}],"usage":{"prompt_tokens":1,"total_tokens":1}}""".stripMargin + } else { + """{"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 bytes = response.getBytes(StandardCharsets.UTF_8) + exchange.getResponseHeaders.add("Content-Type", "application/json") + exchange.sendResponseHeaders(if (failed) 400 else 200, bytes.length) + val output = exchange.getResponseBody + try output.write(bytes) finally output.close() + } + }) + server.start() + val url: String = s"http://127.0.0.1:${server.getAddress.getPort}/v1/chat/completions" + override def close(): Unit = server.stop(0) +} + +private[ai] object AIPromptStorageTestConfig { + def configure(prompt: AIPrompt, path: String, url: String): AIPrompt = { + prompt.setApiType("chat_completions") + .setUrl(url) + .setDeploymentName("local-model") + .setSubscriptionKey("synthetic-key") + .setPromptTemplate("{text}") + .setOutputCol("answer") + .setErrorCol("error") + .setUsageCol("usage") + .setRawResponseCol("raw") + .setOutputBatchSize(2) + .setOutputPath(path) + } +} + +class AIPromptOutputSuite extends AIPromptOutputTestBase { + test("implicit Fabric scalar model preflight rejects unavailable models before creating durable output") { + delta() + val previous = GlobalParams.getGlobalParam(URLKey) + val root = path() + try { + GlobalParams.setGlobalParam(URLKey, "https://fabric.example.invalid/cognitive/openai") + val input = spark.createDataFrame(Collections.singletonList(Row("hello")), + new StructType().add("text", StringType)) + val stage = new TestableAIFabricAuthPrompt("fabric-durable-preflight").setModel("blocked-model") + .setInputCol("text").setOutputPath(root) + stage.transformSchema(input.schema) + val error = intercept[IllegalArgumentException](stage.transform(input)) + assert(error.getMessage.contains("Synthetic Fabric model is unavailable")) + assert(!Files.exists(new File(root).toPath)) + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(URLKey, value) + case None => GlobalParams.resetGlobalParam(URLKey) + } + } + } + + test("storage schema and codec preserve sparse vectors, nested metadata, and nullable values") { + val metadata = new MetadataBuilder().putString("description", "documented metadata").build() + val nested = StructType(Seq(StructField("vector", SQLDataTypes.VectorType, nullable = false, metadata))) + val schema = StructType(Seq( + StructField("vectors", ArrayType(nested, containsNull = false), nullable = false, metadata), + StructField("labels", MapType(StringType, IntegerType, valueContainsNull = false), nullable = false))) + val planned = AIPromptOutput.storageSchema(schema) + val json = AIPromptStorageCodec.schemaJson(planned) + assert(AIPromptStorageCodec.readSchema(json) == planned) + assert(!json.contains("VectorUDT")) + val sparse = Vectors.sparse(4, Array(1, 3), Array(2.0, 0.0)) + val row = Row(Seq(Row(sparse), Row(Vectors.dense(1.0, 2.0))), Map("one" -> 1)) + val input = spark.createDataFrame(Collections.singletonList(row), schema) + val decoded = AIPromptStorageCodec.decode(AIPromptStorageCodec.encode(input, planned), planned) + assert(decoded.schema == planned) + val result = decoded.head().getSeq[Row](0).head.getAs[Vector](0) + assert(result == sparse) + assert(result.isInstanceOf[SparseVector]) + assert(planned("vectors").metadata == metadata) + assert(planned("vectors").dataType.asInstanceOf[ArrayType].containsNull) + val missing = Row(Seq(Option.empty[Row].orNull, Row(Option.empty[Vector].orNull)), + Map("missing" -> Option.empty[Integer].orNull)) + val nullableInput = spark.createDataFrame(Collections.singletonList(missing), planned) + val nullableOutput = AIPromptStorageCodec.decode(AIPromptStorageCodec.encode(nullableInput, planned), planned) + assert(nullableOutput.head() == missing) + } + + test("storage schemas reject arbitrary class metadata and incompatible Delta names") { + intercept[IllegalArgumentException] { + AIPromptStorageCodec.readSchema("""{"type":"udt","class":"do.not.load.ArbitraryClass"}""") + } + intercept[IllegalArgumentException] { + AIPromptOutput.storageSchema(StructType(Seq(StructField("bad name", StringType)))) + } + intercept[IllegalArgumentException] { + AIPromptOutput.storageSchema(StructType(Seq(StructField("x", NullType)))) + } + intercept[IllegalArgumentException] { + AIPromptOutput.storageSchema(StructType(Seq(StructField("matrix", SQLDataTypes.MatrixType)))) + } + intercept[IllegalArgumentException] { + AIPromptOutput.storageSchema(StructType(Seq(StructField("A", StringType), StructField("a", StringType)))) + } + } + + test("vector decoding rejects missing stored elements instead of silently producing zero") { + val logical = StructType(Seq(StructField("vector", SQLDataTypes.VectorType))) + val physical = AIPromptStorageCodec.physical(logical).asInstanceOf[StructType] + val corrupt = Row(Row(false, 2, Seq.empty[Int], Seq[Any](1.0, Option.empty[java.lang.Double].orNull))) + val frame = spark.createDataFrame(Collections.singletonList(corrupt), physical) + val error = intercept[Exception](AIPromptStorageCodec.decode(frame, logical).collect()) + val causes = Iterator.iterate(Option(error: Throwable))(_.flatMap(e => Option(e.getCause))) + .takeWhile(_.nonEmpty).flatten + assert(causes.exists(e => Option(e.getMessage).exists(_.contains("Invalid encoded AIPrompt vector elements")))) + } + + test("missing Delta fails before creating a run") { + assume(!AIPromptDeltaTestSession.hasDelta, "Run in the no-Delta classpath profile.") + val root = path() + val error = intercept[IllegalStateException] { + AIPromptDelta.requireRuntime(spark, AIPromptDelta.qualify(spark, root)) + } + assert(error.getMessage.contains("Delta runtime")) + assert(!new File(root).exists()) + } + + test("fingerprint rotates scalar authentication but retains request settings") { + val schema = StructType(Seq(StructField("text", StringType))) + val first = new AIPrompt("same").setSubscriptionKey("first").setPromptTemplate("{text}") + val second = new AIPrompt("same").setSubscriptionKey("second").setPromptTemplate("{text}") + val output = StructType(Seq(StructField("answer", StringType))) + assert(AIPromptOutput.fingerprint(first, schema, output, spark) == + AIPromptOutput.fingerprint(second, schema, output, spark)) + val differentUid = new AIPrompt("different").setSubscriptionKey("rotated").setPromptTemplate("{text}") + differentUid.restoreDurableOutputDefaults(Map( + "outputCol" -> first.getOutputCol, "errorCol" -> first.getErrorCol, "messagesCol" -> first.getMessagesCol)) + assert(!differentUid.isSet(differentUid.responseIdCol)) + assert(AIPromptOutput.fingerprint(first, schema, output, spark) == + AIPromptOutput.fingerprint(differentUid, schema, output, spark)) + second.setCustomHeaders(Map("Authorization" -> "Bearer rotated")) + assert(AIPromptOutput.fingerprint(first, schema, output, spark) == + AIPromptOutput.fingerprint(second, schema, output, spark)) + second.setCustomHeaders(Map("x-request-behavior" -> "different")) + assert(AIPromptOutput.fingerprint(first, schema, output, spark) != + AIPromptOutput.fingerprint(second, schema, output, spark)) + second.setPromptTemplate("different {text}") + assert(AIPromptOutput.fingerprint(first, schema, output, spark) != + AIPromptOutput.fingerprint(second, schema, output, spark)) + val customHandler = new AIPrompt("same") + customHandler.setHandler(customHandler.getHandler) + intercept[IllegalArgumentException](AIPromptOutput.fingerprint(customHandler, schema, output, spark)) + } + + test("durable results are eager, bounded, terminal, storage-backed and keep duplicate source rows") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server) + val input = Seq("duplicate", "duplicate", "ROW_ERROR").toDF("text") + val result = stage.transform(input) + assert(server.requests.get() == 3) + assert(result.schema == stage.transformSchema(input.schema)) + assert(result.count() == 3) + result.collect() + assert(server.requests.get() == 3) + val status = AIPromptOutput.getOutputStatus(spark, root) + assert(status.isComplete && !status.allSucceeded && status.rowErrorCount == 1) + assert(status.committedCount == 3 && status.pendingCount == 0 && status.ownerToken.isEmpty) + val manifest = AIPromptDelta.control(spark, AIPromptDelta.qualify(spark, root)) + val snapshot = AIPromptDelta.input(spark, AIPromptDelta.qualify(spark, root), manifest) + assert(snapshot.select(manifest.keyColumn).distinct().count() == 3) + val text = spark.read.format("delta").load(AIPromptDelta.child(root, "_control")).toJSON.head() + assert(!text.contains("synthetic-key") && !text.contains("{text}") && !text.contains(server.url)) + val controlVersion = AIPromptDelta.version(spark, root, "_control") + val resumed = AIPromptOutput.resumeObserved(prompt(root, server).setSubscriptionKey("rotated"), spark, + event => fail(s"Completed resume must not acquire ownership or write state: $event")) + assert(resumed.count() == 3 && server.requests.get() == 3) + assert(AIPromptDelta.version(spark, root, "_control") == controlVersion) + assert(AIPromptDelta.control(spark, root) == manifest) + intercept[IllegalArgumentException](stage.transform(input)) + } + } + + test("durable malformed body columns commit terminal row errors instead of aborting initialization") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server).setExtraBodyCol("body") + val input = Seq(("valid", "{}"), ("malformed", "{"), ("nonobject", "[]"), + ("duplicate", """{"x":1,"x":2}""")).toDF("text", "body") + val output = stage.transform(input) + assert(output.count() == 4 && server.requests.get() == 1) + val rows = output.collect() + assert(rows.count(_.getAs[Row]("error") == null) == 1) + assert(rows.filter(_.getAs[Row]("error") != null) + .forall(_.getAs[Row]("error").getString(0).contains("invalid_request"))) + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.isComplete && status.rowErrorCount == 3) + assert(stage.resume(spark).count() == 4 && AIPrompt.loadOutput(spark, root).count() == 4) + assert(server.requests.get() == 1) + } + } + + test("empty input initializes typed readable output without HTTP") { + delta() + withServer { server => + val input = spark.createDataFrame(Collections.emptyList[Row](), + StructType(Seq(StructField("text", StringType, nullable = false)))) + val root = path() + val stage = prompt(root, server) + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + assert(output.count() == 0 && server.requests.get() == 0) + assert(AIPromptOutput.getOutputStatus(spark, root).allSucceeded) + } + } + + test("durable paths and directory checks honor the supplied session filesystem configuration") { + delta() + withServer { server => + val session = spark.newSession() + import session.implicits._ + val root = path() + val original = Option(spark.sparkContext.hadoopConfiguration.get("fs.file.impl")) + session.conf.set("fs.file.impl", classOf[AIPromptSessionFileSystem].getName) + session.conf.set("fs.file.impl.disable.cache", "true") + session.conf.set("aiprompt.test.workingDir", new File(root).toURI.toString) + val qualified = AIPromptDelta.qualify(session, "run") + assert(new File(new java.net.URI(qualified)) == new File(root, "run")) + AIPromptSessionFileSystem.Checks.set(0) + val stage = prompt(qualified, server) + val output = AIPromptOutput.transformObserved(stage, Seq("one").toDF("text"), event => + if (event == "before_control_create") assert(AIPromptSessionFileSystem.Checks.get() > 0)) + assert(output.count() == 1 && server.requests.get() == 1) + assert(AIPrompt.getOutputStatus(session, qualified).allSucceeded) + assert(Option(spark.sparkContext.hadoopConfiguration.get("fs.file.impl")) == original) + } + } + + test("durable malformed Decisions JSON commits terminal row errors without blocking valid rows") { + delta() + Using.resource(new AIRecordingServer((_, _) => 200 -> NativeRequestFixtures.DecisionResponse)) { server => + val session = spark + import session.implicits._ + val root = path() + val stage = new AIPrompt().setApiType("decisions").setUrl(server.url + "/v1").setModel("model") + .setInputCol("input").setInputMode("json").setQuestions(NativeRequestFixtures.DecisionQuestions) + .setOutputCol("answer").setErrorCol("error").setOutputPath(root).setOutputBatchSize(10) + val data = Seq("\"valid\"", "{", "null", """[{"role":"system","content":"bad role"}]""", + """[{"role":"user","content":[]}]""").toDF("input") + val result = stage.transform(data) + assert(result.count() == 5 && server.requests.size() == 1) + val errors = result.select("error").collect().flatMap(row => Option(row.getAs[Row](0))) + assert(errors.length == 4 && errors.forall(_.getString(0).contains("invalid_request"))) + assert(AIPrompt.getOutputStatus(spark, root).rowErrorCount == 4) + assert(stage.resume(spark).count() == 5 && AIPrompt.loadOutput(spark, root).count() == 5) + assert(server.requests.size() == 1) + } + } + + test("durable malformed messages commit row errors while valid messages are sent once") { + delta() + withServer { server => + val root = path() + val schema = StructType(Seq(StructField("messages", ArrayType( + StructType(Seq(StructField("role", StringType), StructField("content", StringType))))))) + val data = spark.createDataFrame(java.util.Arrays.asList( + Row(Seq(Row("user", "valid"))), Row(Seq(Row("user", Option.empty[String].orNull))), + Row(Seq(Option.empty[Row].orNull))), schema) + val stage = prompt(root, server) + stage.clear(stage.promptTemplate) + stage.setMessagesCol("messages") + val result = stage.transform(data) + assert(result.count() == 3 && server.requests.get() == 1) + assert(result.select("error").collect().count(!_.isNullAt(0)) == 2) + assert(AIPrompt.getOutputStatus(spark, root).rowErrorCount == 2) + assert(stage.resume(spark).count() == 3 && server.requests.get() == 1) + } + } + + test("durable extensions reject nested resources and provider retrieval before inference") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val data = Seq("one").toDF("text") + val bodies: Seq[Map[String, Any]] = Seq( + Map("attachments" -> Seq(Map("file_id" -> "synthetic-file"))), + Map("data_sources" -> Seq(Map("type" -> "azure_search", + "parameters" -> Map("endpoint" -> "https://search.example.invalid", "index_name" -> "docs")))), + Map("web_search_options" -> Map.empty[String, Any]), + Map("prompt" -> Map("id" -> "synthetic-stored-prompt"))) + bodies.foreach { body => + val root = path() + intercept[IllegalArgumentException](prompt(root, server).setExtraBody(body).transform(data)) + assert(!new File(root).exists() && server.requests.get() == 0) + } + val bodyRows = bodies.map(body => ("one", AIJson.fromAny(body).compactPrint)).toDF("text", "body") + intercept[org.apache.spark.SparkException] { + prompt(path(), server).setExtraBodyCol("body").transform(bodyRows) + } + assert(server.requests.get() == 0) + val schema = Map("type" -> "object", "properties" -> Map("file_id" -> Map("type" -> "string"), + "prompt" -> Map("type" -> "string"))) + val inert = Map("response_format" -> Map("type" -> "json_schema", + "json_schema" -> Map("name" -> "answer", "schema" -> schema)), "provider_option" -> true) + assert(prompt(path(), server).setExtraBody(inert).transform(data).count() == 1) + assert(server.requests.get() == 1) + } + } + + test("corrupt mandatory output bindings fail status read and resume instead of fabricating success") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server) + stage.transform(Seq("ROW_ERROR").toDF("text")) + val original = AIPromptDelta.control(spark, root).bindings.parseJson.asJsObject.fields + val corrupt = Seq(original - "errorCol", original - "outputCol", + original.updated("errorCol", JsString("answer")), original.updated("errorCol", JsNull), + original.updated("outputCol", original("errorCol"))) + corrupt.foreach { bindings => + val json = JsObject(bindings).compactPrint + spark.sql(s"UPDATE ${AIPromptDelta.table(root, "_control")} SET bindings = ${AIPromptDelta.literal(json)}") + intercept[IllegalArgumentException](AIPrompt.getOutputStatus(spark, root)) + intercept[IllegalArgumentException](AIPrompt.loadOutput(spark, root)) + intercept[IllegalArgumentException](stage.resume(spark)) + } + assert(server.requests.get() == 1) + } + } + + test("durable multibatch ledger workload records Spark job counts without replay") { + delta() + withServer { server => + val group = java.util.UUID.randomUUID().toString + val input = spark.range(8).select(col("id").cast("string").as("text")) + val stage = prompt(path(), server) + spark.sparkContext.setJobGroup(group, "AIPrompt durable ledger workload") + val start = System.nanoTime() + try { + val result = stage.transform(input) + assert(result.count() == 8 && server.requests.get() == 8) + val jobs = spark.sparkContext.statusTracker.getJobIdsForGroup(group).length + val seconds = (System.nanoTime() - start).toDouble / 1e9 + info(f"Durable ledger workload: rows=8 batches=4 jobs=$jobs seconds=$seconds%.2f") + } finally spark.sparkContext.clearJobGroup() + } + } + + test("real Delta restores embedding VectorType and metadata without retaining inference") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server).setApiType("embeddings") + .setUrl(server.url.replace("chat/completions", "embeddings")).setTextCol("text") + stage.clear(stage.promptTemplate) + val metadata = new MetadataBuilder().putString("source", "retained").build() + val input = Seq("embedding input").toDF("text").select(col("text").as("text", metadata)) + val output = stage.transform(input) + assert(output.schema == stage.transformSchema(input.schema)) + assert(output.schema("answer").dataType == SQLDataTypes.VectorType) + assert(output.schema("text").metadata == metadata) + assert(output.head().getAs[Vector]("answer") == Vectors.dense(1.25, 0.0, -2.5)) + assert(AIPromptOutput.loadOutput(spark, root).head().getAs[Vector]("answer") == Vectors.dense(1.25, 0.0, -2.5)) + assert(server.requests.get() == 1) + } + } + + test("an ambiguous acknowledged commit retains ownership and resumes committed errors without replay") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("ok", "ROW_ERROR", "last").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("simulated lost acknowledgement")) + } + val partial = AIPromptOutput.getOutputStatus(spark, root) + assert(partial.committedCount == 2 && partial.pendingCount == 1 && partial.ownerToken.isDefined) + assert(AIPromptOutput.loadOutput(spark, root).count() == 2 && server.requests.get() == 2) + intercept[IllegalArgumentException](stage.resume(spark)) + intercept[IllegalArgumentException] { + AIPromptOutput.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = false) + } + intercept[IllegalArgumentException] { + AIPromptOutput.releaseOutputOwner(spark, root, "wrong-owner", confirmStopped = true) + } + AIPromptOutput.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + assert(prompt(root, server).setSubscriptionKey("new-key").setOutputBatchSize(1).resume(spark).count() == 3) + assert(server.requests.get() == 3) + } + } + + test("every incomplete initialization boundary fails explicitly without HTTP") { + delta() + withServer { server => + val session = spark + import session.implicits._ + Seq("control_created", "input_committed", "results_created").foreach { boundary => + val root = path() + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(prompt(root, server), Seq("ok").toDF("text"), event => + if (event == boundary) throw new IllegalStateException("initialization failure")) + } + val failure = intercept[IllegalArgumentException](AIPromptOutput.loadOutput(spark, root)) + assert(failure.getMessage.contains("initialization is incomplete")) + assert(failure.getMessage.contains("new outputPath")) + val statusFailure = intercept[IllegalArgumentException](AIPrompt.getOutputStatus(spark, root)) + assert(statusFailure.getMessage.contains("new outputPath")) + intercept[IllegalArgumentException](prompt(root, server).resume(spark)) + } + assert(server.requests.get() == 0) + } + } + + test("resume rejects request and output changes before HTTP and restores implicit UID output bindings") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val first = prompt(root, server) + first.clear(first.outputCol) + first.clear(first.errorCol) + first.transform(Seq(java.sql.Timestamp.valueOf("2026-01-01 12:00:00")).toDF("text")) + val second = prompt(root, server) + second.clear(second.outputCol) + second.clear(second.errorCol) + assert(second.uid != first.uid) + assert(second.resume(spark).columns.contains(first.getOutputCol)) + val timeZone = spark.conf.get("spark.sql.session.timeZone") + try { + spark.conf.set("spark.sql.session.timeZone", if (timeZone == "UTC") "America/Los_Angeles" else "UTC") + val failure = intercept[IllegalArgumentException](second.resume(spark)) + assert(failure.getMessage.contains("configuration, public output schema or Spark SQL semantics")) + } finally spark.conf.set("spark.sql.session.timeZone", timeZone) + intercept[IllegalArgumentException](prompt(root, server).setOutputCol("changed").resume(spark)) + intercept[IllegalArgumentException](second.setPromptTemplate("changed {text}").resume(spark)) + assert(server.requests.get() == 1) + } + } + + test("resume restores retained generated messages after input and output collisions with a different UID") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val first = prompt(root, server).setDropPrompt(false).setOutputBatchSize(1) + val base = first.getMessagesCol + first.setOutputCol(base + "_2") + val input = Seq(("one", "user-base-one", "user-next-one"), ("two", "user-base-two", "user-next-two")) + .toDF("text", base, base + "_1") + val expected = first.transformSchema(input.schema) + val retained = base + "_3" + assert(expected.fieldNames.contains(retained)) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(first, input, event => + if (event == "batch_committed") throw new IllegalStateException("interrupted after a committed batch")) + } + val manifest = AIPromptDelta.control(spark, AIPromptDelta.qualify(spark, root)) + assert(manifest.outputType == expected) + assert(AIPromptOutput.loadOutput(spark, root).schema == expected) + val status = AIPromptOutput.getOutputStatus(spark, root) + AIPromptOutput.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val second = prompt(root, server).setDropPrompt(false).setOutputCol(first.getOutputCol) + val untouchedDefault = second.getMessagesCol + assert(untouchedDefault != base && second.uid != first.uid) + val resumed = second.resume(spark) + assert(resumed.schema == expected && resumed.count() == 2) + assert(second.getMessagesCol == untouchedDefault && !second.isSet(second.messagesCol)) + val preserved = resumed.select(AIPromptDelta.column(base), AIPromptDelta.column(base + "_1")).collect() + .map(row => row.getString(0) -> row.getString(1)).toSet + assert(preserved == Set("user-base-one" -> "user-next-one", "user-base-two" -> "user-next-two")) + assert(resumed.select(AIPromptDelta.column(retained)).head().getSeq[Row](0).nonEmpty) + assert(server.requests.get() == 2) + } + } + + test("saved input is evaluated once and unsupported inputs fail before run creation") { + delta() + withServer { server => + val evaluations = spark.sparkContext.longAccumulator("durable-source-evaluations") + val source = udf { id: Long => evaluations.add(1); s"row-$id" } + val input = spark.range(3).select(source(col("id")).as("text")) + prompt(path(), server).transform(input).count() + assert(evaluations.value == 3L) + val root = path() + val bad = input.withColumn("bad name", lit(1)) + intercept[IllegalArgumentException](prompt(root, server).transform(bad)) + assert(!new File(root).exists()) + assert(server.requests.get() == 3) + } + } + + test("credential columns fail before snapshot creation or HTTP even with distinct per-row principals") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val input = Seq(("one", "synthetic-principal-one"), ("two", "synthetic-principal-two")) + .toDF("text", "credential") + .withColumn("headers", map(lit("Authorization"), col("credential"))) + val settings: Seq[AIPrompt => Unit] = Seq( + stage => stage.setSubscriptionKeyCol("credential"), + stage => stage.setAADTokenCol("credential"), + stage => stage.setCustomAuthHeaderCol("credential"), + stage => stage.setVectorParam(stage.customHeaders, "headers")) + settings.foreach { setting => + val root = path() + val stage = prompt(root, server) + setting(stage) + intercept[IllegalArgumentException](stage.transform(input)) + assert(!new File(root).exists()) + } + assert(server.requests.get() == 0) + } + } + + test("durable validation rejects provider resource bindings without a storage dependency") { + val session = spark + import session.implicits._ + val input = Seq("ordinary https://example.invalid text").toDF("text") + val stage = new AIPrompt().setInputCol("text") + stage.validateDurableInput(input.schema, session) + stage.validateDurableRows(input) + intercept[IllegalArgumentException] { + stage.setPreviousResponseId("saved-response").validateDurableInput(input.schema, session) + } + stage.clear(stage.previousResponseId) + Seq("conversation", "file_id", "tools", "retrieve", "run_id").foreach { key => + stage.setExtraBody(Map("include" -> Seq("message.output_text.logprobs"), key -> "external-resource")) + intercept[IllegalArgumentException](stage.validateDurableInput(input.schema, session)) + } + stage.setExtraBodyCol("body") + stage.validateDurableRows(Seq("""{"provider_option":true}""", Option.empty[String].orNull).toDF("body")) + intercept[org.apache.spark.SparkException] { + stage.validateDurableRows(Seq("""{"file_id":"external-resource"}""").toDF("body")) + } + } + + test("durable row validation defers malformed content but rejects external references") { + val session = spark + import session.implicits._ + val stage = new AIPrompt().setMessagesCol("messages") + def check(content: String): Unit = { + val json = s"""{"messages":[{"role":"user","content":$content}]}""" + stage.validateDurableRows(spark.read.json(Seq(json).toDS())) + } + Seq( + "\"An ordinary https://example.invalid address in text\"", + """[{"type":"text","text":"hello"},{"type":"input_text","text":"hello"}]""", + """[{"type":"input_image","image_url":"data:image/png;base64,AQID","detail":"low"}]""", + """[{"type":"image_url","image_url":{"url":"data:image/png;base64,AQID"}}]""", + """[{"type":"input_file","file_data":"data:text/plain;base64,aGVsbG8=","filename":"note.txt"}]""", + """[{"type":"input_file","file_data":"aGVsbG8=","filename":"note.txt"}]""", + """[{"type":"input_file","file_data":"https://example.invalid/not-a-location"}]""", + """[{"type":"audio","data":"unknown"}]""", + """[{"type":"text","text":"hello","resource":"unsupported-field"}]""", + "null") + .foreach(check) + Seq( + """[{"type":"input_image","image_url":"https://example.invalid/image.png"}]""", + """[{"type":"input_file","file_id":"provider-file"}]""", + """[{"type":"input_file","file_url":"https://example.invalid/file"}]""").foreach { content => + intercept[org.apache.spark.SparkException](check(content)) + } + val decisions = new AIPrompt().setApiType("decisions").setInputCol("input").setInputMode("json") + decisions.validateDurableRows(Seq("\"inert URL https://example.invalid\"").toDF("input")) + intercept[org.apache.spark.SparkException] { + decisions.validateDurableRows(Seq( + """[{"role":"user","content":[{"type":"input_image","image_url":"https://example.invalid"}]}]""") + .toDF("input")) + } + } + + test("public durable media validation runs before attachment fetching or inference") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val stage = prompt(root, server).setPromptTemplate("{text} {asset}").setColumnTypes(Map("asset" -> "path")) + intercept[org.apache.spark.SparkException] { + stage.transform(Seq(("look at this", server.url + "/image.png")).toDF("text", "asset")) + } + assert(server.requests.get() == 0) + assert(AIPromptDelta.control(spark, AIPromptDelta.qualify(spark, root)).state == "initializing") + val valid = stage.copy(org.apache.spark.ml.param.ParamMap.empty).setOutputPath(path()) + val result = valid.transform(Seq(("https://example.invalid is inert text", + "data:text/plain;base64,aGVsbG8=")).toDF("text", "asset")) + assert(result.head().getAs[Row]("error") == null) + assert(result.count() == 1 && server.requests.get() == 1) + } + } + + test("permission failure at the real Delta commit retains ownership and never retries inference") { + delta() + assume(java.nio.file.FileSystems.getDefault.supportedFileAttributeViews().contains("posix"), + "Requires POSIX permissions for a real commit failure.") + withServer { server => + val session = spark + import session.implicits._ + val directory = Files.createTempDirectory(tmpDir, "aiprompt-commit-permission-") + val root = directory.resolve("run").toString + val log = directory.resolve("run").resolve("results").resolve("_delta_log") + val permissions = PosixFilePermissions.fromString("rwx------") + try { + Files.setPosixFilePermissions(directory, PosixFilePermissions.fromString("r-x------")) + assume(!Files.isWritable(directory), "Requires a non-root user with effective POSIX permission enforcement.") + } finally { + Files.setPosixFilePermissions(directory, permissions) + } + try { + val failure = intercept[Exception] { + AIPromptOutput.transformObserved(prompt(root, server), Seq("one").toDF("text"), boundary => + if (boundary == "before_batch_commit") { + Files.setPosixFilePermissions(log, PosixFilePermissions.fromString("r-x------")) + }) + } + val causes = Iterator.iterate(Option(failure: Throwable))(_.flatMap(e => Option(e.getCause))) + .takeWhile(_.isDefined).flatten.map(_.toString).mkString(" ") + assert(causes.contains("Permission denied") || causes.contains("AccessDenied")) + assert(server.requests.get() == 1) + val status = AIPromptOutput.getOutputStatus(spark, root) + assert(status.committedCount == 0 && status.pendingCount == 1 && status.ownerToken.isDefined) + assert(AIPromptOutput.loadOutput(spark, root).count() == 0) + assert(server.requests.get() == 1) + } finally { + if (Files.exists(log)) Files.setPosixFilePermissions(log, permissions) + } + } + } + + test("corrupt committed keys and modified input snapshots fail instead of returning empty success") { + delta() + withServer { server => + val session = spark + import session.implicits._ + Seq("results", "_input", "format", "null_format").foreach { table => + val root = path() + prompt(root, server).transform(Seq("one").toDF("text")) + if (table == "format") { + spark.sql(s"UPDATE ${AIPromptDelta.table(root, "_control")} SET formatVersion = 999") + } else if (table == "null_format") { + spark.sql(s"UPDATE ${AIPromptDelta.table(root, "_control")} SET formatVersion = NULL") + } else { + val location = AIPromptDelta.child(root, table) + spark.read.format("delta").load(location).write.format("delta").mode("append").save(location) + } + Seq[() => Any](() => AIPromptOutput.loadOutput(spark, root), + () => AIPromptOutput.getOutputStatus(spark, root), + () => prompt(root, server).resume(spark)).foreach { operation => + val error = intercept[IllegalArgumentException](operation()) + if (table == "format" || table == "null_format") { + assert(error.getMessage.contains("Unsupported AIPrompt output format version")) + } else if (table == "results") { + assert(error.getMessage.contains("Corrupt AIPrompt row identity: keys must be unique and non-null")) + } else { + assert(error.getMessage.contains("The immutable AIPrompt input snapshot has been modified")) + } + } + } + assert(server.requests.get() == 4) + } + } + + test("qualified ABFS storage runs and recovers with the caller's configured identity") { + delta() + val base = AILiveTestEnvironment.get("AIPROMPT_ABFS_TEST_ROOT") + assume(base.isDefined, "ABFS/OneLake qualification requires AIPROMPT_ABFS_TEST_ROOT and caller storage identity.") + assert(Try(new URI(base.get)).toOption.exists(uri => Set("abfs", "abfss").contains(uri.getScheme)), + "ABFS qualification requires an abfs:// or abfss:// test root.") + withServer { server => + val session = spark + import session.implicits._ + val root = base.get.stripSuffix("/") + "/" + java.util.UUID.randomUUID().toString + val stage = prompt(root, server) + try { + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("one", "two", "three").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("interruption")) + } + val status = AIPromptOutput.getOutputStatus(spark, root) + AIPromptOutput.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + assert(stage.resume(spark).count() == 3 && server.requests.get() == 3) + } finally { + val location = new org.apache.hadoop.fs.Path(root) + location.getFileSystem(spark.sparkContext.hadoopConfiguration).delete(location, true) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputTestBase.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputTestBase.scala new file mode 100644 index 00000000000..e57902f6daa --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptOutputTestBase.scala @@ -0,0 +1,39 @@ +// 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.{SparkSessionManagement, TestBase} + +/** Delta suites require a caller-provided test runtime; absence is reported as a skip. */ +private[ai] abstract class AIPromptOutputTestBase extends TestBase { + override lazy val sparkProvider: SparkSessionManagement = AIPromptDeltaTestSession + + protected def path(): String = tmpDir.resolve(java.util.UUID.randomUUID().toString).toAbsolutePath.toString + + protected def delta(): Unit = { + val available = AIPromptDeltaTestSession.hasDelta + if (sys.env.get("SYNAPSEML_REQUIRE_DELTA_TESTS").contains("true")) { + assert(available, "SYNAPSEML_REQUIRE_DELTA_TESTS=true requires Delta on the cognitive test classpath.") + } + assume(available, "This integration test requires a caller-provided Delta test classpath.") + } + + protected def withServer(test: AIPromptRecoveryRecordingServer => Unit): Unit = { + val server = new AIPromptRecoveryRecordingServer + try test(server) finally server.close() + } + + protected def prompt(path: String, server: AIPromptRecoveryRecordingServer): AIPrompt = + AIPromptStorageTestConfig.configure(new AIPrompt(), path, server.url) + + override def beforeAll(): Unit = { + super.beforeAll() + TestBase.stopSparkSession() + sparkProvider.resetSparkSession(numCores = Some(2)) + } + + override def afterAll(): Unit = { + try sparkProvider.stopSparkSession() finally super.afterAll() + } +} 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..d0ce1257d6f --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptPreflightSuite.scala @@ -0,0 +1,779 @@ +// 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 => + val saved = tmpDir.resolve("invalid-timeout-output") + 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 && !Files.exists(saved)) + } + } + } + + 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/AIPromptRecoveryAttachmentSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryAttachmentSuite.scala new file mode 100644 index 00000000000..8a53acd9b20 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryAttachmentSuite.scala @@ -0,0 +1,168 @@ +// 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.nio.file.{Files, Paths} +import java.util.concurrent.{CancellationException, ConcurrentLinkedQueue} +import java.util.concurrent.atomic.AtomicInteger +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptRecoveryAttachmentSuite extends AIPromptOutputTestBase { + Seq("chat_completions", "responses").foreach { api => + test(s"$api attachment size limits reject invalid Param bindings before HTTP or durable storage") { + 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 root = Files.createTempDirectory("aiprompt-file-limits-") + try { + 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") + durable <- Seq(true) + } { + 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) + } + if (durable) stage.setOutputPath(root.resolve("run").toString) + 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 && !Files.exists(root.resolve("run"))) + } + 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) + } finally org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } + + Seq(true).foreach { durable => + test(s"$api applies default attachment limits without making them explicit with durable=$durable") { + if (durable) delta() + 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") + if (durable) stage.setOutputPath(path()) + 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) + if (durable) { + assert(stage.resume(spark).count() == 1 && server.requests.isEmpty) + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && status.rowErrorCount == 1 && status.ownerToken.isEmpty) + } else { + 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 rejects resolver-equivalent columnTypes keys before HTTP or durable storage") { + 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) + val output = path() + 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") + for { + kind <- Seq("path", "text") + durable <- Seq(true) + } { + val stage = base.copy(ParamMap.empty).setColumnTypes(Map("asset" -> kind, "ASSET" -> "path")) + if (durable) stage.setOutputPath(output) + 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 && !Files.exists(Paths.get(output))) + } + } + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + } + Seq(true).foreach { durable => + Seq("chat_completions", "responses").foreach { api => + test(s"missing attachment placeholders preserve usable $api inputs with durable=$durable") { + if (durable) delta() + 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") + if (durable) stage.setOutputPath(path()).setOutputBatchSize(2) + 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")))) + } + if (durable) { + assert(stage.resume(spark).count() == 7 && server.requests.size() == 3) + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && status.rowErrorCount == 4 && status.ownerToken.isEmpty) + } + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryEndpointSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryEndpointSuite.scala new file mode 100644 index 00000000000..eed6a928e40 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryEndpointSuite.scala @@ -0,0 +1,328 @@ +// 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.ml.param.ParamMap +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.nio.file.Files +import java.util.Locale +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptRecoveryEndpointSuite extends AIPromptOutputTestBase { + 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 storage") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val keys = Seq(OpenAIDeploymentNameKey, OpenAIEmbeddingDeploymentNameKey, AIFoundryModel) + val previous = keys.map(key => key -> GlobalParams.getGlobalParam(key)) + val parent = Files.createTempDirectory("aiprompt-deployment-unicode-") + 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; + durable <- Seq(true)) { + 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)) + } + if (durable) stage.setOutputPath(parent.resolve("run").toString) + 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 && !Files.exists(parent.resolve("run"))) + } + } + } finally { + previous.foreach { + case (key, Some(value)) => GlobalParams.setGlobalParam(key, value) + case (key, None) => GlobalParams.resetGlobalParam(key) + } + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + } + for (custom <- Seq(false, true); component <- Seq("path", "query")) { + test(s"endpoint Unicode rejects malformed $component with custom=$custom before HTTP or storage") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val parent = Files.createTempDirectory("aiprompt-endpoint-unicode-") + try { + 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"); durable <- Seq(true)) { + 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) + } + if (durable) stage.setOutputPath(parent.resolve("run").toString) else stage.clear(stage.outputPath) + 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 && !Files.exists(parent.resolve("run"))) + } + } + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + } + test("dated API versions reject malformed scalar and default Unicode before HTTP or storage") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val previous = GlobalParams.getGlobalParam(OpenAIApiVersion) + val root = Files.createTempDirectory("aiprompt-version-unicode-") + 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"); durable <- Seq(true)) { + 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)) + } + if (durable) stage.setOutputPath(root.resolve("run").toString) + 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 && !Files.exists(root.resolve("run"))) + } + } + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } + 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 root = Files.createTempDirectory("aiprompt-responses-sampling-") + 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) + val durable = stage.copy(ParamMap.empty).setOutputPath(root.resolve("run").toString) + assert(durable.transformSchema(data.schema) == AIPromptOutput.storageSchema(schema)) + assert(server.requests.isEmpty && !Files.exists(root.resolve("run"))) + 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) + } + org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } + } + test("invalid endpoint ports fail lazy and durable preflight without creating storage") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val parent = Files.createTempDirectory("aiprompt-port-preflight-") + val root = parent.resolve("run") + try { + Seq(0, 65536, Int.MaxValue).foreach { port => + val endpoint = s"http://127.0.0.1:$port/v1" + Seq(false, true).foreach { custom => + Seq(true).foreach { durable => + val stage = new AIPrompt().setModel("model").setInputCol("text") + if (custom) stage.setCustomUrlRoot(endpoint) else stage.setUrl(endpoint) + if (durable) stage.setOutputPath(root.toString) + Seq[() => Any](() => stage.transformSchema(inputSchema), () => stage.transform(input)).foreach { call => + val error = intercept[IllegalArgumentException](call()) + assert(error.getMessage.contains("port") && !error.getMessage.contains(endpoint)) + assert(!Files.exists(root)) + } + } + } + } + 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) + stage.setOutputPath(root.toString).transformSchema(inputSchema) + assert(!Files.exists(root)) + } + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + + test("underscore endpoint hosts report the unsupported host before lazy or durable execution") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val parent = Files.createTempDirectory("aiprompt-host-preflight-") + val root = parent.resolve("run") + try { + 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) + Seq(true).foreach { durable => + if (durable) stage.setOutputPath(root.toString) else stage.clear(stage.outputPath) + 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)) + assert(!Files.exists(root)) + } + } + } + } + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + Seq("azure.com", "azure.us", "azure.cn").foreach { cloud => + test(s"Decisions rejects $cloud endpoints before lazy or durable execution") { + val input = spark.createDataFrame(Seq(Row("hello")).asJava, inputSchema) + val parent = Files.createTempDirectory("aiprompt-azure-preflight-") + val root = parent.resolve("run") + try { + 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) + durable <- Seq(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) + if (durable) stage.setOutputPath(root.toString) else stage.clear(stage.outputPath) + 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) && !Files.exists(root)) + } + } + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + } + 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) + val parent = Files.createTempDirectory("aiprompt-openai-root-") + val output = parent.resolve("run") + 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 durable = stage.copy(ParamMap.empty).setOutputPath(output.toString) + assert(durable.transformSchema(inputSchema) == AIPromptOutput.storageSchema(schema)) + assert(!Files.exists(output)) + 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) + } + org.apache.commons.io.FileUtils.deleteDirectory(parent.toFile) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryHttpPolicySuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryHttpPolicySuite.scala new file mode 100644 index 00000000000..a76f7e04660 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryHttpPolicySuite.scala @@ -0,0 +1,52 @@ +// 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.services.openai.OpenAIDecisions +import org.apache.spark.ml.param.ParamMap +import org.apache.spark.ml.util.Identifiable +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{MapType, StringType, StructType} +import spray.json._ + +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import scala.collection.JavaConverters._ +import scala.util.Using + +class AIPromptRecoveryHttpPolicySuite extends AIPromptOutputTestBase { + 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("authentication framing hooks fail scalar preflight before HTTP or durable storage") { + 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 => + val root = Files.createTempDirectory("ai-auth-framing-") + try { + cases.foreach { create => + Seq(true).foreach { durable => + val configured = create().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + if (durable) configured.setOutputPath(root.resolve("run").toString) + 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(!Files.exists(root.resolve("run")) && server.requests.isEmpty) + } + } + } finally org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryLoggingSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryLoggingSuite.scala new file mode 100644 index 00000000000..5a84b0bf8ce --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryLoggingSuite.scala @@ -0,0 +1,166 @@ +// 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.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 + +class AIPromptRecoveryLoggingSuite extends AIPromptOutputTestBase { + 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("durable batches construct one HTTP stage per public run without sharing it between runs") { + delta() + val data = spark.createDataFrame( + Seq(Row("first"), Row("second"), Row("third")).asJava, new StructType().add("text", StringType)) + withServer { server => + val stage = prompt(path(), server).setOutputBatchSize(1) + Seq(1, 2).foreach { concurrency => + stage.setOutputPath(path()).setConcurrency(concurrency) + withLogs { logged => + val saved = stage.transform(data) + assert(saved.count() == 3 && saved.filter("error IS NOT NULL").count() == 0) + val constructors = logged.asScala.count(message => + message.contains(".HTTPTransformer") && message.contains("\"method\":\"constructor\"")) + assert(constructors == 1, s"Expected one HTTP constructor per run, observed $constructors") + assert(AIPrompt.getOutputStatus(spark, stage.getOutputPath).pendingCount == 0) + } + } + assert(server.requests.get() == 6) + } + } + test("resume failures use the public recovery execution logger") { + val stage = new AIPrompt() + withLogs { logged => + val error = intercept[IllegalArgumentException](stage.resume(spark)) + assert(error.getMessage.contains("resume requires outputPath") && + error.getMessage.contains("original run directory")) + assert(logged.asScala.exists(_.contains("\"method\":\"resume\""))) + } + } + + for (method <- Seq("resume"); succeeds <- Seq(false, true)) { + test(s"public $method logging uses the execution session with succeeds=$succeeds") { + if (method == "resume" && succeeds) delta() + 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(path(), server) + if (method == "transform") { + stage.clear(stage.outputPath) + if (!succeeds) stage.setInputCol("missing") + } else if (succeeds) { + assert(stage.transform(data).count() == 1) + } else stage.clear(stage.outputPath) + prior match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + withLogs { logged => + def invoke(): DataFrame = if (method == "transform") stage.transform(data) else stage.resume(execution) + 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(s""""method":"$method"""")) + 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() + } + } + } + } + test("bounded durable recovery never depends on response cookies or resends committed rows") { + delta() + val session = spark.newSession() + session.conf.set("spark.sql.shuffle.partitions", "1") + val data = session.createDataFrame((1 to 4).map(number => Row(s"row-$number")).asJava, + new StructType().add("text", StringType)).coalesce(1) + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply, + responseHeaders = Map("Set-Cookie" -> "provider-session=synthetic; Path=/"))) { server => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setConcurrency(1).setSubscriptionKey("synthetic-first") + .setOutputPath(root).setOutputBatchSize(2) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") throw new IllegalStateException("stop after committed batch")) + } + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 2 && status.pendingCount == 2 && server.requests.size() == 2) + AIPrompt.releaseOutputOwner(session, root, status.ownerToken.get, confirmStopped = true) + stage.setSubscriptionKey("synthetic-rotated") + val rows = stage.resume(session).collect() + assert(rows.length == 4 && rows.forall(row => + row.getAs[String]("answer") == "answer" && row.getAs[Row]("error") == null)) + assert(stage.resume(session).count() == 4 && server.requests.size() == 4) + assert(server.requests.asScala.map(_._2).toSet.size == 4) + val headers = server.headers.asScala.toVector + assert(headers.map(_("api-key")) == + Vector.fill(2)("synthetic-first") ++ Vector.fill(2)("synthetic-rotated")) + assert(headers.forall(!_.contains("cookie"))) + assert(AIPrompt.getOutputStatus(session, root).pendingCount == 0) + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPostProcessingSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPostProcessingSuite.scala new file mode 100644 index 00000000000..eae331cff7a --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPostProcessingSuite.scala @@ -0,0 +1,41 @@ +// 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 AIPromptRecoveryPostProcessingSuite extends AIPromptOutputTestBase { + Seq("csv", "regex").foreach { mode => + test(s"durable $mode postprocessing keeps planned schemas and values") { + delta() + withServer { server => + val session = spark + import session.implicits._ + val root = path() + val options = if (mode == "csv") Map("delimiter" -> "n") else { + Map("regex" -> "a(nswer)", "regexGroup" -> "1") + } + val stage = prompt(root, server).setPostProcessingOptions(options) + val input = Seq("hello", "ROW_ERROR").toDF("text") + val result = stage.transform(input) + Seq(result, AIPrompt.loadOutput(spark, root), stage.resume(spark)).foreach { output => + assert(output.schema == stage.transformSchema(input.schema)) + val rows = output.collect() + val answered = rows.find(_.getAs[String]("text") == "hello").get + if (mode == "csv") assert(answered.getAs[Seq[String]]("answer") == Seq("a", "swer")) + else assert(answered.getAs[String]("answer") == "nswer") + assert(answered.getAs[Row]("error") == null) + val failed = rows.find(_.getAs[String]("text") == "ROW_ERROR").get + assert(failed.isNullAt(failed.fieldIndex("answer")) && failed.getAs[Row]("error") != null) + } + assert(server.requests.get() == 2) + assert(AIPrompt.getOutputStatus(spark, root).rowErrorCount == 1) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPreflightSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPreflightSuite.scala new file mode 100644 index 00000000000..a378be932a6 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryPreflightSuite.scala @@ -0,0 +1,276 @@ +// 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 AIPromptRecoveryPreflightSuite extends AIPromptOutputTestBase { + 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)) + Seq("chat_completions", "responses", "decisions").foreach { api => + test(s"$api explicit messages reject unused construction settings before HTTP or durable storage") { + 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 => + val root = Files.createTempDirectory("aiprompt-message-settings-") + try { + 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) => + Seq(true).foreach { durable => + val stage = configure(base.copy(ParamMap.empty)) + if (durable) stage.setOutputPath(root.resolve("run").toString) + assert(intercept[IllegalArgumentException](stage.transformSchema(data.schema)) + .getMessage.contains(name)) + assert(intercept[IllegalArgumentException](stage.transform(data)).getMessage.contains(name)) + assert(server.requests.isEmpty && !Files.exists(root.resolve("run"))) + } + } + 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() + } + } + } finally org.apache.commons.io.FileUtils.deleteDirectory(root.toFile) + } + } + } + Seq("systemone", "decisions", "embeddings").foreach { api => + test(s"$api ignores inactive global temperature in schema, requests and durable identity") { + 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 fingerprint = AIPromptOutput.fingerprint(original, input.schema, expectedSchema, spark) + 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) + assert(AIPromptOutput.fingerprint(stage.effectiveCopy(), input.schema, expectedSchema, spark) == + fingerprint) + 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) + val fingerprint = AIPromptOutput.fingerprint(explicit.effectiveCopy(), input.schema, schema, spark) + 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(AIPromptOutput.fingerprint(explicit.effectiveCopy(), input.schema, schema, spark) == fingerprint) + 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) + } + } + } + Seq("systemone", "decisions", "chat_completions", "responses", "embeddings", "custom", "global-v1") + .foreach { route => + test(s"global API versions are inactive for $route in schema, requests and durable identity") { + 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 fingerprint = AIPromptOutput.fingerprint(stage.effectiveCopy(), data.schema, expectedSchema, spark) + 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) + assert(AIPromptOutput.fingerprint(stage.effectiveCopy(), data.schema, expectedSchema, spark) == + fingerprint) + 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 fingerprint = AIPromptOutput.fingerprint(stage.effectiveCopy(), data.schema, + stage.transformSchema(data.schema), spark) + 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(AIPromptOutput.fingerprint(stage.effectiveCopy(), data.schema, + stage.transformSchema(data.schema), spark) != fingerprint) + } + } finally { + previous match { + case Some(value) => GlobalParams.setGlobalParam(OpenAIApiVersion, value) + case None => GlobalParams.resetGlobalParam(OpenAIApiVersion) + } + } + } + test("invalid concurrent timeout defaults fail before HTTP or durable storage") { + Using.resource(new AIRecordingServer((_, _) => 200 -> chatReply)) { server => + val saved = tmpDir.resolve("invalid-timeout-output") + Seq(0.0, -1.0, 1e-10, Double.NaN, Double.PositiveInfinity).foreach { timeout => + Seq(true).foreach { durable => + val stage = new TestableAIDefaultTimeoutPrompt("invalid-default-timeout").setConcurrentTimeoutDefault(timeout) + .setUrl(server.url + "/v1").setModel("model").setInputCol("text") + if (durable) stage.setOutputPath(saved.toString) + assert(intercept[IllegalArgumentException](stage.transformSchema(input.schema)) + .getMessage.contains("concurrentTimeout")) + assert(intercept[IllegalArgumentException](stage.transform(input)).getMessage.contains("concurrentTimeout")) + assert(server.requests.isEmpty && !Files.exists(saved)) + } + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryResponseFormatSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryResponseFormatSuite.scala new file mode 100644 index 00000000000..dd13f6c9972 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoveryResponseFormatSuite.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.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.io.File +import java.nio.file.Files +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.util.Using + +class AIPromptRecoveryResponseFormatSuite extends AIPromptOutputTestBase { + 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 storage or 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)) + Seq(true).foreach { durable => + val root = path() + val stage = new AIPrompt().setUrl(server.url + "/v1").setModel("model").setInputCol("text") + .setResponseFormatCol("format") + if (durable) stage.setOutputPath(root) + 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(!new File(root).exists() && server.requests.isEmpty) + } + } + } + } + + private val selfContainedSchema = Map("type" -> "object", "properties" -> Map( + "file_id" -> Map("type" -> "string"), "image_url" -> Map("type" -> "string"))) + private val selfContainedFormat: Map[String, Any] = + Map("type" -> "json_schema", "name" -> "self_contained", "schema" -> selfContainedSchema) + + replies.foreach { case (api, reply) => + test(s"durable $api typed response formats check resources outside output schemas") { + import spark.implicits._ + val input = Seq("hello").toDF("text") + val restricted = Seq[Map[String, Any]]( + Map("type" -> "json_object", "file_id" -> "provider-file"), + Map("type" -> "json_schema", "json_schema" -> + ((selfContainedFormat - "type") + ("tools" -> Seq("provider-tool")))), + selfContainedFormat + ("image_url" -> "https://example.invalid/media")) + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputPath(root) + restricted.foreach { value => + intercept[IllegalArgumentException](stage.setResponseFormat(value).transformSchema(input.schema)) + stage.setResponseFormatCol("format") + intercept[org.apache.spark.SparkException] { + stage.validateDurableRows(input.withColumn("format", lit(AIJson.fromAny(value).compactPrint))) + } + } + val maps = Seq(("hello", Map("type" -> "json_object", "file_id" -> "provider-file"))).toDF("text", "format") + intercept[org.apache.spark.SparkException](stage.validateDurableRows(maps)) + stage.setResponseFormat(selfContainedFormat).transformSchema(input.schema) + stage.setResponseFormatCol("format").validateDurableRows(input.withColumn("format", + lit(AIJson.fromAny(selfContainedFormat).compactPrint))) + stage.validateDurableRows(Seq("not json", """{"type":42}""", Option.empty[String].orNull).toDF("format")) + assert(!new File(root).exists() && server.requests.isEmpty) + } + } + + test(s"durable $api typed response format resources fail before inference and valid rows resume") { + delta() + import spark.implicits._ + val resource = Map("type" -> "json_object", "file_id" -> "provider-file") + val invalid = Seq(("hello", resource)).toDF("text", "format") + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + def stage: AIPrompt = new AIPrompt().setApiType(api).setUrl(server.url + "/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error").setOutputPath(path()) + val scalar = stage.setResponseFormat(resource) + intercept[IllegalArgumentException](scalar.transform(invalid)) + assert(!new File(scalar.getOutputPath).exists()) + val column = stage.setResponseFormatCol("format") + intercept[org.apache.spark.SparkException](column.transform(invalid)) + assert(AIPromptDelta.control(spark, AIPromptDelta.qualify(spark, column.getOutputPath)).state == "initializing") + intercept[IllegalArgumentException](column.resume(spark)) + assert(server.requests.isEmpty) + val valid = stage.setResponseFormatCol("format").setOutputBatchSize(2) + val input = Seq(("ok", AIJson.fromAny(selfContainedFormat).compactPrint), ("bad", "not json")) + .toDF("text", "format") + val transformed = valid.transform(input) + Seq(transformed, valid.resume(spark)).foreach { output => + assert(output.schema == valid.transformSchema(input.schema)) + val rows = output.collect().map(row => row.getAs[String]("text") -> row).toMap + assert(rows("ok").getAs[String]("answer") == "answer" && rows("ok").getAs[Row]("error") == null) + assert(rows("bad").getAs[Row]("error").getString(0).contains("invalid_request")) + } + val status = AIPrompt.getOutputStatus(spark, valid.getOutputPath) + assert(status.isComplete && status.rowErrorCount == 1 && server.requests.size() == 1) + assert(format(server.requests.peek()._2, api) == AIRequestAdapter.responseFormat(selfContainedFormat, api)) + } + } + } + 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"))) + } + } + + Seq(true).foreach { durable => + test(s"responses column schema depth boundary retains echoed metadata with durable=$durable") { + if (durable) delta() + 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") + if (durable) stage.setOutputPath(path()).setOutputBatchSize(4) + 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) + if (durable) { + verify(stage.resume(spark).collect()) + val status = AIPrompt.getOutputStatus(spark, stage.getOutputPath) + assert(status.isComplete && status.rowErrorCount == 1 && server.requests.size() == 3) + } + } + } + } + + replies.foreach { case (api, reply) => + + 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") + val fingerprints = Seq(names, names.reverse).map { 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) + val fingerprint = AIPromptOutput.fingerprint(stage, input.schema, outputSchema, spark) + assert(AIPromptOutput.fingerprint(loaded, input.schema, outputSchema, spark) == fingerprint) + 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) + fingerprint + } + if (!perRow) assert(fingerprints.distinct.size == 2) + assert(server.requests.size() == 2) + } + } + } + Seq(true).foreach { durable => + test(s"$api preserves response format descriptions and extensions with durable=$durable") { + if (durable) delta() + 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) + if (durable) stage.setOutputPath(path()).setOutputBatchSize(1) + 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)) + if (durable) { + assert(loaded.resume(spark).count() == 1 && server.requests.size() == 1) + assert(AIPrompt.getOutputStatus(spark, loaded.getOutputPath).allSucceeded) + if (!perRow) { + loaded.setResponseFormat(Map("type" -> "json_schema", + "json_schema" -> (details + ("description" -> "Different instructions")))) + intercept[IllegalArgumentException](loaded.resume(spark)) + assert(server.requests.size() == 1) + } + } + } + } + } + } + + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoverySuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoverySuite.scala new file mode 100644 index 00000000000..56b2d55cf23 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptRecoverySuite.scala @@ -0,0 +1,106 @@ +// 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 AIPromptRecoverySuite extends AIPromptOutputTestBase { + 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("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 original = stage.effectiveCopy() + val fingerprint = AIPromptOutput.fingerprint(original, input.schema, + original.transformSchema(input.schema), spark) + Seq(Right("absent-global-column"), Left("")).foreach { default => + GlobalParams.setGlobalParam(OpenAIDeploymentNameKey, default) + val effective = stage.effectiveCopy() + assert(AIPromptOutput.fingerprint(effective, input.schema, + effective.transformSchema(input.schema), spark) == fingerprint) + 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("JSON-valued body parameters preserve exact numbers and JVM arrays through save load and fingerprinting") { + 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(AIPromptOutput.fingerprint(stage, input.schema, output, spark) == + AIPromptOutput.fingerprint(loaded, input.schema, output, spark)) + 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) + } + } + } +} 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/AIPromptStorageCodecSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageCodecSuite.scala new file mode 100644 index 00000000000..ab648eec78e --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageCodecSuite.scala @@ -0,0 +1,366 @@ +// 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.ml.param.ParamMap +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.execution.FileSourceScanExec +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import spray.json._ + +import java.nio.file.{Files, Paths} +import java.util.Arrays +import scala.util.Using + +class AIPromptStorageCodecSuite extends AIPromptOutputTestBase { + test("durable schema planning rejects negative decimal scales at every nesting position") { + val setting = "spark.sql.legacy.allowNegativeScaleOfDecimal" + val previous = spark.conf.getOption(setting) + try { + spark.conf.set(setting, "true") + val decimal = DecimalType(8, -1) + val types = Seq(decimal, new StructType().add("amount", decimal), ArrayType(decimal), + MapType(StringType, decimal), MapType(decimal, StringType)) + types.foreach { dataType => + val root = path() + val stage = new AIPrompt().setOutputPath(root).setUrl("http://localhost/v1") + .setModel("model").setInputCol("text") + val schema = new StructType().add("text", StringType).add("retained", dataType) + val error = intercept[IllegalArgumentException](stage.transformSchema(schema)) + assert(error.getMessage.contains("nonnegative decimal scales")) + assert(!Files.exists(Paths.get(root))) + } + } finally previous.fold(spark.conf.unset(setting))(spark.conf.set(setting, _)) + } + + test("negative decimal scales stay valid lazily but fail durable transform before I/O") { + delta() + 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 root = path() + val stage = prompt(root, server) + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + val lazyRow = lazyStage.transform(data).head() + assert(lazyRow.getAs[String]("answer") == "answer") + assert(lazyRow.getAs[Any]("retained") == data.head().getAs[Any]("retained")) + val error = intercept[IllegalArgumentException](stage.transform(data)) + assert(error.getMessage.contains("nonnegative decimal scales")) + assert(!Files.exists(Paths.get(root)) && server.requests.get() == index + 1) + } + } + } finally previous.fold(spark.conf.unset(setting))(spark.conf.set(setting, _)) + } + + test("nonnegative decimal scales retain values through durable transform load and resume") { + delta() + val data = spark.range(1).select(lit("hello").as("text"), + lit(120).cast(DecimalType(8, 0)).as("whole"), + array(lit("12.34").cast(DecimalType(8, 2))).as("nested")) + withServer { server => + val root = path() + val stage = prompt(root, server) + Seq(stage.transform(data), AIPrompt.loadOutput(spark, root), stage.resume(spark)).foreach { output => + assert(output.schema == stage.transformSchema(data.schema)) + assert(output.select("whole", "nested").head() == data.select("whole", "nested").head()) + } + assert(server.requests.get() == 1 && AIPrompt.getOutputStatus(spark, root).allSucceeded) + } + } + + Seq("sql alias", "file metadata").foreach { source => + test(s"durable storage accepts Spark internal metadata from $source") { + delta() + val input = if (source == "sql alias") { + spark.sql("SELECT lower('HELLO')").toDF("text") + } else { + val parquet = path() + spark.range(1).select(lit("hello").as("text")).write.parquet(parquet) + spark.read.parquet(parquet).select(col("text"), col("_metadata")) + } + assert(input.schema.fields.exists(_.metadata != Metadata.empty)) + withServer { server => + val root = path() + val stage = prompt(root, server) + val transformed = stage.transform(input) + val loaded = AIPrompt.loadOutput(spark, root) + val resumed = stage.resume(spark) + Seq(transformed, loaded, resumed).foreach { output => + assert(output.schema == stage.transformSchema(input.schema)) + assert(output.select(input.columns.map(col): _*).collect().toSeq == input.collect().toSeq) + assert(output.head().getAs[String]("answer") == "answer") + } + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded && server.requests.get() == 1) + } + } + } + + test("durable schemas remove source-column behavior while preserving nested user metadata") { + delta() + val user = new MetadataBuilder().putString("description", "retained").build() + val generated = new MetadataBuilder().withMetadata(user) + .putString("delta.generationExpression", "lower('HELLO')").build() + val internal = new MetadataBuilder().withMetadata(user).putBoolean("__autoGeneratedAlias", true) + .putLong("delta.identity.start", 1L).putLong("delta.identity.step", 1L) + .putLong("delta.identity.highWaterMark", 2L).putBoolean("delta.identity.allowExplicitInsert", true).build() + val nested = StructType(Seq(StructField("value", StringType, nullable = true, internal))) + val schema = StructType(Seq(StructField("text", StringType, nullable = true, generated), + StructField("nested", ArrayType(nested), nullable = true, user), + StructField("mapped", MapType(StringType, nested), nullable = true, user))) + val input = spark.createDataFrame(Arrays.asList( + Row("hello", Seq(Row("item")), Map("key" -> Row("item"))), + Row("second", Seq(Row("other")), Map("key" -> Row("other")))), schema) + withServer { server => + val root = path() + val stage = prompt(root, server).setOutputBatchSize(1) + val expected = StructType(Seq(StructField("text", StringType, nullable = true, user), + StructField("nested", ArrayType(StructType(Seq(StructField("value", StringType, true, user)))), true, user), + StructField("mapped", MapType(StringType, StructType(Seq(StructField("value", StringType, true, user)))), + true, user))) + assert(AIPromptOutput.storageSchema(schema) == expected) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, input, event => + if (event == "batch_committed") throw new IllegalStateException("stop after first metadata batch")) + } + val partial = AIPrompt.getOutputStatus(spark, root) + assert(partial.committedCount == 1 && partial.pendingCount == 1 && server.requests.get() == 1) + assert(AIPrompt.loadOutput(spark, root).count() == 1 && server.requests.get() == 1) + AIPromptOutput.releaseOutputOwner(spark, root, partial.ownerToken.get, confirmStopped = true) + val recovery = spark.newSession() + Seq(stage.resume(recovery), AIPrompt.loadOutput(recovery, root), stage.resume(recovery)).foreach { output => + assert(output.schema == stage.transformSchema(schema)) + val restored = output.select(input.columns.map(col): _*) + assert(restored.schema == expected) + assert(restored.orderBy("text").collect().toSeq == input.orderBy("text").collect().toSeq) + val changed = new MetadataBuilder().putString("description", "modified").build() + val modified = restored.select(col("text").as("text", changed), col("nested"), col("mapped")) + intercept[IllegalArgumentException](AIPromptDelta.validateSchema(modified, expected)) + } + assert(server.requests.get() == 2 && AIPrompt.getOutputStatus(recovery, root).allSucceeded) + } + } + + test("file authorities reject before public durable I/O even in Spark local mode") { + delta() + withServer { server => + val root = path() + val suffix = new java.io.File(root).toURI.getPath + Seq("file://remote.invalid", "file://localhost").foreach { prefix => + val location = prefix + suffix + val qualified = AIPromptDelta.qualify(spark, location) + val direct = intercept[IllegalArgumentException](AIPromptDelta.requireRuntime(spark, qualified)) + assert(direct.getMessage.contains("same-host local directory")) + val stage = prompt(location, server) + val input = spark.range(1).select(lit("hello").as("text")) + val error = intercept[IllegalArgumentException](stage.transform(input)) + assert(error.getMessage.contains("same-host local directory")) + } + assert(!Files.exists(Paths.get(root)) && server.requests.get() == 0) + } + } + + Seq("retained", "input", "extraBody", "responseFormat").foreach { binding => + test(s"empty input names remain valid for lazy $binding bindings but not durable storage") { + 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)) + val root = path() + val durable = stage.copy(ParamMap.empty).setOutputPath(root) + assert(intercept[IllegalArgumentException](durable.transformSchema(input.schema)) + .getMessage.contains("Delta-compatible column names")) + if (AIPromptDeltaTestSession.hasDelta) { + assert(intercept[IllegalArgumentException](durable.transform(input)) + .getMessage.contains("Delta-compatible column names")) + } else { + assert(intercept[IllegalStateException](durable.transform(input)) + .getMessage.contains("requires a compatible, caller-configured Delta runtime")) + } + assert(!Files.exists(Paths.get(root)) && server.requests.size() == 1) + } + } + } + + test("vector map keys retain SQL identity through lazy and durable transform load and resume") { + delta() + withServer { server => + val root = path() + val stage = prompt(root, 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)) + def verify(output: DataFrame, expected: StructType): Unit = { + assert(output.schema == expected) + 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]) + assert(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)) + } + val lazyStage = stage.copy(ParamMap.empty) + lazyStage.clear(lazyStage.outputPath) + val lazyOutput = lazyStage.transform(input).cache() + try verify(lazyOutput, lazyStage.transformSchema(input.schema)) finally lazyOutput.unpersist() + val transformed = stage.transform(input) + Seq(transformed, AIPrompt.loadOutput(spark, root), stage.resume(spark)).foreach { output => + verify(output, stage.transformSchema(input.schema)) + } + val replacing = lazyStage.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() == 3) + } + } + + Seq(false, true).foreach { withVectors => + test(s"durable nested Unicode names follow SQL resolution with vectors=$withVectors") { + delta() + 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("\u03c3"), lit(2).as("\u03c2")) ++ (if (withVectors) + Seq(map_from_arrays(array(col("vector")), array(lit("weight"))).as("weights")) else Seq.empty) + val payload = struct(fields: _*) + val input = vectors.select(lit("hello").as("text"), payload.as("payload"), + array(payload).as("items"), map(lit("item"), payload).as("mapping")) + withServer { server => + val root = path() + val stage = prompt(root, server) + Seq[() => Any](() => stage.transformSchema(input.schema), () => stage.transform(input)).foreach { call => + assert(intercept[IllegalArgumentException](call()).getMessage.contains("Ambiguous durable fields")) + } + assert(!Files.exists(Paths.get(root)) && server.requests.get() == 0) + spark.conf.set("spark.sql.caseSensitive", "true") + val expected = stage.transformSchema(input.schema) + val transformed = stage.transform(input) + Seq(transformed, AIPrompt.loadOutput(spark, root), stage.resume(spark)).foreach { output => + assert(output.schema == expected) + 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[String]("answer") == "answer" && row.getAs[Row]("error") == null) + } + assert(server.requests.get() == 1) + } + } finally spark.conf.set("spark.sql.caseSensitive", caseSensitive) + } + } + + private def assertPruned(frame: DataFrame): Unit = { + val selected = frame.filter(col("id") === 1).select("id", "answer") + val scans = selected.queryExecution.sparkPlan.collect { case scan: FileSourceScanExec => scan } + assert(scans.size == 1) + assert(scans.head.requiredSchema.fieldNames.toSet == Set("id", "answer")) + assert(scans.head.dataFilters.flatMap(_.references.map(_.name)).contains("id")) + assert(selected.collect().toSeq == Seq(Row(1, "answer"))) + val counted = frame.agg(count(lit(1)).as("rows")) + val countScans = counted.queryExecution.sparkPlan.collect { case scan: FileSourceScanExec => scan } + assert(countScans.forall(_.requiredSchema.isEmpty)) + assert(counted.head().getLong(0) == 2L) + } + + Seq(false, true).foreach { withVectors => + test(s"durable returned output preserves scan pruning and filters with vectors=$withVectors") { + delta() + withServer { server => + val root = path() + val stage = prompt(root, server) + val metadata = new MetadataBuilder().putString("source", "retained").build() + val primitive = StructType(Seq(StructField("id", IntegerType, nullable = false), + StructField("text", StringType, nullable = false, metadata))) + val features = StructType(Seq( + StructField("vectors", ArrayType(SQLDataTypes.VectorType)), + StructField("weights", MapType(StringType, SQLDataTypes.VectorType)))) + val schema = if (withVectors) primitive.add("features", features, nullable = true, metadata) else primitive + val sparse = Vectors.sparse(4, Array(1, 3), Array(2.0, 0.0)) + val missing = Option.empty[Vector].orNull + val rows = if (withVectors) { + Arrays.asList( + Row(1, "first", Row(Seq(sparse, missing), Map("dense" -> Vectors.dense(1.0, 2.0)))), + Row(2, "second", Row(Option.empty[Seq[Vector]].orNull, Map("missing" -> missing)))) + } else Arrays.asList(Row(1, "first"), Row(2, "second")) + val input = spark.createDataFrame(rows, schema) + val transformed = stage.transform(input) + val loaded = AIPrompt.loadOutput(spark, root) + val resumed = stage.resume(spark) + Seq(transformed, loaded, resumed).foreach { output => + assert(output.schema == stage.transformSchema(input.schema)) + assert(output.schema("text").metadata == metadata) + assertPruned(output) + } + assertPruned(AIPromptStorageCodec.encode(loaded, loaded.schema)) + if (withVectors) { + assert(loaded.schema("features").metadata == metadata) + assert(loaded.select("id", "features").orderBy("id").collect().toSeq == + input.select("id", "features").orderBy("id").collect().toSeq) + } + assert(server.requests.get() == 2) + } + } + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageProcessSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageProcessSuite.scala new file mode 100644 index 00000000000..f0af265f7cc --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStorageProcessSuite.scala @@ -0,0 +1,457 @@ +// 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 org.apache.spark.sql.{DataFrame, Dataset, Encoders} +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} + +import java.io.File +import java.net.URLClassLoader +import java.nio.channels.{ClosedByInterruptException, FileLockInterruptionException} +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import scala.annotation.tailrec +import scala.collection.JavaConverters._ +import scala.util.control.NonFatal + +private[ai] object AIPromptStorageProcess { + def mark(file: File): Unit = { + Files.createDirectories(file.toPath.getParent) + Files.write(file.toPath, "ready".getBytes(StandardCharsets.UTF_8)) + } + + @tailrec + def awaitGate(gate: File): Unit = { + if (!gate.exists()) { + Thread.sleep(100) + awaitGate(gate) + } + } + + def main(args: Array[String]): Unit = args match { + case Array("hold-control", path, _, _, markerPath, gatePath) => + AIPromptLocalTransaction(new File(path).toURI.toString) { + mark(new File(markerPath)) + awaitGate(new File(gatePath)) + } + case _ => runSpark(args) + } + + private def runSpark(args: Array[String]): Unit = { + val Array(command, path, url, pauseAt, markerPath, gatePath) = args + AIPromptDeltaTestSession.resetSparkSession(numCores = Some(2)) + val spark = AIPromptDeltaTestSession.spark + val marker = new File(markerPath) + val stage = if (pauseAt == "after_response") { + new TestableAIPromptAfterResponsePause(markerPath, pause = command != "resume") + } else new AIPrompt() + val prompt = AIPromptStorageTestConfig.configure(stage, path, url) + def event(boundary: String): Unit = { + if (boundary == pauseAt) { + mark(marker) + if (gatePath.nonEmpty) awaitGate(new File(gatePath)) else new CountDownLatch(1).await() + } + } + try { + import spark.implicits._ + command match { + case "create" => + AIPromptOutput.transformObserved(prompt, Seq("duplicate", "duplicate", "ROW_ERROR").toDF("text"), event) + case "create-final" => + AIPromptOutput.transformObserved(prompt, Seq("duplicate", "duplicate").toDF("text"), event) + case "resume" => + prompt.setSubscriptionKey("rotated-synthetic-key") + AIPromptOutput.resumeObserved(prompt, spark, event) + } + mark(new File(markerPath + ".completed")) + spark.stop() + } catch { + case error: Throwable => + error.printStackTrace() + mark(new File(markerPath + ".failed")) + spark.stop() + System.exit(2) + } + } +} + +/** Pauses a Spark task after an actual HTTP response, without materializing inference twice. */ +private[ai] class TestableAIPromptAfterResponsePause(marker: String, pause: Boolean) extends AIPrompt { + override private[ai] def effectiveCopy(): AIPrompt = this + + override private[ai] def infer(dataset: Dataset[_], hiddenColumns: Seq[String], http: HTTPTransformer): DataFrame = { + val result = super.infer(dataset, hiddenColumns, http) + val markerPath = marker + if (!pause) result else result.map { row => + AIPromptStorageProcess.mark(new File(markerPath)) + new CountDownLatch(1).await() + row + }(Encoders.row(result.schema)) + } +} + +/** Opt-in because this suite starts and forcibly terminates multiple real Spark driver JVMs. */ +class AIPromptStorageProcessSuite extends AIPromptOutputTestBase { + private val driverWaitMinutes = 10 + private val processes = scala.collection.mutable.ArrayBuffer.empty[Process] + + private def enabled(): Unit = { + assume(sys.env.get("AIPROMPT_RUN_PROCESS_TESTS").contains("true"), + "Set AIPROMPT_RUN_PROCESS_TESTS=true to run real driver kill and acquisition race tests.") + delta() + } + + private def root(): File = Files.createTempDirectory(tmpDir, "aiprompt-process-").toFile + + private def classpath: String = { + @tailrec + def urls(loader: ClassLoader, found: Seq[String]): Seq[String] = { + Option(loader) match { + case None => found + case Some(current) => + val entries = current match { + case paths: URLClassLoader => paths.getURLs.toSeq.filter(_.getProtocol == "file") + .map(url => new File(url.toURI).getAbsolutePath) + case _ => Seq.empty + } + urls(current.getParent, found ++ entries) + } + } + sys.env.getOrElse("AIPROMPT_TEST_CLASSPATH", + (urls(getClass.getClassLoader, Seq.empty) ++ System.getProperty("java.class.path").split(File.pathSeparator)) + .distinct.mkString(File.pathSeparator)) + } + + private def launch(command: String, run: File, server: AIPromptRecoveryRecordingServer, boundary: String, + marker: File, gate: Option[File] = None): Process = { + val java = new File(System.getProperty("java.home"), "bin/java").getAbsolutePath + val scratch = new File(marker.getParentFile, marker.getName + "-jvm") + Files.createDirectories(scratch.toPath) + val classpathFile = new File(scratch, "classpath.args") + val classpathArguments = "-cp\n\"" + classpath.replace("\\", "\\\\").replace("\"", "\\\"") + "\"\n" + Files.write(classpathFile.toPath, classpathArguments.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(java, "-Xmx768m", "-Djava.io.tmpdir=" + scratch.getAbsolutePath, + "@" + classpathFile.getAbsolutePath, + "com.microsoft.azure.synapse.ml.services.ai.AIPromptStorageProcess", + command, run.getAbsolutePath, server.url, boundary, marker.getAbsolutePath, + gate.map(_.getAbsolutePath).getOrElse("")) + .redirectErrorStream(true).redirectOutput(new File(marker.getAbsolutePath + ".log")).start() + processes += process + process + } + + private def waitForMarker(process: Process, marker: File): Unit = { + val deadline = System.nanoTime() + TimeUnit.MINUTES.toNanos(driverWaitMinutes) + @tailrec + def waitUntilReady(): Unit = { + if (!marker.exists()) { + if (!process.isAlive) { + val output = Files.readAllLines(new File(marker.getAbsolutePath + ".log").toPath, StandardCharsets.UTF_8) + .asScala.takeRight(40).mkString("\n") + fail(s"Driver exited before reaching ${marker.getName}:\n$output") + } + if (System.nanoTime() >= deadline) { + stop(process) + fail(s"Driver did not reach ${marker.getName} within $driverWaitMinutes minutes.") + } + Thread.sleep(100) + waitUntilReady() + } + } + waitUntilReady() + } + + private def stop(process: Process): Unit = { + if (process.isAlive) { + val descendants = process.toHandle.descendants() + try descendants.iterator().asScala.foreach(_.destroyForcibly()) finally descendants.close() + process.destroyForcibly() + } + assert(process.waitFor(30, TimeUnit.SECONDS), "Test driver did not stop.") + } + + private def finish(process: Process): Int = { + if (!process.waitFor(driverWaitMinutes, TimeUnit.MINUTES)) { + stop(process) + fail("Test driver did not finish.") + } + process.exitValue() + } + + private def release(run: File): Unit = { + val state = AIPromptOutput.getOutputStatus(spark, run.getAbsolutePath) + AIPromptOutput.releaseOutputOwner(spark, run.getAbsolutePath, state.ownerToken.get, confirmStopped = true) + } + + private def cleanup(actions: Seq[() => Unit])(finish: => Unit): Unit = { + val failures = scala.collection.mutable.ArrayBuffer.empty[Throwable] + var interrupted = false + def attempt(action: () => Unit): Unit = { + try action() catch { + case error: InterruptedException => + interrupted = true + Thread.interrupted() + failures += error + case NonFatal(error) => failures += error + case error: Throwable => + failures += error + throw error + } + } + try { + try actions.foreach(attempt) finally attempt(() => finish) + failures.headOption.foreach(throw _) + } catch { + case error: Throwable => + failures.filterNot(_ eq error).foreach(error.addSuppressed) + throw error + } finally { + if (interrupted) Thread.currentThread().interrupt() + } + } + + override def afterAll(): Unit = + cleanup(processes.map(process => () => stop(process)).toSeq)(super.afterAll()) + + test("child cleanup attempts every driver and shared cleanup before reporting failures") { + Seq(new IllegalStateException("driver did not stop"), new InterruptedException("interrupted wait")).foreach { + first => + val second = new IllegalStateException("another driver did not stop") + val attempted = scala.collection.mutable.ArrayBuffer.empty[String] + val actions = Seq[() => Unit]( + () => { + attempted += "first" + if (first.isInstanceOf[InterruptedException]) Thread.currentThread().interrupt() + throw first + }, + () => { + attempted += "second" + assert(!Thread.currentThread().isInterrupted) + throw second + }, + () => { attempted += "third"; () }) + try { + val error = intercept[Exception] { + cleanup(actions) { + attempted += "shared" + assert(!Thread.currentThread().isInterrupted) + } + } + assert(attempted.toSeq == Seq("first", "second", "third", "shared")) + assert(error eq first) + assert(error.getSuppressed.toSeq == Seq(second)) + assert(Thread.currentThread().isInterrupted == first.isInstanceOf[InterruptedException]) + } finally { + Thread.interrupted() + } + } + } + + test("shared cleanup still runs when a fatal child cleanup error escapes") { + val failure = new LinkageError("child cleanup linkage failure") + var finished = false + val error = intercept[LinkageError] { + cleanup(Seq(() => throw failure)) { finished = true } + } + assert((error eq failure) && finished) + } + + test("fatal child cleanup retains prior and shared cleanup failures") { + Seq(new IllegalStateException("shared failure"), new InterruptedException("shared interruption"), + new LinkageError("shared linkage failure")).foreach { shared => + val prior = new IllegalStateException("prior driver failure") + val fatal = new LinkageError("driver linkage failure") + var laterDriverAttempted = false + try { + val error = intercept[Throwable] { + cleanup(Seq(() => throw prior, () => throw fatal, () => { laterDriverAttempted = true })) { + if (shared.isInstanceOf[InterruptedException]) Thread.currentThread().interrupt() + throw shared + } + } + val expected = if (shared.isInstanceOf[LinkageError]) shared else fatal + val suppressed = if (expected eq fatal) Seq(prior, shared) else Seq(prior, fatal) + assert(error eq expected) + assert(error.getSuppressed.toSeq == suppressed) + assert(!laterDriverAttempted) + assert(Thread.currentThread().isInterrupted == shared.isInstanceOf[InterruptedException]) + } finally { + Thread.interrupted() + } + } + } + + test("administrative release can cancel a separate driver's local control lock wait") { + enabled() + val session = spark + import session.implicits._ + withServer { server => + val directory = root() + val run = new File(directory, "run") + val stage = prompt(run.getAbsolutePath, server) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("one", "two", "three").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop before owner release")) + } + val before = AIPrompt.getOutputStatus(session, run.getAbsolutePath) + assert(before.committedCount == 2 && before.pendingCount == 1 && before.ownerToken.isDefined) + val gate = new File(directory, "release-lock") + val marker = new File(directory, "lock-held") + val holder = launch("hold-control", run, server, "", marker, Some(gate)) + val pool = Executors.newFixedThreadPool(2) + val caller = new AtomicReference[Thread]() + val failure = new AtomicReference[Exception]() + val interrupted = new AtomicBoolean() + try { + waitForMarker(holder, marker) + val waiting = pool.submit(new Callable[Unit] { + override def call(): Unit = { + caller.set(Thread.currentThread()) + try { + AIPrompt.releaseOutputOwner(session, run.getAbsolutePath, before.ownerToken.get, confirmStopped = true) + } catch { + case error: Exception => failure.set(error) + } finally interrupted.set(Thread.currentThread().isInterrupted) + } + }) + eventually(timeout(Span(60, Seconds))) { + val stack = Option(caller.get()).toSeq.flatMap(_.getStackTrace) + assert(stack.exists(_.getClassName == AIPromptLocalTransaction.getClass.getName)) + assert(stack.exists(frame => + (frame.getClassName == "sun.nio.ch.FileChannelImpl" && frame.getMethodName == "lock") || + (frame.getClassName == "java.lang.Thread" && frame.getMethodName == "sleep"))) + } + val cancel = pool.submit(new Callable[Unit] { + override def call(): Unit = caller.get().interrupt() + }) + cancel.get(10, TimeUnit.SECONDS) + waiting.get(10, TimeUnit.SECONDS) + assert(holder.isAlive && !gate.exists()) + assert(failure.get().isInstanceOf[InterruptedException] || + failure.get().isInstanceOf[FileLockInterruptionException] || + failure.get().isInstanceOf[ClosedByInterruptException]) + assert(interrupted.get()) + assert(AIPrompt.getOutputStatus(session, run.getAbsolutePath) == before) + assert(server.requests.get() == 2) + AIPromptStorageProcess.mark(gate) + assert(finish(holder) == 0) + AIPrompt.releaseOutputOwner(session, run.getAbsolutePath, before.ownerToken.get, confirmStopped = true) + assert(stage.resume(session).count() == 3 && server.requests.get() == 3) + val complete = AIPrompt.getOutputStatus(session, run.getAbsolutePath) + assert(complete.allSucceeded && complete.ownerToken.isEmpty) + } finally { + cleanup(Seq( + () => { AIPromptStorageProcess.mark(gate); stop(holder) }, + () => { + pool.shutdownNow() + assert(pool.awaitTermination(30, TimeUnit.SECONDS)) + })) { () } + } + } + } + + test("driver death at every initialization write never fabricates a resumable run") { + enabled() + val server = new AIPromptRecoveryRecordingServer + try { + Seq("control_created", "input_committed", "results_created").foreach { boundary => + val directory = root() + val run = new File(directory, "run") + val marker = new File(directory, "paused") + val process = launch("create", run, server, boundary, marker) + waitForMarker(process, marker) + stop(process) + val error = intercept[IllegalArgumentException] { + AIPromptOutput.loadOutput(spark, run.getAbsolutePath) + } + assert(error.getMessage.contains("initialization is incomplete")) + } + assert(server.requests.get() == 0) + } finally server.close() + } + + test("killed driver recovery skips committed batches and final lost acknowledgements") { + enabled() + val server = new AIPromptRecoveryRecordingServer + try { + Seq("create", "create-final").foreach { command => + val directory = root() + val run = new File(directory, "run") + val marker = new File(directory, "paused") + val before = server.requests.get() + val process = launch(command, run, server, "batch_committed", marker) + waitForMarker(process, marker) + stop(process) + val state = AIPromptOutput.getOutputStatus(spark, run.getAbsolutePath) + assert(state.committedCount == 2 && state.state != "complete") + assert(server.requests.get() == before + 2) + assert(AIPromptOutput.loadOutput(spark, run.getAbsolutePath).count() == 2) + release(run) + val resumed = launch("resume", run, server, "", new File(directory, "resumed")) + assert(finish(resumed) == 0) + val total = if (command == "create") 3 else 2 + assert(AIPromptOutput.loadOutput(spark, run.getAbsolutePath).count() == total) + assert(server.requests.get() == before + total) + } + } finally server.close() + } + + test("driver death after paid responses and before commit permits replay of only uncommitted input") { + enabled() + val server = new AIPromptRecoveryRecordingServer + try { + val directory = root() + val run = new File(directory, "run") + val marker = new File(directory, "response") + val process = launch("create", run, server, "after_response", marker) + waitForMarker(process, marker) + stop(process) + assert(server.requests.get() > 0) + assert(AIPromptOutput.getOutputStatus(spark, run.getAbsolutePath).committedCount == 0) + release(run) + val paid = server.requests.get() + assert(finish(launch("resume", run, server, "after_response", new File(directory, "resumed"))) == 0) + assert(server.requests.get() == paid + 3) + val status = AIPrompt.getOutputStatus(spark, run.getAbsolutePath) + assert(status.isComplete && status.committedCount == 3 && status.ownerToken.isEmpty) + } finally server.close() + } + + test("separate drivers racing first creation and later acquisition have one inferencing winner") { + enabled() + val server = new AIPromptRecoveryRecordingServer + try { + Seq("create", "resume").foreach { command => + val directory = root() + val run = new File(directory, "run") + if (command == "resume") { + val ready = new File(directory, "ready") + val initializer = launch("create", run, server, "ready", ready) + waitForMarker(initializer, ready) + stop(initializer) + release(run) + } + val gate = new File(directory, "gate") + val left = new File(directory, "left") + val right = new File(directory, "right") + val boundary = if (command == "create") "before_control_create" else "before_acquire" + val one = launch(command, run, server, boundary, left, Some(gate)) + val two = launch(command, run, server, boundary, right, Some(gate)) + waitForMarker(one, left) + waitForMarker(two, right) + val before = server.requests.get() + AIPromptStorageProcess.mark(gate) + val outcomes = Seq(finish(one), finish(two)) + assert(outcomes.count(_ == 0) == 1 && outcomes.count(_ != 0) == 1) + assert(server.requests.get() == before + 3) + val status = AIPromptOutput.getOutputStatus(spark, run.getAbsolutePath) + assert(status.isComplete && status.ownerToken.isEmpty) + } + } finally server.close() + } +} diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStoredIdentitySuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStoredIdentitySuite.scala new file mode 100644 index 00000000000..0b5e85ff31a --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptStoredIdentitySuite.scala @@ -0,0 +1,390 @@ +// 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.build.BuildInfo +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, OpenAIPrompt} +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.internal.SQLConfInjections +import org.apache.spark.sql.types._ +import spray.json.{JsArray, JsNumber, JsObject, JsString} + +import java.io.File +import java.util.Collections +import scala.collection.JavaConverters._ +import scala.collection.immutable.ListMap +import scala.util.Using + +class AIPromptStoredIdentitySuite extends AIPromptOutputTestBase { + test("rendered JSON number kinds define public requests and durable identities") { + 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")) + def configured(value: Either[String, Double]): AIPrompt = { + val stage = new AIPrompt("rendered-number").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))) + } + def fingerprint(stage: AIPrompt): String = + AIPromptOutput.fingerprint(stage, data.schema, stage.transformSchema(data.schema), spark) + 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 = configured(value) + assert(fingerprint(stage) == fingerprint(configured(Left(rendered)))) + assert(fingerprint(stage) != fingerprint(configured(Left(rendered + ".0")))) + assert(stage.transform(data).head().getAs[String]("answer") == "answer") + val request = server.requests.poll()._2 + assert(AIJson.objectString(request).fields("x").compactPrint == rendered) + } + assert(server.requests.isEmpty) + } + } + + test("equivalent JSON numeric spellings retain exact durable fingerprints") { + val data = spark.range(1).select(lit("hello").as("text")) + def configured(name: String, number: String): AIPrompt = { + val value = ListMap[String, Any]("number" -> BigDecimal(number), + "nested" -> Seq(ListMap("number" -> BigDecimal(number)))) + val stage = new AIPrompt("numeric-identity").setUrl("http://localhost/v1").setModel("model") + .setInputCol("text").setOutputCol("answer").setErrorCol("error") + name match { + case "extraBody" => stage.setExtraBody(value) + case "responseFormat" => stage.setResponseSchema(ListMap("type" -> "object", + "properties" -> ListMap("value" -> ListMap("type" -> "number")), "values" -> value)) + case "questions" => stage.setApiType("decisions").setQuestions(Seq(ListMap( + "type" -> "predicate", "name" -> "check", "instructions" -> "Check", "values" -> value))) + } + } + def fingerprint(stage: AIPrompt): String = + AIPromptOutput.fingerprint(stage, data.schema, stage.transformSchema(data.schema), spark) + val equivalent = Seq("2.50" -> "2.5", "1.00" -> "1.0", "1e2" -> "100.0", "-0.00" -> "0.0", + "-2.50" -> "-2.5", "0.1234567890123456789012345678900" -> "0.12345678901234567890123456789", + "1.00e400" -> "1e400", "1.00e-400" -> "1e-400", "100e2147483645" -> "1000e2147483644") + Seq("extraBody", "responseFormat", "questions").foreach { name => + equivalent.foreach { case (original, rebound) => + assert(fingerprint(configured(name, original)) == fingerprint(configured(name, rebound)), + s"$name: $original and $rebound must have the same recovery identity") + } + assert(fingerprint(configured(name, "0.1")) != + fingerprint(configured(name, "0.10000000000000000000000000001"))) + Seq("0", "1", "32", "-32").foreach { integer => + assert(fingerprint(configured(name, integer)) != fingerprint(configured(name, integer + ".0")), + s"$name: integer $integer and decimal $integer.0 must have distinct recovery identities") + } + } + } + + test("durable base64 chunking identity rejects changes between batches and across recovery sessions") { + delta() + val session = spark.newSession() + val setting = "spark.sql.chunkBase64String.enabled" + session.conf.set(setting, "true") + val data = session.range(3).withColumn("text", lit("long input " * 12)) + val expression = "base64(encode(text, 'UTF-8'))" + val original = data.selectExpr(expression).head().getString(0) + session.conf.set(setting, "false") + val changed = data.selectExpr(expression).head().getString(0) + assert(original.contains("\r\n") && original.replace("\r\n", "") == changed.replace("\r\n", "")) + // Earlier Spark patches always wrap; newer runtimes consult this registered setting. + if (SQLConfInjections.registeredKeys.contains(setting)) assert(original != changed) + session.conf.set(setting, "true") + val reply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + val stage = AIPromptStorageTestConfig.configure(new AIPrompt(), root, server.url + "/v1") + .setPromptTemplate("{" + expression + "}").setOutputBatchSize(1) + def contents: Vector[String] = server.requests.asScala.toVector.map { case (_, body) => + AIJson.objectString(body).fields("messages").asInstanceOf[JsArray].elements.last + .asJsObject.fields("content").asInstanceOf[JsString].value + } + val mismatch = intercept[IllegalArgumentException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "batch_committed") session.conf.set(setting, "false")) + } + assert(mismatch.getMessage.contains("Spark SQL semantics")) + val status = AIPrompt.getOutputStatus(session, root) + assert(status.committedCount == 1 && status.pendingCount == 2 && status.ownerToken.isEmpty) + assert(contents == Vector(original)) + val resumed = spark.newSession() + resumed.conf.set(setting, "false") + val version = AIPromptDelta.version(resumed, root, "_control") + assert(intercept[IllegalArgumentException](stage.resume(resumed)).getMessage.contains("Spark SQL semantics")) + assert(contents == Vector(original) && AIPrompt.getOutputStatus(resumed, root).ownerToken.isEmpty) + assert(AIPromptDelta.version(resumed, root, "_control") == version) + resumed.conf.set(setting, "true") + assert(stage.resume(resumed).count() == 3 && contents == Vector.fill(3)(original)) + assert(AIPrompt.getOutputStatus(resumed, root).allSucceeded) + } + } + + test("corrupt required control identities reject before public durable reads or owner changes") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("hello").toDF("text"), event => + if (event == "ready") throw new IllegalStateException("stop before requests")) + } + val owner = AIPrompt.getOutputStatus(spark, root).ownerToken.get + val table = AIPromptDelta.table(root, "_control") + val actions: Seq[() => Any] = Seq( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true)) + Seq("runId", "keyColumn", "fingerprint", "libraryVersion", "sparkVersion", "bindings", + "inputSchema", "outputSchema").foreach { column => + val original = spark.sql(s"SELECT $column FROM $table").head().getString(0) + Seq("NULL", "''").foreach { invalid => + spark.sql(s"UPDATE $table SET $column = $invalid") + val version = AIPromptDelta.version(spark, root, "_control") + val checks = if (column == "runId") actions else Seq(() => AIPromptDelta.control(spark, root)) + checks.foreach { check => + val error = intercept[IllegalArgumentException](check()) + assert(error.getMessage.contains("Corrupt AIPrompt control table") && error.getCause == null) + } + assert(AIPromptDelta.version(spark, root, "_control") == version && server.requests.get() == 0) + } + spark.sql(s"UPDATE $table SET $column = ${AIPromptDelta.literal(original)}") + } + Seq("", " ", "\t").foreach { invalid => + spark.sql(s"UPDATE $table SET owner = ${AIPromptDelta.literal(invalid)}") + val version = AIPromptDelta.version(spark, root, "_control") + actions.foreach { check => + val error = intercept[IllegalArgumentException](check()) + assert(error.getMessage.contains("Corrupt AIPrompt control table") && error.getCause == null) + } + assert(AIPromptDelta.version(spark, root, "_control") == version && server.requests.get() == 0) + } + spark.sql(s"UPDATE $table SET owner = ${AIPromptDelta.literal(owner)}") + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.contains(owner)) + AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true) + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + assert(stage.resume(spark).count() == 1 && server.requests.get() == 1) + } + } + + test("durable recovery rejects missing scalar authentication before ownership and permits credential rotation") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server).setTelemHeaders(Map.empty[String, String]) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("one", "two", "three").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after first batch")) + } + val status = AIPrompt.getOutputStatus(spark, root) + assert(status.committedCount == 2 && server.requests.get() == 2) + AIPrompt.releaseOutputOwner(spark, root, status.ownerToken.get, confirmStopped = true) + val version = AIPromptDelta.version(spark, root, "_control") + val missing = prompt(root, server).setTelemHeaders(Map.empty[String, String]) + missing.clear(missing.subscriptionKey) + Seq(false, true).foreach { blank => + if (blank) missing.setSubscriptionKey(" ").setAADToken(" ").setCustomAuthHeader(" ") + .setCustomHeaders(Map("Authorization" -> " ")).setTelemHeaders(Map("api-key" -> "ignored")) + val error = intercept[IllegalArgumentException] { + AIPromptOutput.resumeObserved(missing, spark, event => fail(s"Missing auth acquired ownership: $event")) + } + assert(error.getMessage.contains("configuration")) + assert(AIPromptDelta.version(spark, root, "_control") == version && server.requests.get() == 2) + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + } + missing.setAADToken("rotated-synthetic-token") + assert(missing.resume(spark).count() == 3 && server.requests.get() == 3) + assert(AIPrompt.getOutputStatus(spark, root).allSucceeded) + } + } + + test("durable authentication presence follows scalar precedence inherited headers and implicit Fabric routing") { + val schema = new StructType().add("text", StringType) + val stages = Seq( + new AIPrompt("auth-presence") -> "api-key", + new TestableAIHeaderPrompt("auth-presence") -> "x-model-key", + new TestableAISharedHeaderPrompt("auth-presence") -> "authorization", + new TestableAIFabricAuthPrompt("auth-presence") -> "api-key") + stages.foreach { case (source, authName) => + val base = source.setPromptTemplate("{text}").setTelemHeaders(Map.empty[String, String]).effectiveCopy() + def fingerprint(stage: AIPrompt): String = AIPromptOutput.fingerprint(stage, schema, schema, spark) + val configured = fingerprint(base.copy(org.apache.spark.ml.param.ParamMap.empty) + .setSubscriptionKey("synthetic-key")) + val setters: Seq[AIPrompt => Unit] = Seq( + _.setAADToken("synthetic-token"), _.setCustomAuthHeader("Custom synthetic-token"), + _.setCustomHeaders(Map(authName.toUpperCase(java.util.Locale.ROOT) -> "synthetic-key")), + _.setSubscriptionKey(" ").setAADToken("synthetic-token")) + setters.foreach { set => + val rotated = base.copy(org.apache.spark.ml.param.ParamMap.empty) + set(rotated) + assert(fingerprint(rotated) == configured) + } + val implicitFabric = source.isInstanceOf[TestableAIFabricAuthPrompt] + assert((fingerprint(base) == configured) == implicitFabric) + val blank = base.copy(org.apache.spark.ml.param.ParamMap.empty).setSubscriptionKey(" ") + .setAADToken(" ").setCustomAuthHeader(" ").setCustomHeaders(Map(authName -> " ")) + .setTelemHeaders(Map(authName -> "ignored")) + assert(fingerprint(blank) == fingerprint(base)) + } + } + + Seq("systemone", "decisions").foreach { api => + Seq(Map.empty[String, String], Map("same" -> "telemetry")).foreach { telemetry => + test(s"$api class identity is independent of telemetry $telemetry for pending and completed recovery") { + delta() + import spark.implicits._ + val reply = if (api == "systemone") NativeRequestFixtures.SystemResponse + else NativeRequestFixtures.DecisionResponse + val questions = if (api == "systemone") NativeRequestFixtures.SystemQuestions + else NativeRequestFixtures.DecisionQuestions + scala.util.Using.resource(new AIRecordingServer((_, _) => 200 -> reply)) { server => + val root = path() + def configure(stage: AIPrompt): AIPrompt = stage.setApiType(api).setTelemHeaders(telemetry) + .setUrl(server.url).setModel("model").setQuestions(questions).setInputCol("text") + .setOutputCol("answer").setErrorCol("error").setOutputPath(root).setOutputBatchSize(1) + def direct(): AIPrompt = configure(if (api == "systemone") new AISystemOne() else new OpenAIDecisions()) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(direct(), Seq("one", "two").toDF("text"), event => + if (event == "batch_committed") throw new IllegalStateException("stop after first outcome")) + } + val owner = AIPrompt.getOutputStatus(spark, root).ownerToken.get + AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true) + Seq(false, true).foreach { complete => + val version = AIPromptDelta.version(spark, root, "_control") + val mismatch = intercept[IllegalArgumentException](configure(new AIPrompt()).resume(spark)) + assert(mismatch.getMessage.contains("stage class")) + assert(AIPromptDelta.version(spark, root, "_control") == version) + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + assert(server.requests.size() == (if (complete) 2 else 1)) + assert(direct().resume(spark).count() == 2 && server.requests.size() == 2) + } + } + } + } + } + + Seq(false, true).foreach { complete => + test(s"resume rejects changed stored build or Spark identity before ownership when complete is $complete") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + val data = Seq("hello").toDF("text") + if (complete) stage.transform(data) else { + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, data, event => + if (event == "ready") throw new IllegalStateException("stop before requests")) + } + val owner = AIPrompt.getOutputStatus(spark, root).ownerToken.get + AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true) + } + val table = AIPromptDelta.table(root, "_control") + val original = spark.sql(s"SELECT * FROM $table").head() + val versions = Seq("libraryVersion" -> BuildInfo.version, "sparkVersion" -> spark.version) + versions.foreach { case (column, expected) => + assert(original.getAs[String](column) == expected) + spark.sql(s"UPDATE $table SET $column = 'incompatible-version'") + val version = AIPromptDelta.version(spark, root, "_control") + val error = intercept[IllegalArgumentException](stage.resume(spark)) + assert(error.getMessage.contains("build or Spark runtime") && error.getCause == null) + assert(AIPromptDelta.version(spark, root, "_control") == version) + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.isEmpty) + assert(AIPrompt.loadOutput(spark, root).count() == (if (complete) 1 else 0)) + assert(server.requests.get() == (if (complete) 1 else 0)) + spark.sql(s"UPDATE $table SET $column = ${AIPromptDelta.literal(expected)}") + } + assert(stage.resume(spark).count() == 1 && server.requests.get() == 1) + } + } + } + + private val malformedSchemas = Seq( + "non-struct" -> "\"string\"", + "field-name" -> """{"struct":[{"name":5,"type":"string","metadata":{}}]}""", + "field-value" -> """{"struct":[null]}""", + "metadata" -> """{"struct":[{"name":"text","type":"string","metadata":{"secret":[1,"private"]}}]}""", + "metadata-array" -> """{"struct":[{"name":"text","type":"string","metadata":{"secret":[null]}}]}""", + "metadata-number" -> """{"struct":[{"name":"text","type":"string","metadata":{"n":9223372036854775808}}]}""", + "decimal" -> """{"struct":[{"name":"text","type":{"precision":1.5,"scale":0},"metadata":{}}]}""", + "decimal-precision" -> """{"struct":[{"name":"text","type":{"precision":99,"scale":0},"metadata":{}}]}""", + "decimal-scale" -> """{"struct":[{"name":"text","type":{"precision":1,"scale":2},"metadata":{}}]}""", + "deep" -> ("{\"array\":" * 110 + "\"string\"" + "}" * 110)) + + malformedSchemas.foreach { case (name, json) => + test(s"malformed stored $name schema fails with a controlled argument error") { + val error = intercept[IllegalArgumentException](AIPromptStorageCodec.readSchema(json)) + assert(error.getCause == null && !error.getMessage.contains("private")) + } + } + + test("corrupt schema rejects all public durable reads and owner release without writes or sends") { + delta() + import spark.implicits._ + withServer { server => + val root = path() + val stage = prompt(root, server) + intercept[IllegalStateException] { + AIPromptOutput.transformObserved(stage, Seq("hello").toDF("text"), event => + if (event == "ready") throw new IllegalStateException("stop before requests")) + } + val owner = AIPrompt.getOutputStatus(spark, root).ownerToken.get + val table = AIPromptDelta.table(root, "_control") + Seq("inputSchema", "outputSchema").foreach { column => + val original = spark.sql(s"SELECT $column FROM $table").head().getString(0) + malformedSchemas.foreach { case (_, json) => + spark.sql(s"UPDATE $table SET $column = ${AIPromptDelta.literal(json)}") + val version = AIPromptDelta.version(spark, root, "_control") + val actions: Seq[() => Any] = Seq( + () => AIPrompt.getOutputStatus(spark, root), + () => AIPrompt.loadOutput(spark, root), + () => stage.resume(spark), + () => AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true)) + actions.foreach { action => + val error = intercept[IllegalArgumentException](action()) + assert(error.getCause == null && !error.getMessage.contains("private")) + } + assert(AIPromptDelta.version(spark, root, "_control") == version && server.requests.get() == 0) + } + spark.sql(s"UPDATE $table SET $column = ${AIPromptDelta.literal(original)}") + } + assert(AIPrompt.getOutputStatus(spark, root).ownerToken.contains(owner)) + AIPrompt.releaseOutputOwner(spark, root, owner, confirmStopped = true) + assert(stage.resume(spark).count() == 1 && server.requests.get() == 1) + } + } + + test("schema depth limits reject before durable initialization and supported metadata round trips") { + delta() + withServer { server => + def nestedStruct(levels: Int): DataType = + (1 to levels).foldLeft[DataType](StringType)((value, _) => new StructType().add("child", value)) + val accepted = new StructType().add("text", StringType).add("nested", nestedStruct(32)) + assert(AIPromptStorageCodec.readSchema(AIPromptStorageCodec.schemaJson(accepted)) == + AIPromptOutput.storageSchema(accepted)) + val metadata = new MetadataBuilder().putLongArray("values", Array(1L, 2L)) + .putMetadata("nested", new MetadataBuilder().putString("label", "value").build()).build() + val leaf = StructType(Seq(StructField("child", StringType, nullable = false, metadata))) + val withMetadata = (1 until 32).foldLeft[DataType](leaf)((value, _) => + new StructType().add("child", value)) + val arrays = (1 to 110).foldLeft[DataType](StringType)((value, _) => ArrayType(value)) + Seq(arrays, nestedStruct(40), withMetadata).foreach { nested => + val root = path() + val schema = new StructType().add("text", StringType).add("nested", nested) + val data = spark.createDataFrame(Collections.emptyList[Row](), schema) + val error = intercept[IllegalArgumentException](prompt(root, server).transform(data)) + assert(error.getMessage.contains("AIPrompt durable schema") && error.getCause == null) + assert(!new File(root).exists() && server.requests.get() == 0) + } + val supported = StructType(Seq(StructField("text", StringType, nullable = false, metadata))) + assert(AIPromptStorageCodec.readSchema(AIPromptStorageCodec.schemaJson(supported)) == + AIPromptOutput.storageSchema(supported)) + } + } +} 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..c3fc9f40732 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIPromptTestFixtures.scala @@ -0,0 +1,146 @@ +// 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.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(in.readAllBytes(), 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/AIRecoveryJsonSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIRecoveryJsonSuite.scala new file mode 100644 index 00000000000..5ae6123b893 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIRecoveryJsonSuite.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.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 + +class AIRecoveryJsonSuite extends AIPromptOutputTestBase { + private val missing: Null = None.orNull + private val chatReply = """{"choices":[{"index":0,"message":{"role":"assistant","content":"answer"}}]}""" + test("JSON-valued defaults preserve exact numbers explicitness and durable fingerprints 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) + val output = stage.transformSchema(data.schema) + assert(AIPromptOutput.fingerprint(stage, data.schema, output, spark) == + AIPromptOutput.fingerprint(loaded, data.schema, output, spark)) + 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 HTTP and durable fingerprints") { + 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") + val fingerprints = Seq(names, names.reverse).map { 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) + val output = stage.transformSchema(data.schema) + val fingerprint = AIPromptOutput.fingerprint(stage, data.schema, output, spark) + assert(AIPromptOutput.fingerprint(loaded, data.schema, output, spark) == fingerprint) + 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) + fingerprint + } finally { + org.apache.commons.io.FileUtils.deleteDirectory(new java.io.File(directory).getParentFile) + } + } + assert(fingerprints.distinct.size == 2) + } + } +} 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..2eafae68ba6 --- /dev/null +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/ai/AIValueCodecSuite.scala @@ -0,0 +1,221 @@ +// 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("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 48f7f3bfca6..436d81b57bd 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 @@ -8,19 +8,63 @@ import com.microsoft.azure.synapse.ml.io.http.ErrorUtils import com.sun.net.httpserver.{HttpExchange, HttpHandler, HttpServer} import org.apache.commons.io.FileUtils import org.apache.commons.io.IOUtils +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{FSDataInputStream, FileSystem, Path, RawLocalFileSystem} +import org.apache.spark.{TaskContext, TaskKilledException} import org.apache.spark.ml.param.ParamMap import org.apache.spark.sql.Row -import org.apache.spark.sql.functions.{col, struct, to_json} +import org.apache.spark.sql.functions.{col, lit, struct, to_json} import org.apache.spark.sql.types.{ArrayType, StringType, StructField, StructType} +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.{Seconds, Span} import spray.json._ -import java.io.File -import java.net.InetSocketAddress +import java.io.{File, IOException} +import java.net.{InetSocketAddress, URI} +import java.nio.channels.ClosedByInterruptException import java.nio.charset.StandardCharsets import java.nio.file.Files -import java.util.concurrent.ConcurrentLinkedQueue -import java.util.Locale +import java.util.concurrent.{Callable, ConcurrentHashMap, ConcurrentLinkedQueue, CountDownLatch, + ExecutionException, Executors, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import java.util.{Locale, UUID} import scala.collection.JavaConverters._ +import scala.util.Using + +private[openai] object AttachmentReadFailures { + class Gate { + val entered = new CountDownLatch(1) + val release = new CountDownLatch(1) + val completed = new CountDownLatch(1) + val context = new AtomicReference[TaskContext]() + val taskFailure = new AtomicReference[Throwable]() + val completionInterrupted = new AtomicBoolean() + } + + val Active = new ConcurrentHashMap[String, Gate]() + + def fail(id: String): Nothing = { + Option(Active.get(id)).foreach { gate => + val context = TaskContext.get() + gate.context.set(context) + context.addTaskFailureListener((_: TaskContext, error: Throwable) => gate.taskFailure.set(error)) + context.addTaskCompletionListener[Unit] { _ => + gate.completionInterrupted.set(Thread.currentThread().isInterrupted) + gate.completed.countDown() + } + gate.entered.countDown() + require(gate.release.await(30, TimeUnit.SECONDS), "Attachment cancellation test gate timed out") + } + if (id == "interrupted.png") throw new ClosedByInterruptException() + throw new IOException("Synthetic attachment read failure") + } +} + +class AttachmentReadFailureFileSystem extends RawLocalFileSystem { + override def getUri: URI = URI.create("attachmentfailure:///") + + override def open(path: Path, bufferSize: Int): FSDataInputStream = AttachmentReadFailures.fail(path.getName) +} class OpenAIPromptMultimodalRequestSuite extends TestBase { @@ -390,6 +434,76 @@ class OpenAIPromptMultimodalRequestSuite extends TestBase { } } + test("attachment file reads propagate cancellation without converting ordinary IO failures") { + val location = URI.create("attachmentfailure:///") + val conf = new Configuration() + conf.setClass("fs.attachmentfailure.impl", classOf[AttachmentReadFailureFileSystem], classOf[FileSystem]) + Using.resource(FileSystem.get(location, conf)) { _ => + def read(name: String): Unit = { + OpenAIAttachmentUtils.prepareFile(location.resolve(name).toString, + None, Set("png"), Set.empty, Set.empty) + () + } + assert(!Thread.currentThread().isInterrupted) + try { + val error = intercept[InterruptedException](read("interrupted.png")) + assert(Thread.currentThread().isInterrupted) + assert(error.getMessage == "Attachment read interrupted" && error.getCause == null) + } finally Thread.interrupted() + val error = intercept[IllegalArgumentException](read("ordinary.png")) + assert(error.getMessage == "Unable to read attachment 'ordinary.png'" && error.getCause == null) + assert(!Thread.currentThread().isInterrupted) + } + } + + Seq("chat_completions", "responses").foreach { api => + test(s"legacy $api attachment task cancellation escapes the UDF without interrupting cleanup") { + val context = spark.sparkContext + val group = UUID.randomUUID().toString + val gate = new AttachmentReadFailures.Gate() + val location = URI.create(s"attachmentfailure:///$group") + val conf = new Configuration() + conf.setClass("fs.attachmentfailure.impl", classOf[AttachmentReadFailureFileSystem], classOf[FileSystem]) + AttachmentReadFailures.Active.put(group, gate) + Using.resource(FileSystem.get(location, conf)) { _ => + withEchoServer { (baseUrl, bodies) => + val data = spark.range(1).select(lit("Describe this image.").as("prompt"), + lit(location.toString).as("image")) + val result = prompt(baseUrl, api).transform(data) + val executor = Executors.newSingleThreadExecutor() + val job = executor.submit(new Callable[Unit] { + override def call(): Unit = { + context.setJobGroup(group, "Legacy attachment task cancellation", interruptOnCancel = false) + try result.collect() finally context.clearJobGroup() + } + }) + try { + assert(gate.entered.await(30, TimeUnit.SECONDS)) + context.cancelJobGroup(group) + eventually(timeout(Span(10, Seconds))) { + assert(gate.context.get().isInterrupted()) + } + gate.release.countDown() + assert(gate.completed.await(10, TimeUnit.SECONDS)) + intercept[ExecutionException](job.get(10, TimeUnit.SECONDS)) + assert(!gate.completionInterrupted.get()) + val killed = Iterator.iterate(gate.taskFailure.get())(_.getCause).takeWhile(_ != null) + .collectFirst { case error: TaskKilledException => error } + assert(killed.exists(_.reason == "Attachment read canceled")) + assert(bodies.isEmpty) + } finally { + gate.release.countDown() + context.cancelJobGroup(group) + job.cancel(true) + executor.shutdownNow() + assert(executor.awaitTermination(15, TimeUnit.SECONDS)) + AttachmentReadFailures.Active.remove(group) + } + } + } + } + } + test("OpenAIPrompt preserves existing row errors over attachment failures") { assert(!spark.conf.get("spark.sql.caseSensitive").toBoolean) val existingError = Row("upstream error", null) // scalastyle:ignore null 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/PythonMethodExtractor.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/PythonMethodExtractor.scala new file mode 100644 index 00000000000..a8eec9b42a3 --- /dev/null +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/codegen/PythonMethodExtractor.scala @@ -0,0 +1,112 @@ +// 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.codegen + +import scala.annotation.tailrec + +private[codegen] object PythonMethodExtractor { + private val MethodPrefix = """^def\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(""" + private val MethodStart = MethodPrefix.r + val MethodPattern = (MethodPrefix + """([^)]*)\)(?:\s*->\s*([^:]+))?:\s*$""").r + private val ScopePrefix = """^(?:(?:async\s+)?def|class)\s+""".r + + case class Method(name: String, signature: String, decorators: Seq[String]) + + // Class statements and nested scope bodies can contain multiline expressions and strings. + private class ExpressionState { + private var quote = "" + private var delimiterDepth = 0 + private var explicitContinuation = false + + def continues: Boolean = quote.nonEmpty || delimiterDepth > 0 || explicitContinuation + + private def afterStringCharacter(line: String, index: Int): Int = { + if (line.charAt(index) == '\\') { + index + 2 + } else if (line.startsWith(quote, index)) { + val nextIndex = index + quote.length + quote = "" + nextIndex + } else { + index + 1 + } + } + + @tailrec + private def appendCode(line: String, index: Int, code: StringBuilder): Unit = { + if (index < line.length) { + val char = line.charAt(index) + val nextIndex = if (quote.nonEmpty) { + afterStringCharacter(line, index) + } else if (char == '#') { + line.length + } else if (char == '\'' || char == '"') { + val tripleQuote = char.toString * 3 + quote = if (line.startsWith(tripleQuote, index)) tripleQuote else char.toString + index + quote.length + } else { + if ("([{".contains(char)) delimiterDepth += 1 + else if (")]}".contains(char)) delimiterDepth -= 1 + code.append(char) + index + 1 + } + appendCode(line, nextIndex, code) + } + } + + def scan(line: String): String = { + val code = new StringBuilder + appendCode(line, 0, code) + val statement = code.toString.trim + explicitContinuation = statement.endsWith("\\") + statement + } + } + + def extract(source: String): Seq[Method] = { + var nestedScopeIndent: Option[Int] = None + val expression = new ExpressionState + var decoratorIndent: Option[Int] = None + var decorators = Vector.empty[String] + source.split("\n").toSeq.flatMap { line => + val statement = line.trim + val indent = line.prefixLength(_.isWhitespace) + if (expression.continues) { + // A continuation can close at the declaration's indentation without ending its scope. + expression.scan(line) + None + } else if (statement.isEmpty || statement.startsWith("#")) { + None + } else { + nestedScopeIndent = nestedScopeIndent.filter(indent > _) + if (nestedScopeIndent.isDefined) { + expression.scan(line) + None + } else if (statement.startsWith("@")) { + if (!decoratorIndent.contains(indent)) decorators = Vector.empty + decoratorIndent = Some(indent) + decorators :+= expression.scan(line) + None + } else { + val method = MethodStart.findFirstMatchIn(statement).map { matched => + val bindingDecorators = if (decoratorIndent.contains(indent)) { + decorators.filter(Set("@staticmethod", "@classmethod")) + } else { + Vector.empty + } + Method(matched.group(1), statement, bindingDecorators) + } + decorators = Vector.empty + decoratorIndent = None + // Control blocks do not introduce a Python scope; functions and classes do. + if (ScopePrefix.findFirstIn(statement).isDefined) { + nestedScopeIndent = Some(indent) + } + expression.scan(line) + method + } + } + } + } +} 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 3dc1c4b4b0b..6919fc56d9b 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) @@ -56,9 +54,6 @@ trait BaseWrappable extends Params { } private[codegen] object PythonStubCodegen { - val PythonMethodPattern = - """^\s*def\s+([A-Za-z_][A-Za-z0-9_]*)\(([^)]*)\)(?:\s*->\s*([^:]+))?:\s*$""".r - private val TypeInfoByStage = new WeakHashMap[Params, Map[String, PythonTypeInfo]]() def pythonTypeInfo(stage: Params, param: Param[_]): PythonTypeInfo = synchronized { @@ -76,11 +71,13 @@ private[codegen] object PythonStubCodegen { trait PythonWrappable extends BaseWrappable { import GenerationUtils._ - import PythonStubCodegen.PythonMethodPattern + import PythonMethodExtractor.{Method, MethodPattern} private def getPythonTypeInfo(p: Param[_]): PythonTypeInfo = PythonStubCodegen.pythonTypeInfo(thisStage, p) + protected def pyParamInputType(p: Param[_]): Option[String] = None + def pyAdditionalMethods: String = { "" } @@ -129,7 +126,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: @@ -161,15 +158,14 @@ trait PythonWrappable extends BaseWrappable { method.getName == name && method.getParameterCount == parameterCount) private def validateServiceParamAliases(): Unit = { val paramNames = thisStage.params.map(_.name).toSet - val additionalMethods = pyAdditionalMethods + val additionalMethods = PythonMethodExtractor.extract(pyAdditionalMethods).map(_.name).toSet thisStage.params.collect { case p: ServiceParam[_] => p }.foreach { p => val capName = p.name.capitalize val alias = s"${p.name}Col" require(!paramNames.contains(alias), s"Service parameter ${p.name} cannot use Python alias $alias because that Param already exists") Seq(s"get${capName}Col", s"set${capName}Col").foreach { methodName => - val methodPattern = s"(?m)^\\s*def\\s+$methodName\\s*\\(".r - require(methodPattern.findFirstIn(additionalMethods).isEmpty, + require(!additionalMethods(methodName), s"Generated Python method $methodName conflicts with a hand-written method") } } @@ -348,7 +344,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( @@ -397,7 +393,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: ... @@ -471,8 +467,8 @@ trait PythonWrappable extends BaseWrappable { } private def pyStubAdditionalMethods: String = { - pyAdditionalMethods.split("\n").flatMap { - case PythonMethodPattern(methodName, args, returnType) if !methodName.startsWith("_") => + PythonMethodExtractor.extract(pyAdditionalMethods).flatMap { + case Method(_, MethodPattern(methodName, args, returnType), decorators) if !methodName.startsWith("_") => val fluent = isFluentStubMethod(methodName) val methodArgs = args.split(",").filter(_.trim.nonEmpty) val valueArgCount = methodArgs.count(arg => !Set("self", "cls").contains(arg.trim)) @@ -484,7 +480,7 @@ trait PythonWrappable extends BaseWrappable { pyStubAdditionalArgument(argument, fluent, inferredType) }.mkString(", ") val typedReturn = pyStubAdditionalReturnType(methodName, returnType, fluent) - Some(s"def $methodName($typedArgs) -> $typedReturn: ...") + Some((decorators :+ s"def $methodName($typedArgs) -> $typedReturn: ...").mkString("\n")) case _ => None }.distinct.mkString("\n") @@ -616,7 +612,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..ad5820326e9 --- /dev/null +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/CancellableRequests.scala @@ -0,0 +1,243 @@ +// 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.internal.Logging + +import java.util.concurrent.{Callable, CancellationException, ConcurrentHashMap, ExecutionException, Executors, Future, + FutureTask, 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 + }) + 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() + if (context.exists(_.isInterrupted())) { + throw new TaskKilledException("HTTP request task canceled") + } + if (cancelled.get()) { + throw new CancellationException("HTTP request canceled") + } + } + + def submit[T](action: => T): Future[T] = { + checkTask() + pool.submit(new Callable[T] { + override def call(): T = { + Active.set(Scope.this) + try { + check() + action + } finally Active.remove() + } + }) + } + + @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 => 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.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 c0e0cacf56a..1c1dc1341cd 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,42 +60,94 @@ 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 = { - val rData = new HTTPResponseData(response) - response.close() - rData + try { + 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) + } + } finally response.close() } 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 = { @@ -100,6 +156,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], @@ -121,7 +224,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 @@ -194,87 +301,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 @@ -428,6 +604,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/com/microsoft/azure/synapse/ml/param/GlobalParams.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/GlobalParams.scala index 8d888df85d9..eae653f1e26 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/GlobalParams.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/GlobalParams.scala @@ -17,23 +17,23 @@ object GlobalParams { key.asInstanceOf[GlobalKey[Any]] } - def setGlobalParam[T](key: GlobalKey[T], value: T): Unit = { + def setGlobalParam[T](key: GlobalKey[T], value: T): Unit = synchronized { GlobalParams(untypedKey(key)) = value } - def getGlobalParam[T](key: GlobalKey[T]): Option[T] = { + def getGlobalParam[T](key: GlobalKey[T]): Option[T] = synchronized { GlobalParams.get(untypedKey(key)).map(_.asInstanceOf[T]) } - def resetGlobalParam[T](key: GlobalKey[T]): Unit = { + def resetGlobalParam[T](key: GlobalKey[T]): Unit = synchronized { GlobalParams -= untypedKey(key) } - def getParam[T](p: Param[T]): Option[T] = { + def getParam[T](p: Param[T]): Option[T] = synchronized { ParamToKeyMap.get(p).flatMap(GlobalParams.get).map(_.asInstanceOf[T]) } - def registerParam[T](p: Param[T], key: GlobalKey[T]): Unit = { + def registerParam[T](p: Param[T], key: GlobalKey[T]): Unit = synchronized { ParamToKeyMap(p) = untypedKey(key) } } 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 72a8c82ce94..0447e4251c5 100644 --- a/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala +++ b/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala @@ -37,7 +37,7 @@ private[ml] class ComplexParamsWriter(instance: Params) extends MLWriter { val complexParamLocs = ComplexParamsWriter.getComplexParamLocations(instance, path) val complexParamJson = ComplexParamsWriter.getComplexMetadata(complexParamLocs) ComplexParamsWriter.saveMetadata(instance, path, sparkSession, complexParamJson) - ComplexParamsWriter.saveComplexParams(path, complexParamLocs, shouldOverwrite) + ComplexParamsWriter.saveComplexParams(path, complexParamLocs, shouldOverwrite, sparkSession) } } @@ -65,7 +65,11 @@ private[ml] object ComplexParamsWriter { def saveComplexParams(basePath: String, complexParamLocs: Map[ParamPair[_], Path], shouldOverwrite: Boolean): Unit = { - val spark = SparkSession.builder().getOrCreate() + saveComplexParams(basePath, complexParamLocs, shouldOverwrite, SparkSession.builder().getOrCreate()) + } + + def saveComplexParams(basePath: String, complexParamLocs: Map[ParamPair[_], Path], + shouldOverwrite: Boolean, spark: SparkSession): Unit = { complexParamLocs.foreach { case (ParamPair(p: ComplexParam[Any], v), loc) => p.save(v, spark, new Path(basePath, loc), shouldOverwrite) @@ -106,6 +110,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 @@ -115,9 +127,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) ~ @@ -159,7 +171,7 @@ private[ml] class ComplexParamsReader[T] extends MLReader[T] { val instance = cls.getConstructor(classOf[String]).newInstance(metadata.uid).asInstanceOf[Params] metadata.getAndSetParams(instance) - ComplexParamsReader.getAndSetComplexParams(instance, metadata, path) + ComplexParamsReader.getAndSetComplexParams(instance, metadata, path, sparkSession) instance.asInstanceOf[T] } } @@ -167,7 +179,11 @@ private[ml] class ComplexParamsReader[T] extends MLReader[T] { private[ml] object ComplexParamsReader { def getAndSetComplexParams(instance: Params, metadata: Metadata, basePath: String): Unit = { - val spark = SparkSession.builder().getOrCreate() + getAndSetComplexParams(instance, metadata, basePath, SparkSession.builder().getOrCreate()) + } + + def getAndSetComplexParams(instance: Params, metadata: Metadata, + basePath: String, spark: SparkSession): Unit = { implicit val format: DefaultFormats.type = DefaultFormats val complexParamLocs = (metadata.metadata \ "complexParamLocs") match { case JNothing => 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..292b8e8f6ff --- /dev/null +++ b/core/src/main/scala/org/apache/spark/sql/internal/SQLConfInjections.scala @@ -0,0 +1,17 @@ +// 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 + +import scala.collection.JavaConverters._ + +object SQLConfInjections { + // Enumerating values would also evaluate unrelated optional Hive defaults. + def registeredKeys: Seq[String] = SQLConf.getConfigEntries().asScala.iterator.map(_.key).toVector + + 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/codegen/PyCodegenFixtures.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenFixtures.scala new file mode 100644 index 00000000000..05b7578549a --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenFixtures.scala @@ -0,0 +1,368 @@ +// 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.codegen + +import org.apache.spark.ml.classification.{ProbabilisticClassificationModel, ProbabilisticClassifier} +import org.apache.spark.ml.linalg.{Vector, Vectors} +import org.apache.spark.ml.{Estimator, Model, Transformer} +import org.apache.spark.ml.param.{DoubleParam, Param, ParamMap, StringArrayParam} +import org.apache.spark.sql.{DataFrame, Dataset} +import org.apache.spark.sql.types.StructType +import spray.json.DefaultJsonProtocol._ + +// Nesting does not hide these JVM classes; discovery must exclude their test artifact. +private[codegen] object PyCodegenFixtures { + + class TypedPythonStage(override val uid: String = "typedPythonStage") + extends Transformer with Wrappable { + + override protected lazy val classNameHelper: String = "TypedPythonStage" + + val count = new DoubleParam(this, "count", "numeric value") + val labels = new StringArrayParam(this, "labels", "labels") + val model = new com.microsoft.azure.synapse.ml.param.ServiceParam[String]( + this, "model", "model name") + val text = new Param[String](this, "text", "text value") + + def getModel: String = "" + def setAlias(value: String): this.type = this + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|def setAlias(self, value): + | return self + | + |def clear(self, param): + | return None + | + |def copy(self, extra=None): + | return self + | + |def _identity(method): + | return method + | + |@staticmethod + |# Keep the binding decorator across another decorator and a blank line. + |@_identity + | + |def loadOutput(spark, outputPath): + | return None + | + |@_identity + |@classmethod + | + |# Comments do not end a decorator block. + |def fromPath(cls, path): + | return cls() + | + |def describe(self): + | def check_keys(item): + | return item + | def getModelCol(): + | return check_keys("") + | return getModelCol() + |""".stripMargin + + override def transform(dataset: Dataset[_]): DataFrame = dataset.toDF() + + override def transformSchema(schema: StructType): StructType = schema + + override def copy(extra: ParamMap): TypedPythonStage = defaultCopy(extra) + } + + class ForeignParamPythonStage extends TypedPythonStage("foreignParamPythonStage") { + + override protected lazy val classNameHelper: String = "ForeignParamPythonStage" + + override val text = new Param[String]("otherStage", "text", "text value") + } + + class ConflictingServiceAliasStage(override val uid: String = "conflictingServiceAliasStage") + extends TypedPythonStage(uid) { + + override protected lazy val classNameHelper: String = "ConflictingServiceAliasStage" + + val modelCol = new Param[String](this, "modelCol", "conflicting real parameter") + } + + class ConflictingServiceAliasMethodStage(override val uid: String = "conflictingServiceAliasMethodStage") + extends TypedPythonStage(uid) { + + override protected lazy val classNameHelper: String = "ConflictingServiceAliasMethodStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|def getModelCol(self): + | return "model" + |""".stripMargin + } + + class ConditionalPythonStage extends TypedPythonStage("conditionalPythonStage") { + + override protected lazy val classNameHelper: String = "ConditionalPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|if True: + | @staticmethod + | # A conditional method is still a member of the generated class. + | def conditionalValue(value): + | def hiddenConditional(): + | return value + | return hiddenConditional() + | + | class NestedHelper: + | def getModelCol(self): + | return "nested" + | def nestedOnly(self): + | return None + | + |try: + | def tryValue(self): + | return None + |except RuntimeError: + | def fallbackValue(self): + | return None + |""".stripMargin + } + + class ConflictingConditionalAliasMethodStage extends TypedPythonStage("conflictingConditionalAliasMethodStage") { + + override protected lazy val classNameHelper: String = "ConflictingConditionalAliasMethodStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|if True: + | try: + | def getModelCol(self): + | return "model" + | except RuntimeError: + | pass + |""".stripMargin + } + + class MultilineFunctionPythonStage extends TypedPythonStage("multilineFunctionPythonStage") { + + override protected lazy val classNameHelper: String = "MultilineFunctionPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|def describe( + | self, + | label=")", + | # A closing delimiter in a comment does not finish this header: ) + |): + | def getModelCol(): + | return label + | return getModelCol() + | + |def afterScope(self): + | return None + |""".stripMargin + } + + class MultilineClassPythonStage extends TypedPythonStage("multilineClassPythonStage") { + + override protected lazy val classNameHelper: String = "MultilineClassPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|class NestedHelper( + | object, + |): + | def setModelCol(self, value): + | return value + | def nestedOnly(self): + | return None + | + |def afterScope(self): + | return None + |""".stripMargin + } + + class MultilineStringPythonStage(quote: String) extends TypedPythonStage("multilineStringPythonStage") { + + override protected lazy val classNameHelper: String = "MultilineStringPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + s"""|def describe( + | self, + | label=$quote( + | ' " [ { # still inside the string + |# A string line can end its own literal: $quote, + |): + | def getModelCol(): + | return label + | return getModelCol() + | + |def afterScope(self): + | return None + |""".stripMargin + } + + class ConflictingMultilineStringAliasMethodStage(quote: String) extends MultilineStringPythonStage(quote) { + + override protected lazy val classNameHelper: String = "ConflictingMultilineStringAliasMethodStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|def getModelCol(self): + | return "model" + |""".stripMargin + } + + class BodyStringPythonStage(scopeHeader: String) extends TypedPythonStage("bodyStringPythonStage") { + + override protected lazy val classNameHelper: String = "BodyStringPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + s"""|$scopeHeader + | label = ''' + |This is string content, not a dedent. + |''' + | def getModelCol(): + | return "nested" + | def nestedOnly(): + | return None + | + |def afterScope(self): + | return None + |""".stripMargin + } + + class ConflictingBodyStringAliasMethodStage(scopeHeader: String) + extends TypedPythonStage("conflictingBodyStringAliasMethodStage") { + + override protected lazy val classNameHelper: String = "ConflictingBodyStringAliasMethodStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + s"""|$scopeHeader + | label = ''' + |def notAHeader( + |''' + | + |def getModelCol(self): + | return "model" + |""".stripMargin + } + + class ClassStatementStringPythonStage(controlBlock: Boolean, methodName: String) + extends TypedPythonStage("classStatementStringPythonStage") { + + override protected lazy val classNameHelper: String = "ClassStatementStringPythonStage" + + override def pyAdditionalMethods: String = { + val prefix = if (controlBlock) "if True:\n " else "" + super.pyAdditionalMethods + + s"""|${prefix}label = ''' + |def notAHeader( + |''' + | + |def $methodName(self): + | return None + |""".stripMargin + } + } + + class DecoratedPythonStage extends TypedPythonStage("decoratedPythonStage") { + + override protected lazy val classNameHelper: String = "DecoratedPythonStage" + + override def pyAdditionalMethods: String = + super.pyAdditionalMethods + + """|def _with_label(label): + | return lambda method: method + | + |@staticmethod # no instance binding + |def commentedStatic(value): + | return value + | + |@classmethod # bind to the class + |def commentedClass(cls, value): + | return cls() + | + |@staticmethod + |@_with_label( + | {"labels": ["(", ")"]}, # ) + |) + |def decoratedStatic(value): + | return value + | + |@classmethod + |@_with_label( + | '''( + |# still inside the label: ''', + |) + |def decoratedClass(cls, value): + | return cls() + | + |@staticmethod + |@_with_label \ + | ("label") + |def continuedStatic(value): + | return value + | + |def plainValue(self): + | return None + |""".stripMargin + } + + class TypedPythonModel(override val uid: String = "typedPythonModel") + extends Model[TypedPythonModel] with Wrappable { + + override protected lazy val classNameHelper: String = "TypedPythonModel" + + override def transform(dataset: Dataset[_]): DataFrame = dataset.toDF() + + override def transformSchema(schema: StructType): StructType = schema + + override def copy(extra: ParamMap): TypedPythonModel = defaultCopy(extra) + } + + class TypedPythonEstimator(override val uid: String = "typedPythonEstimator") + extends Estimator[TypedPythonModel] with Wrappable { + + override protected lazy val classNameHelper: String = "TypedPythonEstimator" + + override protected def companionModelClassName: String = + "com.microsoft.azure.synapse.ml.codegen.TypedPythonModel" + + override def fit(dataset: Dataset[_]): TypedPythonModel = new TypedPythonModel() + + override def transformSchema(schema: StructType): StructType = schema + + override def copy(extra: ParamMap): TypedPythonEstimator = defaultCopy(extra) + } + + class TypedPythonClassificationModel(override val uid: String = "typedPythonClassificationModel") + extends ProbabilisticClassificationModel[Vector, TypedPythonClassificationModel] with Wrappable { + + override protected lazy val classNameHelper: String = "TypedPythonClassificationModel" + + override def numClasses: Int = 2 + + override def predictRaw(features: Vector): Vector = Vectors.dense(0.0, 1.0) + + override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = rawPrediction + + override def copy(extra: ParamMap): TypedPythonClassificationModel = defaultCopy(extra) + } + + class TypedPythonClassifier(override val uid: String = "typedPythonClassifier") + extends ProbabilisticClassifier[Vector, TypedPythonClassifier, TypedPythonClassificationModel] with Wrappable { + + override protected lazy val classNameHelper: String = "TypedPythonClassifier" + + override protected def companionModelClassName: String = + "com.microsoft.azure.synapse.ml.codegen.TypedPythonClassificationModel" + + override protected def train(dataset: Dataset[_]): TypedPythonClassificationModel = + new TypedPythonClassificationModel() + + override def copy(extra: ParamMap): TypedPythonClassifier = defaultCopy(extra) + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenSuite.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenSuite.scala index 90f46e7cac2..b2da525d60f 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenSuite.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/codegen/PyCodegenSuite.scala @@ -4,14 +4,7 @@ package com.microsoft.azure.synapse.ml.codegen import org.apache.commons.io.{FileUtils, IOUtils} -import org.apache.spark.ml.classification.{ProbabilisticClassificationModel, ProbabilisticClassifier} -import org.apache.spark.ml.linalg.{Vector, Vectors} -import org.apache.spark.ml.{Estimator, Model, Transformer} -import org.apache.spark.ml.param.{DoubleParam, Param, ParamMap, StringArrayParam} -import org.apache.spark.sql.{DataFrame, Dataset} -import org.apache.spark.sql.types.StructType import org.scalatest.funsuite.AnyFunSuite -import spray.json.DefaultJsonProtocol._ import java.io.File import java.lang.reflect.Modifier @@ -19,125 +12,6 @@ import java.nio.charset.StandardCharsets import java.nio.file.Files import java.util.zip.ZipFile -// Nesting does not hide these JVM classes; discovery must exclude their test artifact. -private[codegen] object PyCodegenFixtures { - - class TypedPythonStage(override val uid: String = "typedPythonStage") - extends Transformer with Wrappable { - - override protected lazy val classNameHelper: String = "TypedPythonStage" - - val count = new DoubleParam(this, "count", "numeric value") - val labels = new StringArrayParam(this, "labels", "labels") - val model = new com.microsoft.azure.synapse.ml.param.ServiceParam[String]( - this, "model", "model name") - val text = new Param[String](this, "text", "text value") - - def getModel: String = "" - def setAlias(value: String): this.type = this - - override def pyAdditionalMethods: String = - super.pyAdditionalMethods + - """|def setAlias(self, value): - | return self - | - |def clear(self, param): - | return None - | - |def copy(self, extra=None): - | return self - |""".stripMargin - - override def transform(dataset: Dataset[_]): DataFrame = dataset.toDF() - - override def transformSchema(schema: StructType): StructType = schema - - override def copy(extra: ParamMap): TypedPythonStage = defaultCopy(extra) - } - - class ForeignParamPythonStage extends TypedPythonStage("foreignParamPythonStage") { - - override protected lazy val classNameHelper: String = "ForeignParamPythonStage" - - override val text = new Param[String]("otherStage", "text", "text value") - } - - class ConflictingServiceAliasStage(override val uid: String = "conflictingServiceAliasStage") - extends TypedPythonStage(uid) { - - override protected lazy val classNameHelper: String = "ConflictingServiceAliasStage" - - val modelCol = new Param[String](this, "modelCol", "conflicting real parameter") - } - - class ConflictingServiceAliasMethodStage(override val uid: String = "conflictingServiceAliasMethodStage") - extends TypedPythonStage(uid) { - - override protected lazy val classNameHelper: String = "ConflictingServiceAliasMethodStage" - - override def pyAdditionalMethods: String = - super.pyAdditionalMethods + - """|def getModelCol(self): - | return "model" - |""".stripMargin - } - - class TypedPythonModel(override val uid: String = "typedPythonModel") - extends Model[TypedPythonModel] with Wrappable { - - override protected lazy val classNameHelper: String = "TypedPythonModel" - - override def transform(dataset: Dataset[_]): DataFrame = dataset.toDF() - - override def transformSchema(schema: StructType): StructType = schema - - override def copy(extra: ParamMap): TypedPythonModel = defaultCopy(extra) - } - - class TypedPythonEstimator(override val uid: String = "typedPythonEstimator") - extends Estimator[TypedPythonModel] with Wrappable { - - override protected lazy val classNameHelper: String = "TypedPythonEstimator" - - override protected def companionModelClassName: String = - "com.microsoft.azure.synapse.ml.codegen.TypedPythonModel" - - override def fit(dataset: Dataset[_]): TypedPythonModel = new TypedPythonModel() - - override def transformSchema(schema: StructType): StructType = schema - - override def copy(extra: ParamMap): TypedPythonEstimator = defaultCopy(extra) - } - - class TypedPythonClassificationModel(override val uid: String = "typedPythonClassificationModel") - extends ProbabilisticClassificationModel[Vector, TypedPythonClassificationModel] with Wrappable { - - override protected lazy val classNameHelper: String = "TypedPythonClassificationModel" - - override def numClasses: Int = 2 - - override def predictRaw(features: Vector): Vector = Vectors.dense(0.0, 1.0) - - override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = rawPrediction - - override def copy(extra: ParamMap): TypedPythonClassificationModel = defaultCopy(extra) - } - - class TypedPythonClassifier(override val uid: String = "typedPythonClassifier") - extends ProbabilisticClassifier[Vector, TypedPythonClassifier, TypedPythonClassificationModel] with Wrappable { - - override protected lazy val classNameHelper: String = "TypedPythonClassifier" - - override protected def companionModelClassName: String = - "com.microsoft.azure.synapse.ml.codegen.TypedPythonClassificationModel" - - override protected def train(dataset: Dataset[_]): TypedPythonClassificationModel = - new TypedPythonClassificationModel() - - override def copy(extra: ParamMap): TypedPythonClassifier = defaultCopy(extra) - } -} - class PyCodegenSuite extends AnyFunSuite { import PyCodegenFixtures._ @@ -231,6 +105,18 @@ class PyCodegenSuite extends AnyFunSuite { assert(process.waitFor() === 0, output) } + private def assertAliasConflict(root: File, stage: TypedPythonStage): Unit = { + val sourceFile = new File(root, "fixture.py") + writeUtf8(sourceFile, "class Fixture:\n" + + stage.pyAdditionalMethods.split("\n").map(" " + _).mkString("\n") + "\n") + assertPythonCompiles(sourceFile) + val error = intercept[IllegalArgumentException] { + stage.makePyFile(codegenConfig(root)) + } + assert(error.getMessage.contains( + "Generated Python method getModelCol conflicts with a hand-written method")) + } + private def wheelEntryContent(wheel: File, path: String): Option[String] = { val archive = new ZipFile(wheel) try { @@ -339,6 +225,164 @@ class PyCodegenSuite extends AnyFunSuite { } } + test("generated additional method stubs retain binding across comments, blanks, and stacked decorators") { + withTempDir { root => + val conf = codegenConfig(root) + new TypedPythonStage().makePyFile(conf) + val stubFile = new File(packageDir(conf.pySrcDir, "/codegen"), "TypedPythonStage.pyi") + val stub = readUtf8(stubFile) + assert(stub.contains("@staticmethod\n def loadOutput(spark: Any, outputPath: Any) -> Any: ...")) + assert(stub.contains("@classmethod\n def fromPath(cls, path: Any) -> Any: ...")) + assert(!stub.contains("@classmethod\n def describe")) + assert(stub.contains("def describe(self) -> Any: ...")) + assert(!stub.contains("def check_keys(")) + assert(!stub.contains("def _identity(")) + assert(occurrences(stub, "def getModelCol(") == 1) + assertPythonCompiles(stubFile) + } + } + + test("class-body control blocks expose methods but not nested function or class members") { + withTempDir { root => + val conf = codegenConfig(root) + new ConditionalPythonStage().makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "ConditionalPythonStage.pyi") + val stub = readUtf8(stubFile) + assert(stub.contains("@staticmethod\n def conditionalValue(value: Any) -> Any: ...")) + assert(stub.contains("def tryValue(self) -> Any: ...")) + assert(stub.contains("def fallbackValue(self) -> Any: ...")) + assert(!stub.contains("def hiddenConditional(")) + assert(!stub.contains("def nestedOnly(")) + assert(occurrences(stub, "def getModelCol(") == 1) + assertPythonCompiles(new File(folder, "ConditionalPythonStage.py")) + assertPythonCompiles(stubFile) + } + } + + Seq(new MultilineFunctionPythonStage, new MultilineClassPythonStage).foreach { stage => + val className = stage.getClass.getSimpleName + test(s"multiline scope headers exclude nested members from $className") { + withTempDir { root => + val conf = codegenConfig(root) + stage.makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, s"$className.pyi") + val stub = readUtf8(stubFile) + assert(stub.contains("def afterScope(self) -> Any: ...")) + assert(!stub.contains("def nestedOnly(")) + assert(occurrences(stub, "def getModelCol(") == 1) + assert(occurrences(stub, "def setModelCol(") == 1) + assertPythonCompiles(new File(folder, s"$className.py")) + assertPythonCompiles(stubFile) + } + } + } + + Seq("'''", "\"\"\"").foreach { quote => + test(s"multiline string headers retain subsequent class methods with $quote") { + withTempDir { root => + val conf = codegenConfig(root) + new MultilineStringPythonStage(quote).makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "MultilineStringPythonStage.pyi") + assertPythonCompiles(new File(folder, "MultilineStringPythonStage.py")) + assertPythonCompiles(stubFile) + val stub = readUtf8(stubFile) + assert(stub.contains("def afterScope(self) -> Any: ...")) + assert(occurrences(stub, "def getModelCol(") == 1) + } + } + + test(s"multiline string headers detect real accessor collisions with $quote") { + withTempDir { root => + assertAliasConflict(root, new ConflictingMultilineStringAliasMethodStage(quote)) + } + } + } + + test("binding decorator inline comments preserve static and class method stubs") { + withTempDir { root => + val conf = codegenConfig(root) + new DecoratedPythonStage().makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "DecoratedPythonStage.pyi") + assertPythonCompiles(new File(folder, "DecoratedPythonStage.py")) + assertPythonCompiles(stubFile) + val stub = readUtf8(stubFile) + assert(stub.contains("@staticmethod\n def commentedStatic(value: Any) -> Any: ...")) + assert(stub.contains("@classmethod\n def commentedClass(cls, value: Any) -> Any: ...")) + } + } + + Seq("def bodyString(self):", "class BodyString:").foreach { scopeHeader => + test(s"body strings do not expose nested accessors in $scopeHeader") { + withTempDir { root => + val conf = codegenConfig(root) + val stage = new BodyStringPythonStage(scopeHeader) + val sourceFile = new File(root, "fixture.py") + writeUtf8(sourceFile, "class Fixture:\n" + + stage.pyAdditionalMethods.split("\n").map(" " + _).mkString("\n") + "\n") + assertPythonCompiles(sourceFile) + stage.makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "BodyStringPythonStage.pyi") + val stub = readUtf8(stubFile) + assert(stub.contains("def afterScope(self) -> Any: ...")) + assert(!stub.contains("def nestedOnly(")) + assert(occurrences(stub, "def getModelCol(") == 1) + assertPythonCompiles(new File(folder, "BodyStringPythonStage.py")) + assertPythonCompiles(stubFile) + } + } + + test(s"body strings do not hide real accessor collisions after $scopeHeader") { + withTempDir { root => + assertAliasConflict(root, new ConflictingBodyStringAliasMethodStage(scopeHeader)) + } + } + } + + Seq(false, true).foreach { controlBlock => + test(s"class statement strings retain subsequent methods with control block $controlBlock") { + withTempDir { root => + val conf = codegenConfig(root) + new ClassStatementStringPythonStage(controlBlock, "afterString").makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "ClassStatementStringPythonStage.pyi") + assertPythonCompiles(new File(folder, "ClassStatementStringPythonStage.py")) + assertPythonCompiles(stubFile) + val stub = readUtf8(stubFile) + assert(stub.contains("def afterString(self) -> Any: ...")) + assert(!stub.contains("def notAHeader(")) + } + } + + test(s"class statement strings preserve alias checks with control block $controlBlock") { + withTempDir { root => + assertAliasConflict(root, new ClassStatementStringPythonStage(controlBlock, "getModelCol")) + } + } + } + + test("binding decorators survive multiline decorator expressions") { + withTempDir { root => + val conf = codegenConfig(root) + new DecoratedPythonStage().makePyFile(conf) + val folder = packageDir(conf.pySrcDir, "/codegen") + val stubFile = new File(folder, "DecoratedPythonStage.pyi") + assertPythonCompiles(new File(folder, "DecoratedPythonStage.py")) + assertPythonCompiles(stubFile) + val stub = readUtf8(stubFile) + assert(stub.contains("@staticmethod\n def decoratedStatic(value: Any) -> Any: ...")) + assert(stub.contains("@classmethod\n def decoratedClass(cls, value: Any) -> Any: ...")) + assert(stub.contains("@staticmethod\n def continuedStatic(value: Any) -> Any: ...")) + assert(stub.contains("def plainValue(self) -> Any: ...")) + assert(!stub.contains("@staticmethod\n def plainValue")) + assert(!stub.contains("@classmethod\n def plainValue")) + } + } + test("generated wrappers without service parameters have empty service metadata") { withTempDir { root => val conf = codegenConfig(root) @@ -361,11 +405,13 @@ class PyCodegenSuite extends AnyFunSuite { test("service column accessors cannot shadow hand-written Python methods") { withTempDir { root => - val error = intercept[IllegalArgumentException] { - new ConflictingServiceAliasMethodStage().makePyFile(codegenConfig(root)) - } - assert(error.getMessage.contains( - "Generated Python method getModelCol conflicts with a hand-written method")) + assertAliasConflict(root, new ConflictingServiceAliasMethodStage()) + } + } + + test("service column accessors cannot shadow methods in class-body control blocks") { + withTempDir { root => + assertAliasConflict(root, new ConflictingConditionalAliasMethodStage()) } } 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 f6126265bb8..acb57611c2d 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 @@ -14,9 +14,10 @@ import org.apache.commons.io.FileUtils import org.apache.hadoop.fs.Path import org.apache.spark.ml.param.{Param, ParamMap, Params} import org.apache.spark.ml.util._ -import org.apache.spark.ml.{ComplexParamsReadable, ComplexParamsWritable, ObjectSerializer, Serializer, Transformer} +import org.apache.spark.ml.{ComplexParamsReadable, ComplexParamsWritable, ExplicitComplexParamsWritable, + ObjectSerializer, Serializer, Transformer} import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.{DataFrame, Dataset} +import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} import java.io.{ ByteArrayInputStream, @@ -123,6 +124,24 @@ class MixedParamTest(override val uid: String) extends TestEstimatorBase(uid) object MixedParamTest extends ComplexParamsReadable[MixedParamTest] +class TestableSessionParamTest(override val uid: String) extends TestEstimatorBase(uid) with ComplexParamsWritable { + val byteArray = new ByteArrayParam(this, "byteArray", "session-dependent test payload") { + override def save(value: Array[Byte], session: SparkSession, path: Path, overwrite: Boolean): Unit = { + require(session.conf.getOption(TestableSessionParamTest.ConfigKey).contains(uid), "Incorrect writer session") + super.save(value, session, path, overwrite) + } + } +} + +object TestableSessionParamTest extends ComplexParamsReadable[TestableSessionParamTest] { + val ConfigKey: String = "synapseml.test.persistence.session" +} + +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 @@ -144,6 +163,24 @@ class ValidateComplexParamSerializer extends TestBase { } } + 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) + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + try { + 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)) + } finally previous.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + test("Complex Param serialization should work on all complex, all normal, or mixed") { spark 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..5aefcfebbbb --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyBoundedRequests.scala @@ -0,0 +1,796 @@ +// 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.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)(_.readAllBytes()) + 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..5aa004f3544 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyCancellableRequests.scala @@ -0,0 +1,505 @@ +// 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.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("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) + 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 && failure.get().isInstanceOf[TaskKilledException]) + } + } + } + } + + 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 && failure.get().isInstanceOf[TaskKilledException]) + } + } + } + + 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/io/split1/VerifyResponseBodyInspection.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala index 7626e6a5fad..f74ba7a9309 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/io/split1/VerifyResponseBodyInspection.scala @@ -11,11 +11,26 @@ import org.apache.http.entity.BasicHttpEntity import org.apache.http.message.{BasicHttpResponse, BasicStatusLine} import java.io.{ByteArrayInputStream, IOException, InputStream} -import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} import scala.io.Source class VerifyResponseBodyInspection extends TestBase { + test("response conversion closes a failed body before propagating its read error") { + val closes = new AtomicInteger() + val failure = new IOException("Synthetic response conversion failure") + val entity = new BasicHttpEntity() + entity.setContent(new InputStream { + override def read(): Int = throw failure + }) + val response = new BasicHttpResponse(HttpVersion.HTTP_1_1, 200, "OK") with CloseableHttpResponse { + override def close(): Unit = { closes.incrementAndGet(); () } + } + response.setEntity(entity) + assert(intercept[IOException](HandlingUtils.convertAndClose(response)) eq failure) + assert(closes.get() == 1) + } + test("response inspection replays bytes after an input failure") { val content = """{"error":{"code":"RateLimitExceeded"}}""".getBytes("UTF-8") val input = new InputStream { 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..3c3e3e97dcf 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,9 @@ 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") + .filterNot(_.getName == "Quickstart - Durable AI results.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..2560c9e58f5 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,28 @@ 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 opt-in AIPrompt notebooks from automatic Databricks inference coverage") { + val names = Set("AIPrompt.ipynb", "Quickstart - Durable AI results.ipynb") + val examples = DatabricksUtilities.NotebookFiles.filter(file => names(file.getName)) + assert(examples.map(_.getName).toSet == names) + assert(examples.forall(file => !DatabricksUtilities.CPUNotebooks.contains(file))) + assert(DatabricksUtilities.CPUNotebooks.nonEmpty) + } + + test("Exclude opt-in AIPrompt notebooks without constructing a Synapse cloud runner") { + val examples = Array( + new File("ExploreAlgorithmsOpenAIAIPrompt.py"), + new File("ExploreAlgorithmsOpenAIQuickstartDurableAIresults.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..9121c39c494 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,9 @@ 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", + "ExploreAlgorithmsOpenAIQuickstartDurableAIresults", "Deploying", // New issue "TuningHyperOpt", // New issue "IsolationForests", // New issue @@ -64,13 +67,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/com/microsoft/azure/synapse/ml/param/VerifyGlobalParams.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyGlobalParams.scala index 35e349887bf..7ee29f3bdda 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyGlobalParams.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyGlobalParams.scala @@ -8,6 +8,64 @@ import org.apache.spark.ml.param.{Param, ParamMap, Params} import org.apache.spark.ml.util.Identifiable import org.scalatest.TestData +import java.io.{File, FileOutputStream} +import java.net.URLClassLoader +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} +import java.util.jar.{Attributes, JarOutputStream, Manifest} +import scala.util.Using + +object GlobalParamsConcurrencyProbe { + private val Workers = 8 + private val Entries = 2000 + private val Rounds = 4 + private case object SharedKey extends GlobalKey[String] + private case class WorkerKey(value: Int) extends GlobalKey[Int] + + def main(args: Array[String]): Unit = { + val executor = Executors.newFixedThreadPool(Workers) + GlobalParams.setGlobalParam(SharedKey, "shared") + try { + (0 until Rounds).foreach { round => + val start = new CountDownLatch(1) + val work = (0 until Workers).map { worker => + executor.submit(new Callable[Unit] { + override def call(): Unit = { + start.await() + val offset = (round * Workers + worker) * Entries + if (args.head == "registration") { + val params = (0 until Entries).map { index => + new Param[String](s"concurrent-${offset + index}", "value", "concurrent registration") + } + params.foreach(GlobalParams.registerParam(_, SharedKey)) + params.foreach { param => + require(GlobalParams.getParam(param).contains("shared"), s"Lost registration: ${param.parent}") + } + } else { + val keys = (0 until Entries).map(index => WorkerKey(offset + index)) + keys.foreach(key => GlobalParams.setGlobalParam(key, key.value)) + keys.foreach { key => + require(GlobalParams.getGlobalParam(key).contains(key.value), s"Lost global value: ${key.value}") + } + keys.foreach { key => + GlobalParams.resetGlobalParam(key) + require(GlobalParams.getGlobalParam(key).isEmpty, s"Reset failed: ${key.value}") + } + } + } + }) + } + start.countDown() + work.foreach(_.get(30, TimeUnit.SECONDS)) + } + println("GLOBAL_PARAMS_CONCURRENCY_OK") + } finally { + executor.shutdownNow() + } + } +} + class VerifyGlobalParams extends TestBase { // Test keys @@ -98,4 +156,44 @@ class VerifyGlobalParams extends TestBase { val result = GlobalParams.getParam(params.unregisteredParam) assert(result.isEmpty) } + + test("global values preserve a present null until reset") { + val value = Option.empty[String].orNull + GlobalParams.setGlobalParam(TestStringKey, value) + assert(GlobalParams.getGlobalParam(TestStringKey) == Some(value)) + GlobalParams.resetGlobalParam(TestStringKey) + assert(GlobalParams.getGlobalParam(TestStringKey).isEmpty) + } + + Seq("registration", "values").foreach { mode => + test(s"concurrent global $mode completes without losing entries in an isolated JVM") { + val loaderPaths = Iterator.iterate(getClass.getClassLoader)(_.getParent) + .takeWhile(_ != null) + .collect { case loader: URLClassLoader => loader } + .flatMap(_.getURLs.iterator).filter(_.getProtocol == "file") + .map(url => new File(url.toURI)).toVector + val systemPaths = System.getProperty("java.class.path").split(File.pathSeparator).map(new File(_)) + val paths = (loaderPaths ++ systemPaths).map(_.getCanonicalFile).distinct + val manifest = new Manifest() + manifest.getMainAttributes.put(Attributes.Name.MANIFEST_VERSION, "1.0") + manifest.getMainAttributes.put(Attributes.Name.CLASS_PATH, paths.map(_.toURI.toASCIIString).mkString(" ")) + val classpath = tmpDir.resolve(s"global-$mode-classpath.jar").toFile + Using.resource(new JarOutputStream(new FileOutputStream(classpath), manifest))(_ => ()) + val log = tmpDir.resolve(s"global-$mode.log") + val java = new File(System.getProperty("java.home"), "bin/java").getAbsolutePath + val probe = GlobalParamsConcurrencyProbe.getClass.getName.stripSuffix("$") + val process = new ProcessBuilder(java, "-cp", classpath.getAbsolutePath, probe, mode) + .redirectErrorStream(true).redirectOutput(log.toFile).start() + val completed = try process.waitFor(45, TimeUnit.SECONDS) finally { + if (process.isAlive) { + process.destroyForcibly() + assert(process.waitFor(10, TimeUnit.SECONDS), "GlobalParams probe did not terminate") + } + } + val output = new String(Files.readAllBytes(log), StandardCharsets.UTF_8) + assert(completed, s"GlobalParams concurrency probe timed out:\n$output") + assert(process.exitValue() == 0, output) + assert(output.contains("GLOBAL_PARAMS_CONCURRENCY_OK"), output) + } + } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/VerifyUDFTransformerPersistence.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/VerifyUDFTransformerPersistence.scala index 32b8c2dace2..686a3f90af3 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/VerifyUDFTransformerPersistence.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/VerifyUDFTransformerPersistence.scala @@ -5,6 +5,7 @@ package com.microsoft.azure.synapse.ml.stages import com.microsoft.azure.synapse.ml.core.test.base.TestBase import org.apache.spark.ml.Serializer +import org.apache.spark.sql.SparkSession import org.apache.spark.sql.expressions.UserDefinedFunction import org.apache.spark.sql.functions.udf @@ -16,6 +17,32 @@ private object UDFPersistenceTestHelpers { class VerifyUDFTransformerPersistence extends TestBase { + Seq(false, true).foreach { allowed => + test(s"complex readers use their assigned Spark session deserialization policy with allowed=$allowed") { + import spark.implicits._ + + 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 { + val path = tmpDir.resolve(s"reader-session-$allowed.model").toString + new UDFTransformer().setUDF(UDFPersistenceTestHelpers.StringToIntegerUDF) + .setInputCol("words").setOutputCol("out").write.session(readerSession).save(path) + if (allowed) { + val loaded = UDFTransformer.read.session(readerSession).load(path) + assert(loaded.transform(Seq("safe").toDF("words")).select("out").as[Int].collect().sameElements(Array(1))) + } else { + val error = intercept[SecurityException](UDFTransformer.read.session(readerSession).load(path)) + assert(error.getMessage.contains(config)) + } + } finally previous.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + test("Persisted UDFs require explicit trusted legacy opt-in") { import spark.implicits._ 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..c0c46b6049d --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/AIPrompt.ipynb @@ -0,0 +1,894 @@ +{ + "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", + "| 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." + ] + }, + { + "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 run_example(name, stage, data):\n", + " rows = stage.transform(data).collect()\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", + "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", + "\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 = run_example(\"responses/store-and-response-id\", stored, countries)\n", + " response_id = first_rows[0][stored.getResponseIdCol()]\n", + " assert response_id\n", + " created_response_ids.append(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", + " for response_id in created_response_ids:\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", + " assert deleted.json()[\"deleted\"] is True\n", + " deleted_response_ids.append(response_id)\n", + " print(f\"Deleted {len(created_response_ids)} response(s) created by this example.\")" + ] + }, + { + "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/Durable AI results.md b/docs/Explore Algorithms/OpenAI/Durable AI results.md new file mode 100644 index 00000000000..eb5b37442ae --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/Durable AI results.md @@ -0,0 +1,529 @@ +--- +title: Durable AI results +hide_title: true +sidebar_label: Durable AI results +description: Opt-in Delta snapshots, committed outcomes, and explicit recovery of uncommitted AI requests. +--- + +# Durable AI results + +**Draft recovery follow-up:** this extracted implementation is not runtime-qualified. +It builds on [the focused native APIs](Unified%20AI%20prompting.md); prerequisite +test passes do not qualify recovery. Managed Spark application-restart recovery +remains unqualified. See [issue #2800](https://github.com/microsoft/SynapseML/issues/2800). + +The [opt-in recovery notebook](Quickstart%20-%20Durable%20AI%20results.ipynb) starts +with all execution flags disabled. The examples below assume `spark`, an ordinary +configured `prompt`, its input DataFrame `tickets`, and the `AIPrompt` import from +the native guide. Storage and model operations require explicit caller authorization. + +## Durable Delta results + + +Set `outputPath` only when you intend to persist **both inputs and outcomes**. +Without it, no Delta dependency is needed and normal Spark replay applies: +`show`, `collect`, another action, recomputation or task retries can repeat paid +requests. Provider `store=true` is unrelated to Spark durability. +Output paths must not contain SQL substitution markers such as `${team}`, even +when Spark variable substitution is disabled. These paths are rejected before +storage is created because SQL control operations and DataFrame reads must +address the same literal directory. Resolve application variables before calling +`setOutputPath`. + +```python +durable = prompt.copy({}).setOutputPath("Files/aiprompt/ticket-run") +durable.setOutputBatchSize(1000) +saved = durable.transform(tickets) # Eager, bounded Delta commits. +# These actions read storage, not inference lineage. +saved.count() +saved.show() +status = AIPrompt.getOutputStatus(spark, "Files/aiprompt/ticket-run") +reopened = AIPrompt.loadOutput(spark, "Files/aiprompt/ticket-run") +remaining = durable.resume(spark) # Same configuration; newly bound credentials. +``` + +```scala +import org.apache.spark.ml.param.ParamMap + +val path = "Files/aiprompt/ticket-run" +val durable = prompt.copy(ParamMap.empty).setOutputPath(path).setOutputBatchSize(1000) +val saved = durable.transform(tickets) // Eager; subsequent actions read Delta. +val reopened = AIPrompt.loadOutput(spark, path) +val status = AIPrompt.getOutputStatus(spark, path) +val remaining = durable.resume(spark) // Requires matching configuration and no active owner. +``` + +**Durable stages are one-shot jobs, not reusable Pipeline stages.** A run directory +can be initialized only once. Keep a stage with `outputPath` outside `Pipeline`, +`PipelineModel`, cross-validation and other tuning wrappers. `Pipeline.fit` may +eagerly run a preceding transformer before fitting an estimator, spending inference +requests during fitting; later scoring, folds or repeated calls then fail because +the run directory already exists. Saving or copying that stage does not create a +fresh run. + +Use the standalone `durable.transform(tickets)` call above, reopen its committed +results with `AIPrompt.loadOutput`, handle terminal row errors, and fit or apply a +downstream pipeline to that stored DataFrame. The downstream pipeline must not +include the durable inference stage. Use a new `outputPath` for each new inference +job; use `resume` only to continue the original saved input snapshot. + +`AIPrompt.load(stagePath)` restores a saved SparkML stage and its Params; +`AIPrompt.loadOutput(spark, runPath)` reads committed outcomes without model +credentials. These paths contain different artifacts and are not interchangeable. +Scala run-management methods are on the `AIPrompt` companion for all three stage +types; the generated Python classes also inherit those convenience methods. +Use the saved native stage's reader (or `AIPrompt.load`) for SparkML persistence; +loading new native artifacts through the legacy `OpenAIPrompt` reader is outside this feature. + +The caller must supply a Delta-capable runtime and a fresh run directory. +The preserved local test setup used the following caller-provided JVM artifact. +It is not a mandatory inference dependency or a qualification claim for this extracted revision: + +| Spark / Scala | Caller-provided Delta artifact | +| --- | --- | +| 3.5 / 2.12 | `io.delta:delta-spark_2.12:3.2.0` | + +Other Spark/Scala ports require separate source adaptation and qualification; +this table does not establish their compatibility. Durable mode currently requires +the `org.apache.spark.sql.delta.sources.DeltaDataSource` class, +`DeltaSparkSessionExtension` in `spark.sql.extensions`, and `DeltaCatalog` in +`spark.sql.catalog.spark_catalog`. The extension and catalog must be configured +before the Spark session is created. + +The opt-in `aiprompt_fabric_recovery_e2e.py` probe is intended to check Fabric OneLake +interruption, owner release and resume within the same managed Spark 3.5 +application using bundled Delta and synthetic HTTP responses. The extracted revision +has not passed this probe; prior combined-tree evidence does not qualify it. Recovery after +stopping and restarting a managed Spark application remains unqualified. +It is not part of CI: run it against the exact SynapseML build and managed runtime +before treating that combination as qualified. This storage check does not +establish provider availability. Bundled Delta alone does not establish compatibility: +runtimes with different provider packages or catalog/extension arrangements are +not supported by these checks. Do not replace managed-runtime jars to bypass +them. Accepted path schemes are `hdfs`, `abfs` and `abfss`, plus `file` only +in Spark local mode with a local directory on the same host. Qualified `file` +URIs must have no authority: host-qualified paths (including `localhost`) and +UNC paths are rejected before creating run storage. Non-local Spark +clusters reject `file` paths before creating run storage: node-local data files +and ownership locks cannot coordinate across machines. Local-mode recovery also +requires access to that same host and directory; it does not protect against +loss of that disk. Shared FUSE/NFS mounts are not a substitute for supported +cross-host storage. The `dbfs`, `s3`, `s3a` and `gs` schemes are not qualified. +This is not a claim of Databricks or Synapse runtime qualification. Lazy +inference does not require these durable storage capabilities. + +Durable schemas widen nullability and remove Spark-internal alias/file metadata +and Delta generated-column/identity metadata, including in nested fields. Those +source-column behaviors do not belong to a materialized snapshot and Delta may +remove them on read. Ordinary user field metadata and input values are retained; +lazy inference does not apply this storage normalization. + +`_control/`, `_input/` and `results/` belong to this run. There is no implicit +overwrite, cleanup, schema evolution, append of a second input, or background +job. Only bounded pending keys reach the driver; input rows remain distributed. +Repeated source rows get distinct persisted identities. Committed row errors +are terminal; retry selected errors by starting an explicit new run. +Internal row identities are excluded from prompt-template expressions, including +wildcards such as `{to_json(struct(*))}`, during both initial inference and resume. +Attachment filename interpolation does not add fields to wildcard expansion: +`{asset}` uses the attachment filename, while `{to_json(struct(*))}` retains only +the public input columns and their original values. +Durable prompt templates additionally require stable built-in expressions over input +columns and literals. Unsupported templates are rejected before run storage is created. +SQL subqueries, UDFs (even deterministic ones), nondeterministic +expressions and runtime-dependent leaves such as `current_timestamp()` are +rejected: external tables, registered functions and session/runtime values could +change between batches or resumes. Materialize those values into input columns +before `transform` and reference those columns in the template; the input +snapshot then preserves them. These additional stability restrictions do not apply +to lazy mode. +Materialize `aes_encrypt` results as well: Spark can mark encryption deterministic +even when it generates a random initialization vector. Durable templates reject all +`aes_encrypt` forms, including explicit-IV and ECB forms; materializing ciphertext +avoids recomputing it during recovery without requiring a fixed IV. +For every API, durable initialization evaluates the template over the committed +input snapshot before sending any inference requests. Before validation, the +snapshot must still have its original schema and data: only intervening `OPTIMIZE` +compaction commits are accepted. The validated input version then remains pinned +through initialization; detected out-of-band changes stop inference rather than +adopt an unchecked snapshot. A row-specific SQL error +(for example, a failing date conversion or ANSI cast) leaves the run in +`initializing`, with no resumable results and no requests sent. Correct the input +or template and use a new output directory; do not resume that incomplete run. +Malformed per-row body extensions also become committed terminal errors without +aborting initialization. Resuming a completed run checks the saved configuration +and ledger integrity, then reads committed output without re-evaluating prompt +templates or parsing input JSON again. It does not acquire an owner or modify the +control table. + +**Keep the private `_input` table immutable after initialization.** Exclude it +from manual and automatic maintenance: do not `OPTIMIZE`, `VACUUM`, append, +update, or change its schema. Reads, status checks, and resume require its latest +version to equal the pinned version, even for a completed run. A later compaction +or logged `VACUUM` commit is rejected even when the pinned files remain readable; +vacuuming can also remove files required by the pinned snapshot. There is no +fallback to a different snapshot. The pre-validation compaction allowance above +does not permit maintenance after the input version is pinned. + +**Defer `results/` maintenance until the run is complete and has no owner.** +Exclude it from manual and automatic `OPTIMIZE` and `VACUUM` while the run is +unfinished or owned. Between batches, the writer requires the results version +to remain unchanged; even data-preserving compaction can stop inference with +"committed outcomes changed outside the current owner" and retain ownership. +If that happens, stop the previous Spark application, establish that its requests +and tasks have quiesced, then use the stop-confirmed administrative release and +resume procedure below. After completion, maintenance must still preserve the +results schema, row identities and outcomes, and respect Delta retention and +readers of older snapshots. Never append, update or delete result rows, and leave +the private `_control/` table managed by the run APIs. This allowance for completed +results does not relax the `_input/` restrictions. + +Recognized credential values are not written into the manifest or used in its +fingerprint. For these credentials, the fingerprint records only whether +authentication is configured, not its value or source. Scalar keys, tokens and +recognized custom authentication headers may rotate or replace each +other on resume. Implicit Fabric authentication also counts as configured +without fetching a token during fingerprinting. Removing all authentication from +an authenticated run fails before ownership acquisition or inference; blank +credentials and auth-named telemetry do not satisfy this check. The same +configuration check applies to completed resume; `loadOutput` remains available +without inference credentials. This does not validate a key's validity or scope: +provider authentication failures still become terminal row errors. +Authentication-named entries in scalar custom and telemetry headers +are excluded from the fingerprint; other header settings remain significant. +The excluded names are the stage's subscription-key and AAD header names +(`api-key` and `Authorization` by default), compared without regard to casing. +Subclasses that override these names use their own names for credential rotation; +headers no longer used for authentication remain part of the fingerprint. +Credentials under unrecognized header names, such as a gateway's +`Ocp-Apim-Subscription-Key` in `customHeaders`, are treated as ordinary settings: +their values contribute to the persisted SHA-256 fingerprint, do not count as +configured authentication, and cannot change on resume. Hashing is not a +substitute for protecting low-entropy secrets. Use recognized authentication +bindings or override the stage's authentication header names when rotation is +required. +Recreate the original stage class as well as its settings: the fingerprint records +the concrete class independently of telemetry, even when telemetry is disabled +or customized. `AISystemOne` and `AIPrompt(apiType="systemone")` +(likewise `OpenAIDecisions` and the Decisions facade) are inference-equivalent, +not interchangeable recovery constructors. +Use the same SynapseML build and Spark runtime versions to resume. The manifest +records both version strings, and recovery rejects a different version before +acquiring ownership or sending requests, including for completed runs. +Library-owned defaults (including prompt, output, telemetry and dated API-version +defaults) also participate in recovery identity. The manifest stores their hash, +not the original configuration; the hash cannot identify which default changed. +Finish pending runs before upgrading, or retain the original environment for +recovery. If it cannot be restored, `loadOutput` still reads committed results; +do not replay the full input. +Other effective configuration must match its fingerprint. This includes the +resolved URL of an automatically selected Fabric endpoint: its workload host, +workspace and capacity are part of the run's identity. Finish durable runs before +moving the workspace to another capacity or changing its workload routing. +If that endpoint changes, `resume` rejects the configuration; automatic migration +of a pending run is not supported. `loadOutput` can still read committed results. +To continue elsewhere, start a new run with deliberately selected unprocessed +source rows rather than replaying the entire original input. + +Object +order in `responseFormat`, `extraBody` and `questions` is fingerprint-significant, including +structured-output property order. Their JSON numbers are compared by exact value +within the rendered integer or decimal kind; equivalent +decimal spellings match only when serialization retains that kind. Array order +and string values remain significant. Fingerprinting preserves UTF-16 code units: +an invalid string setting is not equivalent to a question mark, a replacement +character or literal JSON escape text. For example, changing a malformed +`systemPrompt` to `"?"` cannot resume a run whose committed rows contain +`invalid_request` errors for that original setting. SQL +expression settings also participate: time zone, ANSI and legacy semantics, +parser and function options, Base64 line wrapping, JSON/CSV parsing (including JSON partial results), +the corrupt-record field name, JSON generation, decimal precision, case sensitivity, +timestamp type, map-key deduplication and assignment policy. Changing these can +change pending prompt values even when the public schema stays the same. Restore +them on the executing session before resuming; another active session does not +determine the fingerprint. This includes registered SQL defaults, not just +explicitly configured settings. The same identity check runs before preparing +each pending batch, not only on resume. A detected change stops before that +batch sends requests and releases the quiescent owner; restore the original +settings before resuming. Do not concurrently mutate SQL semantic settings on +the executing session during a run. Resume on the original Spark runtime: a runtime +update that adds semantic settings or changes their defaults can reject recovery, +even with unchanged stage parameters. The check deliberately fails closed rather +than submit pending prompts with changed semantics. If the original runtime +cannot be restored, `loadOutput` still reads committed results; do not blindly +resubmit the full input. Spark execution tuning such as `spark.sql.shuffle.partitions` +does not participate. Stage transport settings do: `concurrency`, `timeout` and +`concurrentTimeout` must retain their original values and explicit/default bindings +on resume. The durable contract conservatively freezes these settings because +they affect request execution and failure outcomes. HTTP `timeout` exhaustion +can become a terminal row error, whereas `concurrentTimeout` expiry fails the +Spark task rather than creating a row error. If it causes a durable batch write +to fail, ownership remains held; confirm the previous execution has stopped +before releasing that owner to recover. Spark retries can resend requests in +the uncommitted batch; earlier committed batches remain protected from replay. +`concurrentTimeout` cannot be increased on `resume`. Tune these settings before +starting a run. To change them after an interruption, start an explicit new run +with a fresh `outputPath` and deliberately select the rows to submit; `resume` +neither retunes pending requests nor retries committed errors. +Only the hash is stored, not these settings' plaintext. Do not +include sensitive columns in the input unless you intend to persist them. +Credential, custom-header and telemetry-header column bindings are rejected in +durable mode, before input persistence. Use scalar headers instead. Remote media, +provider file IDs, previous responses and retrieval/tool state are rejected +before inference: saved inputs must be self-contained. This includes nested +resource bindings and `data_sources`, `web_search_options` and stored `prompt` +references in body extensions, as well as known `image_url` and `file_url` +locations and object-shaped `url` carriers (including `file_data.url`), in scalar or +per-row extensions, typed `responseFormat`/`responseFormatCol` values and native +`questions`/`questionsCol` payloads. Arrays in these +media locations are checked recursively. Only a direct scalar string `file_data` +is treated as literal content; strings inside `file_data` arrays must be inline +data URLs, and object `file_data.url` values remain location checks. These conservative +checks apply even to malformed media shapes and native JSON state. +System One JSON +state and Decisions JSON input are checked on the committed snapshot, whether +supplied by `inputCol`, a rendered `promptTemplate`, or System One's `stateCol`, +before any inference request. These checks cover all successfully parsed values, +even malformed message or question shapes; restricted dependencies abort the +snapshot rather than becoming row errors while valid sibling rows are sent. +In System One JSON state, resource-shaped fields such as +`file_id`, `image_url` and `tools` are conservatively subject to these restrictions. +At any depth, non-null `previous_response_id`, `conversation`, `conversation_id`, +`file_id`, `file_ids`, `tools`, `tool_choice`, `tool_resources`, `functions`, +`function_call`, `retrieve`, `run_id`, `data_sources`, `web_search_options`, `prompt` +and `tool_call_id` are rejected, even if intended as ordinary user-defined names. +Rename those JSON members or use text mode for literal state; ordinary text state +is not interpreted as resource bindings. Diagnostics identify restricted member +names without including their values. Question names, +System One criteria labels, and field names inside an output JSON schema are +descriptions, not resource bindings; inert URLs in ordinary text remain valid. +Only criteria labels are exempt: nested objects and arrays used as criteria values +still undergo resource checks, even when the question shape is malformed. +Typed formats use the same checks after API-specific normalization; output-schema +properties remain exempt, but neighboring format extension fields do not. +Scalar format dependencies fail before storage creation; per-row dependencies +abort snapshot validation before inference. Malformed per-row formats remain +terminal row errors and are not sent or replayed on resume. +Output-schema exemptions apply only to object-member paths in the matching API's +request body: `response_format.json_schema.schema` for Chat and +`text.format.schema` for Responses. Array-wrapped paths, `response_format.schema`, +other APIs, and similarly named paths inside native state or input are not exempt. +Top-level `extraBody.metadata` is also an inert user-defined map: names such as +`conversation`, `prompt`, `tools` or `image_url` there do not create dependencies. +Resource checks still apply elsewhere in the request, including nested extensions. +Responses `include` options that request extra output, such as +`message.output_text.logprobs` and `reasoning.encrypted_content`, are supported in +scalar and per-row body extensions. They do not enable tools or bypass the independent +resource restrictions. Requested provider fields remain available in `rawResponseCol`. +Optional resource and message ID fields containing JSON null are not external +dependencies; actual IDs and remote media remain rejected. +Malformed per-row message shapes, native JSON input and body extensions become +committed terminal row errors rather than aborting initialization. Messages that +exceed JSON parser limits or otherwise cannot be checked for external resources +instead abort initialization before inference; they are never sent unchecked. +Unsupported content +forms are rejected by request validation without sending them to the provider. +Path attachment cells use the same trimming and missing-value handling as lazy +inference: null, empty and whitespace-only cells remain row-level missing inputs +when no attachment is available. Their direct filename placeholders render as +empty strings, so an absent optional attachment does not discard another usable +attachment. Null non-path template values retain Spark's null behavior. +Inline data URLs retain their case-insensitive +scheme/base64 handling, and filename interpolation uses the same normalized +path as attachment construction; actual remote paths are still rejected before inference. +Empty or malformed `data:` media values are self-contained and do not abort the +resource check. They follow the same row-level request validation and provider +error handling as lazy inference; non-`data:` media locations remain rejected. +String-valued `file_data` is literal inline content, not a media location. It is +preserved unchanged rather than interpreted as a path or downloaded, even when +its text resembles a URL. This exemption does not extend to strings inside +array-valued `file_data`; those are checked as media locations. +Supply the encoding required by the provider, normally +a base64 data URL; preserving raw strings does not guarantee provider acceptance. +Invalid file content follows ordinary request validation and provider error +handling, and a committed provider error is not resent on resume. +Simple attachment-column placeholders follow Spark's column resolver, including +case-insensitive aliases and quoted identifiers. Escaped literal braces and SQL +expressions retain their usual template meaning; only direct column placeholders +are replaced with attachment filenames. +Arbitrary input UDTs and nested empty structs (including inside arrays and maps) +are rejected before storage creation. Decimal fields must have nonnegative scales, +including inside structs, arrays and map keys or values: Delta's Parquet storage +cannot represent negative scales even when Spark's legacy option permits them. +This durable restriction does not change lazy inference. Spark ML vectors use a reversible storage +codec, including vectors in retained map keys and values. A zero-column logical input root is supported because the stored +snapshot includes an internal row key. Durable schemas use Delta-compatible nullability. +Sibling field names at every nesting level must also be unambiguous under the +executing session's Spark column resolver; Delta statistics collection resolves +nested fields by name. This check is separate from Delta's case-insensitive +storage-name restriction, which still applies with `spark.sql.caseSensitive=true`. +Stored schema JSON, including field metadata, uses the same bounded JSON parser +as native requests; excessive nesting is rejected before creating a run. +The limit applies to the encoded representation: each struct field adds three +JSON levels, while array elements and map keys or values add one. Field metadata +shares that depth budget; it is not a separate allowance. +Malformed stored schemas reject reads, recovery and owner release without changing the run. +Missing or +incompatible mandatory output bindings are corruption errors, never evidence +of zero failed rows. + +After an interruption or error in an initialized run, including a failed cell in +the same application, first inspect status and partial results. Successful completion clears ownership. +Initial execution and resumed execution both report `running` once initialized. +This state records execution progress, not application liveness; it does not +replace the owner-token and quiescence checks below. +The count-based `isComplete` flag means every input row has a committed outcome; +`allSucceeded` additionally requires zero row errors. After the last batch commits, +these flags can be true while `state` is still `running` and `ownerToken` remains +present. Neither flag proves task quiescence or authorizes owner release. +Follow the stop-confirmed release procedure below; `resume` then finalizes the run +without repeating committed rows, even if no rows were pending. +A driver-side batch-preparation failure releases only its matching owner token +when all prior actions finished successfully and the next distributed action has +not started. Correct the preparation error before calling `resume`; committed +rows are not repeated. Cancellation, failed Spark jobs and uncertain storage +outcomes retain ownership because task/request quiescence cannot be proven. +A failed or disconnected notebook cell is not proof that its executor work stopped. +If status still has an owner, **stop the old Spark application and establish that +no work remains** before using +`releaseOutputOwner(spark, path, expectedOwnerToken, confirmStopped=True)`. +Falsely confirming stopped execution does not fence in-flight writes and can corrupt the ledger. +Use the exact token from status. There is no automatic expiry, takeover or +owner-independent force flag. Recreate the same stage configuration, bind fresh +scalar credentials and call `resume(spark)`; never pass a replacement input. +An incomplete initialization is an explicit error rather than an empty success. +It does not establish that initialization failed: the creating application may +still be writing or validating the input snapshot. Wait for initialization to +finish, or confirm that the creating application has stopped before abandoning +the path, starting a replacement run, or cleaning up. This also applies while +control-table data files exist but their Delta commit is not yet visible. +Run-management calls distinguish a missing path from an existing directory that +is not an AIPrompt run. Check the supplied path in either case; these errors do +not mean an existing run needs to restart and do not authorize deleting unrelated +data. A new `transform` can still create a run at a missing or empty path, but it +never overwrites an unrelated non-empty directory. +On local storage, a failure that leaves only the control lock file can retry +`transform` on that still-uninitialized directory. If a failed control-table write +leaves data files without a committed Delta log, run-management calls and another +`transform` report incomplete initialization rather than suggesting that the +directory can be resumed. They retain those files unchanged; after confirming +stopped execution, use a new `outputPath`. +If initialization fails after +the control table is created but before the run is ready, `getOutputStatus`, +`loadOutput` and `resume` explicitly reject the incomplete run. This includes a +snapshot rejected for external input dependencies. After confirming stopped +execution, correct the input and use a new `outputPath`; releasing ownership +does not make an incomplete run resumable. +The original directory is retained for inspection, not overwritten or automatically +deleted. Stop the original Spark application and establish that no work remains +before any administrative cleanup of that directory. +Local control transactions serialize both same-driver threads and separate drivers; +the temporary file lock does not replace or expire durable ownership. Waiting for +another local driver's control lock is interruptible, including on Windows. +Cancelling that wait does not run the control transaction or clear the durable owner. +Without cancellation, acquisition waits for the current lock holder; there is no +lock-wait deadline or automatic takeover. +On remote storage, ownership races can surface Delta concurrency or path-already-exists +errors rather than an AIPrompt-specific exception. Inspect status and partial results +and follow the stopped-execution rules above. These errors do not authorize takeover, +deleting the run directory, or blindly restarting inference. + +The committed results table is the recovery ledger. An uncommitted interrupted +batch may run again; ambiguous commits are reconciled from Delta on resume. +Neither Delta nor this feature provides exactly-once external inference. +Provider ambiguity, Spark retries and speculation may still incur repeat charges. + +In Fabric, use the default Lakehouse's `Files/...` path or a copied qualified +ABFS path with the existing Spark storage identity. This creates a run directory, +not a catalog table. Path qualification and directory checks use the supplied +Spark session's effective filesystem configuration, including session overrides. +Local tests and OneLake require their respective Delta +transaction guarantees; arbitrary filesystem configurations, particularly +single-driver S3 configurations, do not establish cross-driver ownership. + +## Scaling responsibly + + +Before starting a run, tune partition count and `concurrency` to the deployment's rate limit. +In lazy mode, repartition the input DataFrame deliberately. Durable mode reads a +saved snapshot, so input repartitioning does not determine inference parallelism. +Each durable batch is explicitly repartitioned to the smaller of its row count +and the executing session's `spark.sql.shuffle.partitions`. Use a positive +integer in the qualified runtimes and set that Spark option deliberately +(for example, `spark.conf.set("spark.sql.shuffle.partitions", "8")`); +it can also be changed on resume without changing saved request settings. +On runtimes that accept `auto`, durable execution uses a fixed fallback of at most +200 partitions per batch, further capped by that batch's row count. This is not +cluster-adaptive sizing and does not create one partition per row for large batches. +Use an explicit integer to tune it for your quota and available executor slots. +This fallback is unit-tested; runtime acceptance of `auto` is not qualified here. +Explicit repartitioning prevents AQE from coalescing these partitions, and Delta +appends can produce a file per non-empty partition. Both the automatic fallback +and large explicit counts can therefore increase scheduling and small-file overhead. +Do not compact those result files during the run; follow the maintenance +restrictions above after completion. +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. Settle this tuning before starting a durable run: +`concurrency` cannot change on resume; a different value requires a fresh run. + +`outputBatchSize` bounds a +durable commit, not provider quota. Start with small batches to bound the +uncommitted work exposed to a crash, then measure throughput and commit overhead. +Do not collect the input dataset or loop over rows on the driver. +The implementation collects only the selected batch's numeric keys on the driver, +then builds two literal-key `IN` predicates to select and verify that batch. +Larger batches therefore add driver memory, query-planning and plan-log overhead, +even before inference. The maximum of 100,000 rows is a bound, not a +production sizing recommendation or a qualified performance target. Increase +from small batches only after measuring this cost in your runtime. +Durable execution reuses each validated post-commit ledger and skips repeated +input count/distinct validation, not all input-key scans. Every batch joins +against input keys to select pending rows and to verify committed keys, and scans +the growing results ledger for integrity checks. In the worst case, total key-scan +work grows with input size times the number of batches. Choose `outputBatchSize` +to balance this overhead against the uncommitted work exposed to a crash; +many tiny batches can be expensive. `getOutputStatus`, `loadOutput` and successful +`resume` attempts also run full input and result integrity checks. An unfinished +run that already has an owner is rejected before scanning its input or results; +the later ownership acquisition still checks for races. +After these integrity checks, returned result DataFrames preserve Spark's column +pruning and filter pushdown for columns without vectors. Only top-level columns +containing vectors need conversion, including nested vectors; unrelated raw +responses or retained messages need not be read by a narrow projection. Vector +payload validation runs when a vector-bearing column is read, not when it is +pruned from the query. Full schema and row-key integrity checks are unchanged. +Vector-bearing containers are converted with Spark SQL functions, with UDFs only +at vector leaves. Retained vector map keys are also encoded around inference's +row conversions in both lazy and durable modes, 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; durable +mode separately permits only Spark ML vector UDTs. + +In durable mode all committed row errors are terminal, including exhausted +transport failures and 429/5xx responses: `resume` does not retry them. There is +no failure-rate circuit breaker; later batches continue after row errors, so an +outage or quota limit can leave many rows without model output. Inspect the error +column and `rowErrorCount`; a complete run means all rows have committed outcomes, +not that every inference succeeded. Retry selected failed rows in a new run with +a fresh output path after the service recovers, accounting for possible duplicate +provider work when the original request's outcome was uncertain. diff --git a/docs/Explore Algorithms/OpenAI/Quickstart - Durable AI results.ipynb b/docs/Explore Algorithms/OpenAI/Quickstart - Durable AI results.ipynb new file mode 100644 index 00000000000..a2c1760b767 --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/Quickstart - Durable AI results.ipynb @@ -0,0 +1,135 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Durable AI results and recovery\n", + "Use a Delta-capable Spark runtime. This opt-in notebook persists input rows and outcomes in a fresh caller-owned run directory; apply normal storage access and retention policies. It never creates a Lakehouse or installs storage dependencies.\n", + "\n", + "**This extracted follow-up is not runtime-qualified. Its managed-runtime probe only targets controlled interruption, owner release and resume within the same Spark application. Recovery after stopping and restarting a Fabric Spark application remains unqualified.** The recovery protocol below is not evidence that a managed application restart has been validated. See [the guide's runtime qualification](Durable%20AI%20results.md#durable-delta-results).\n", + "\n", + "Without `outputPath`, Spark actions may repeat model calls. With `outputPath`, `transform` is eager and returns a Delta-backed DataFrame. Committed batches survive a session; uncommitted work, provider ambiguity and Spark retries can still repeat paid inference. This is **not exactly-once inference**.\n", + "\n", + "**Keep durable inference outside reusable SparkML pipelines and tuning wrappers.** Each `outputPath` is a one-shot run directory. `Pipeline.fit` may run inference before fitting a downstream estimator; later `PipelineModel.transform`, folds or repeated calls then fail on that existing directory. Run the standalone transform below, read committed results with `loadOutput`, handle row errors, and give that stored DataFrame to a downstream pipeline that does not contain the durable stage. New inference jobs require fresh paths." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from synapse.ml.services.ai import AIPrompt\n", + "\n", + "RUN_NEW = False\n", + "RUN_READ = False\n", + "RUN_RESUME = False\n", + "# Set a unique fresh path for a new run; keep exactly this path for recovery.\n", + "# Fabric: default Lakehouse Files/... or the copied qualified ABFS path.\n", + "RUN_PATH = 'Files/aiprompt/my-ticket-run'\n", + "\n", + "def configured_stage():\n", + " return (AIPrompt().setApiType('responses')\n", + " .setUrl(os.environ['MODEL_ENDPOINT'])\n", + " .setDeploymentName(os.environ['MODEL_DEPLOYMENT'])\n", + " .setSubscriptionKey(os.environ['MODEL_API_KEY'])\n", + " .setPromptTemplate('Summarize this ticket: {ticket}')\n", + " .setOutputCol('answer').setErrorCol('error').setUsageCol('usage')\n", + " .setRawResponseCol('raw').setOutputPath(RUN_PATH).setOutputBatchSize(2))\n", + "\n", + "if RUN_NEW:\n", + " tickets = spark.createDataFrame([\n", + " (1, 'Checkout is unavailable.'),\n", + " (2, 'Please explain my invoice.'),\n", + " (3, 'How do I update my address?'),\n", + " ], ['ticket_id', 'ticket'])\n", + " saved = configured_stage().transform(tickets)\n", + " assert saved.count() == 3 # Reads committed Delta, not live inference.\n", + " saved.select('ticket_id', 'answer', 'error').show(truncate=False)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Read from another session without model credentials\n", + "Storage identity is still required. Status derives counts from a pinned committed results version. All rows having terminal outcomes does not mean all model calls succeeded. Row errors are terminal and are not automatically resubmitted.\n", + "\n", + "A new Spark session is not necessarily a new Spark application. This example does not qualify reads or recovery after a managed application restart; the managed qualification boundary above still applies.\n", + "\n", + "`isComplete` means every row has a committed outcome; `allSucceeded` also requires zero row errors. These count flags can be true while `state` is still `running` and `ownerToken` remains present after the last batch commits. Neither flag proves task quiescence or permits owner release. Follow the stop-confirmed recovery procedure below; `resume` then finalizes without repeating committed rows, even when no rows are pending.\n", + "\n", + "`AIPrompt.load(stagePath)` restores a saved SparkML stage and its settings. `AIPrompt.loadOutput(spark, RUN_PATH)` instead reads committed outcomes; a stage directory and a run directory are not interchangeable." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if RUN_READ:\n", + " status = AIPrompt.getOutputStatus(spark, RUN_PATH)\n", + " print(status)\n", + " partial_or_complete = AIPrompt.loadOutput(spark, RUN_PATH)\n", + " partial_or_complete.select('ticket_id', 'answer', 'error').show(truncate=False)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Recover after interruption\n", + "For Fabric, the probe targets controlled interruption and recovery within one Spark application. Stopping and restarting that application remains unqualified; do not treat the administrative sequence below as a validated managed-restart walkthrough. This limitation does not relax the stop/quiescence requirement when ownership is retained.\n", + "\n", + "Inspect status after any interruption or error, including a failed cell in the same application. Successful completion clears ownership. A driver-side batch-preparation failure also releases its matching token when all prior actions finished successfully and no new distributed work started; fix the preparation error before resuming. Cancellation, failed Spark jobs and uncertain storage outcomes retain ownership because task/request quiescence is not proven. If status still has an owner, **first stop the previous Spark application and prove its requests/tasks have quiesced.** A failed or disconnected notebook cell is not proof. Copy the exact token from status, and only then perform the administrative release below. A mismatched token fails; there is no automatic takeover or force flag.\n", + "\n", + "Recreate the original stage class, identical non-credential request settings and output names. The concrete stage class is recorded independently of telemetry, even when telemetry is disabled or customized, so direct stages and the facade are not interchangeable recovery constructors. Stage transport settings `concurrency`, `timeout` and `concurrentTimeout` must also retain their original values and explicit/default bindings: they participate in the fingerprint even though Spark tuning such as `spark.sql.shuffle.partitions` does not. To change transport settings, start a new run with a fresh `outputPath` and deliberately select the rows to submit; resume cannot retune pending requests or retry committed errors. Rebind scalar credentials; rotation does not change the fingerprint. Resume uses the saved input snapshot, never a replacement DataFrame. Do not include credential columns, remote media, file IDs, previous-response IDs or provider retrieval/tool state in a durable run. Incomplete initialization must be investigated, not treated as a completed empty result." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Deliberately not executable without an operator supplying the exact token:\n", + "# AIPrompt.releaseOutputOwner(spark, RUN_PATH,\n", + "# expectedOwnerToken='', confirmStopped=True)\n", + "\n", + "if RUN_RESUME:\n", + " resumed = configured_stage().resume(spark)\n", + " print(AIPrompt.getOutputStatus(spark, RUN_PATH))\n", + " resumed.select('ticket_id', 'answer', 'error').show(truncate=False)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Scale and operate\n", + "Before starting a run, set `spark.sql.shuffle.partitions` deliberately, for example `spark.conf.set('spark.sql.shuffle.partitions', '8')`, and tune per-partition `concurrency` to the model quota. Each durable batch is repartitioned to `min(batch row count, spark.sql.shuffle.partitions)`; repartitioning the original input does not control inference parallelism after snapshot storage. Approximate in-flight requests are active batch partitions times `concurrency`, limited by executor slots and available rows. The Spark partition setting can change on resume, but stage transport tuning is fixed for the lifetime of a durable run. A durable batch limits rows at risk before the next atomic commit; smaller batches improve recovery granularity but incur more Delta overhead. Only bounded pending keys are collected internally, not source rows. Repeated user rows keep distinct stored identities. Inspect committed outputs instead of calling `count` or `show` on a live inference plan before writing it.\n", + "\n", + "With the default handler, each active partition can create up to `concurrency` HTTP worker threads and buffer that many pending request/response entries. `concurrency` defaults to `1`; explicit positive values have no additional library ceiling or executor-wide clamp. Budget for all simultaneously running partitions on each executor, including inline-media/response memory and Spark's own threads. High values can exhaust container thread/process or memory limits and cause executor loss and repeated requests. Start at the default and measure before raising it; do not carry large legacy `OpenAIPrompt` settings over unchanged. Choose the value before the durable run starts because changing it requires a fresh run, not resume.\n", + "\n", + "The directory contains `_control/`, `_input/` and `results/`; it is not automatically a catalog table. Keep the complete run together for recovery. No automatic cleanup or schema evolution is performed.\n", + "\n", + "**Exclude `_input/` from all manual and automatic maintenance after initialization. Defer `results/` `OPTIMIZE` and `VACUUM` until status is complete and has no owner.** Even data-preserving results compaction between batches can stop inference and retain ownership. If that happens, follow the stop-confirmed release and resume procedure above; a failed cell alone is not proof of quiescence. Maintenance of completed results must preserve schema, row identities and outcomes, and respect Delta retention and readers of older snapshots. Never append, update or delete result rows; leave `_control/` managed by the run APIs. The `_input/` restriction still applies after completion. See [the guide](Durable%20AI%20results.md) for Scala APIs, vector storage and schema nullability." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "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..0561bdb970c --- /dev/null +++ b/docs/Explore Algorithms/OpenAI/Unified AI prompting.md @@ -0,0 +1,906 @@ +--- +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. + +| `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. + +Without `outputPath`, 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 in this mode 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. + +The separate [durable-results guide](Durable%20AI%20results.md) describes opt-in +Delta snapshots and explicit recovery. That draft follow-up remains unqualified; +it does not promise exactly-once requests or managed Spark application-restart support. + +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 without discarding other +usable inputs. 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`. +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/website/sidebars.js b/website/sidebars.js index dc8e05825af..ae1355730aa 100644 --- a/website/sidebars.js +++ b/website/sidebars.js @@ -49,6 +49,10 @@ module.exports = { items: [ "Explore Algorithms/OpenAI/Langchain", "Explore Algorithms/OpenAI/OpenAI", + "Explore Algorithms/OpenAI/Unified AI prompting", + "Explore Algorithms/OpenAI/AIPrompt", + "Explore Algorithms/OpenAI/Durable AI results", + "Explore Algorithms/OpenAI/Quickstart - Durable AI results", "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",