This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch 
gh-readonly-queue/main/pr-8617-82ab13483dff640e580850d6f534964ae2d68029
in repository https://gitbox.apache.org/repos/asf/texera.git

commit 09a04ceb75592a66a41d0f06b6ea6084427bcd64
Author: Prateek Ganigi <[email protected]>
AuthorDate: Mon Oct 5 19:28:34 2026 +0000

    fix(workflow-operator): degrade instead of raising on malformed chat 
responses (#8617)
    
    ### What changes were proposed in this PR?
    
    The text-generation and QA/ranking codegens read a chat provider's reply
    by indexing straight into it:
    
        return body["choices"][0]["message"]["content"]
    
    That assumes the response is a dict, that `choices` is non-empty, and
    that the first choice carries `message.content`. When a provider returns
    a 200 with anything else, the indexing raises, and since the parse step
    runs once per row, a single malformed response aborts the entire run.
    `TextGenCodegen` was the most exposed: it had no `isinstance` check
    either, so even a plain list body raised.
    
    All four extractions (one in `TextGenCodegen`, three in
    `QaRankingCodegen`) now guard on a truthy `choices` and use `.get`
    chaining, falling back to the raw JSON body. That is the same idiom
    `ImageTaskCodegen` adopted in #7920, and it matches how the native
    `hf-inference` shapes beside it already degrade via `body.get("answer",
    json.dumps(body))`. A bad response now lands in the result column as raw
    JSON instead of killing the run.
    
    ### Any related issues?
    
    Closes #8486
    
    ### How was this PR tested?
    
    145 tests pass in the Hugging Face suites,
    `PythonCodeRawInvalidTextSpec` py-compiles the generated Python for all
    117 operators, and `scalafmtCheck` and `scalafix --check` are clean.
    Added a degradation test to `TextGenCodegenSpec` and
    `QaRankingCodegenSpec`, and updated the assertions in
    `HuggingFaceInferenceOpDescSpec` that pinned the old form.
    
    The emitted Python was also run directly against the failing shapes,
    comparing this branch with main across 33 task/response combinations: 16
    of them raised before (`IndexError` on an empty `choices` list,
    `KeyError` on a missing `message`/`content`, `TypeError` on a list body)
    and none raise now. Well-formed chat responses and native `hf-inference`
    responses parse identically to before.
    
    ### Was this PR authored or co-authored using generative AI tooling?
    
    Yes, this PR was co-authored with Claude in compliance with ASF policy.
---
 .../codegen/HuggingFaceCodegenBase.scala           |  33 ++++
 .../huggingFace/codegen/QaRankingCodegen.scala     |  19 +-
 .../huggingFace/codegen/TextGenCodegen.scala       |   9 +-
 .../src/test/resources/python/hf_parse_probe.py    |  76 +++++++
 .../HuggingFaceInferenceOpDescSpec.scala           |  16 +-
 .../codegen/HuggingFaceCodegenBaseSpec.scala       |  15 ++
 .../codegen/HuggingFaceParseBehaviorSpec.scala     | 220 +++++++++++++++++++++
 .../huggingFace/codegen/QaRankingCodegenSpec.scala |  21 +-
 .../huggingFace/codegen/TextGenCodegenSpec.scala   |  31 ++-
 9 files changed, 421 insertions(+), 19 deletions(-)

diff --git 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala
 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala
index 1db86f5ab1..81a4b58a62 100644
--- 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala
+++ 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala
@@ -260,6 +260,39 @@ object HuggingFaceCodegenBase {
        |        summary = "; ".join(errors) if errors else "no providers 
available"
        |        return last_resp, summary
        |
+       |    def _chat_message_content(self, body):
+       |        '''Return the assistant text from a chat-completions response, 
or None
+       |        when the body is not that shape. Providers differ and 
malformed 200s
+       |        happen, so every level is type-checked rather than indexed: 
indexing
+       |        a non-dict raises AttributeError, which _parse_response's
+       |        except (KeyError, IndexError, TypeError) does not catch, so 
that row
+       |        ends up reading "Request failed" instead of the body. Callers 
fall
+       |        back to their native shape, or to json.dumps(body), on None.
+       |        '''
+       |        if not isinstance(body, dict):
+       |            return None
+       |        choices = body.get("choices")
+       |        if not isinstance(choices, list) or not choices:
+       |            return None
+       |        first = choices[0]
+       |        if not isinstance(first, dict):
+       |            return None
+       |        message = first.get("message")
+       |        if not isinstance(message, dict):
+       |            return None
+       |        content = message.get("content")
+       |        if isinstance(content, str):
+       |            return content
+       |        if isinstance(content, list):
+       |            # Some OpenAI-compatible providers return content as a 
list of
+       |            # parts ({"type": "text", "text": ...}); join the text of 
those.
+       |            parts = [
+       |                part["text"] for part in content
+       |                if isinstance(part, dict) and 
isinstance(part.get("text"), str)
+       |            ]
+       |            return "".join(parts) if parts else None
+       |        return None
+       |
        |    def _chat_content_for_task(self, pipeline_payload, prompt_value):
        |        '''Reformulate a structured task (question-answering,
        |        table-question-answering, zero-shot-classification,
diff --git 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegen.scala
 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegen.scala
index bdc266080a..d4b4da95f1 100644
--- 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegen.scala
+++ 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegen.scala
@@ -76,21 +76,24 @@ object QaRankingCodegen extends TaskCodegen {
 
   override def parsePython(ctx: CodegenContext): String =
     """            if task == "question-answering":
+      |                # Third-party chat providers answer via 
choices[0].message;
+      |                # hf-inference returns the native {"answer": ...} shape.
+      |                content = self._chat_message_content(body)
+      |                if content is not None:
+      |                    return content
       |                if isinstance(body, dict):
-      |                    # Third-party chat providers answer via 
choices[0].message;
-      |                    # hf-inference returns the native {"answer": ...} 
shape.
-      |                    if "choices" in body:
-      |                        return body["choices"][0]["message"]["content"]
       |                    return body.get("answer", json.dumps(body))
       |                return json.dumps(body)
       |            elif task == "table-question-answering":
+      |                content = self._chat_message_content(body)
+      |                if content is not None:
+      |                    return content
       |                if isinstance(body, dict):
-      |                    if "choices" in body:
-      |                        return body["choices"][0]["message"]["content"]
       |                    return body.get("answer", json.dumps(body))
       |                return json.dumps(body)
       |            elif task in ("zero-shot-classification", 
"sentence-similarity", "text-ranking"):
-      |                if isinstance(body, dict) and "choices" in body:
-      |                    return body["choices"][0]["message"]["content"]
+      |                content = self._chat_message_content(body)
+      |                if content is not None:
+      |                    return content
       |                return json.dumps(body)""".stripMargin
 }
diff --git 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegen.scala
 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegen.scala
index b836de9e12..b1650d5aab 100644
--- 
a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegen.scala
+++ 
b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegen.scala
@@ -28,7 +28,9 @@ package org.apache.texera.amber.operator.huggingFace.codegen
   * (Cerebras, Groq, Sambanova, Together, …) accepts.
   *
   * The parse step pulls `body["choices"][0]["message"]["content"]` out of
-  * the response.
+  * the response, degrading to raw JSON when a provider returns a shape that
+  * does not carry readable text. This keeps the result column string-valued
+  * for malformed or non-string chat content.
   */
 object TextGenCodegen extends TaskCodegen {
 
@@ -50,5 +52,8 @@ object TextGenCodegen extends TaskCodegen {
 
   override def parsePython(ctx: CodegenContext): String =
     """            if task == "text-generation":
-      |                return 
body["choices"][0]["message"]["content"]""".stripMargin
+      |                content = self._chat_message_content(body)
+      |                if content is not None:
+      |                    return content
+      |                return json.dumps(body)""".stripMargin
 }
diff --git 
a/common/workflow-operator/src/test/resources/python/hf_parse_probe.py 
b/common/workflow-operator/src/test/resources/python/hf_parse_probe.py
new file mode 100644
index 0000000000..a15fefc299
--- /dev/null
+++ b/common/workflow-operator/src/test/resources/python/hf_parse_probe.py
@@ -0,0 +1,76 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+"""Run the generated Hugging Face response parser against real response shapes.
+
+Reads {"source": <generated operator Python>, "task": <task tag>,
+"bodies": [<parsed response>, ...]} on stdin and writes
+{"results": [{"value": <cell text>} | {"raised": "<ExcType>"}, ...]} on stdout.
+
+The generated module cannot be imported as-is (it depends on the pytexera
+runtime), so the parser and its helpers are lifted out of the source and
+executed on their own. That keeps the assertion about the emitted code itself
+rather than a transcription of it.
+"""
+import io
+import json
+import re
+import sys
+
+# Exactly what _parse_response can reach: itself, the chat-content helper, and
+# the data-URL helper used by the image-to-image branch. Anything else in the
+# generated class belongs to the request loop, not to parsing.
+WANTED = ("_parse_response", "_chat_message_content", "_url_to_data_url")
+
+
+def lift(source, name):
+    """Return the text of a 4-space-indented method, or '' when absent."""
+    out, started = [], False
+    for line in source.split("\n"):
+        stripped = line.strip()
+        if stripped.startswith("def " + name + "("):
+            started = True
+        elif started and (line.startswith("    def ") or (line and not 
line.startswith(" "))):
+            break
+        if started:
+            out.append(line)
+    return "\n".join(out)
+
+
+def main():
+    request = json.load(sys.stdin)
+    methods = [m for m in (lift(request["source"], n) for n in WANTED) if 
m.strip()]
+    module = "import json\nclass Parser:\n    TASK = %r\n%s\n" % (
+        request["task"],
+        "\n".join(methods),
+    )
+    namespace = {}
+    exec(compile(module, "<generated>", "exec"), namespace)  # noqa: S102 - 
the point of the probe
+    parser = namespace["Parser"]()
+
+    results = []
+    for body in request["bodies"]:
+        try:
+            results.append({"value": parser._parse_response(body)})
+        except Exception as exc:  # noqa: BLE001 - reporting the type is the 
assertion
+            results.append({"raised": type(exc).__name__})
+    json.dump({"results": results}, sys.stdout)
+
+
+if __name__ == "__main__":
+    main()
diff --git 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala
 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala
index 56b4b65ebf..5592944018 100644
--- 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala
+++ 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala
@@ -111,7 +111,9 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec 
with Matchers {
     code should include("self.MAX_NEW_TOKENS")
     code should include("self.TEMPERATURE")
     // Parse — text-gen pulls choices[0].message.content out of the response.
-    code should include("""body["choices"][0]["message"]["content"]""")
+    code should include(
+      """content = self._chat_message_content(body)"""
+    )
   }
 
   it should "send the provider-specific model id on provider-scoped chat 
routes" in {
@@ -228,7 +230,9 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec 
with Matchers {
       safeTemp = 0.0
     )
     TextGenCodegen.payloadPython(ctx) should include("self.MODEL_ID")
-    TextGenCodegen.parsePython(ctx) should 
include("""body["choices"][0]["message"]["content"]""")
+    TextGenCodegen.parsePython(ctx) should include(
+      """content = self._chat_message_content(body)"""
+    )
   }
 
   "image task family" should
@@ -568,7 +572,9 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec 
with Matchers {
     code should include("Context column")
     code should include("""payload = {"inputs": {"question": prompt_value, 
"context": ctx_val}}""")
     code should include("""body.get("answer", json.dumps(body))""")
-    code should include("""body["choices"][0]["message"]["content"]""")
+    code should include(
+      """content = self._chat_message_content(body)"""
+    )
   }
 
   it should "route table-question-answering with a precomputed table payload" 
in {
@@ -577,7 +583,9 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec 
with Matchers {
     code should include("table_dict = {}")
     code should include("""payload = {"inputs": {"query": prompt_value, 
"table": table_dict}}""")
     code should include("""body.get("answer", json.dumps(body))""")
-    code should include("""body["choices"][0]["message"]["content"]""")
+    code should include(
+      """content = self._chat_message_content(body)"""
+    )
   }
 
   it should "route zero-shot-classification with candidate labels" in {
diff --git 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBaseSpec.scala
 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBaseSpec.scala
index 9acae7db85..f293515dae 100644
--- 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBaseSpec.scala
+++ 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBaseSpec.scala
@@ -182,4 +182,19 @@ class HuggingFaceCodegenBaseSpec extends AnyFlatSpec with 
Matchers {
     */
   private def chatContentHelper(rendered: String): String =
     rendered.split("def _chat_content_for_task")(1).split("    def ")(0)
+
+  // Review feedback on #8617 (@Copilot): a truthy `choices` is not enough — 
the
+  // value may not be a list, its first item may not be a dict, and `message` 
or
+  // `content` may be null or a list of parts. Every level is type-checked in 
one
+  // shared helper so all eight chat extractions degrade identically.
+  it should "emit a type-checked chat-content helper" in {
+    val out = HuggingFaceCodegenBase.render(makeCtx(), StubCodegen)
+    out should include("def _chat_message_content(self, body):")
+    val helper = out.split("def _chat_message_content")(1).split("    def ")(0)
+    helper should include("isinstance(choices, list)")
+    helper should include("isinstance(first, dict)")
+    helper should include("isinstance(message, dict)")
+    helper should include("isinstance(content, str)")
+    helper should include("isinstance(content, list)")
+  }
 }
diff --git 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceParseBehaviorSpec.scala
 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceParseBehaviorSpec.scala
new file mode 100644
index 0000000000..ac45bead3d
--- /dev/null
+++ 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceParseBehaviorSpec.scala
@@ -0,0 +1,220 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.operator.huggingFace.codegen
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.typesafe.config.ConfigFactory
+import org.apache.texera.amber.pybuilder.PyStringTypes.EncodableString
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, Path}
+import java.util.concurrent.TimeUnit
+import scala.jdk.CollectionConverters._
+import scala.util.Try
+
+/**
+  * Executes the *generated* response parser instead of matching its source 
text.
+  *
+  * Every other Hugging Face spec asserts on the emitted Python as a string, 
so a
+  * parser that raises at runtime still passes them (review feedback on #8617).
+  * This spec renders the operator, lifts `_parse_response` and its helpers 
out of
+  * the result, and runs them under a real interpreter against the response 
shapes
+  * providers actually return — valid, empty, missing, null and wrong-typed.
+  *
+  * The contract asserted here: `_parse_response` always returns a string and
+  * never raises, because an uncaught exception is caught by the per-row 
handler
+  * and replaces the cell with "Request failed" instead of the raw body.
+  */
+class HuggingFaceParseBehaviorSpec extends AnyFlatSpec with Matchers {
+
+  private val mapper = new ObjectMapper()
+
+  private def makeCtx(task: EncodableString): CodegenContext =
+    CodegenContext(
+      hfApiToken = "token",
+      modelId = "Qwen/Qwen2.5-72B-Instruct",
+      promptColumn = "prompt",
+      resultColumn = "hf_response",
+      task = task,
+      systemPrompt = "You are a helpful assistant.",
+      safeMaxTokens = 256,
+      safeTemp = 0.7
+    )
+
+  /** Same resolution order as PythonCodeRawInvalidTextSpec: udf.conf, then 
PATH. */
+  private lazy val pythonExe: Option[String] = {
+    def fromConfig: Option[String] =
+      Try(ConfigFactory.parseResources("udf.conf").resolve()).toOption
+        .orElse(Try(ConfigFactory.load()).toOption)
+        .flatMap(c => Try(c.getConfig("python").getString("path")).toOption)
+        .map(_.trim)
+        .filter(_.nonEmpty)
+
+    def isRunnable(exe: String): Boolean =
+      Try(new ProcessBuilder(exe, 
"--version").redirectErrorStream(true).start()).toOption
+        .exists { p =>
+          if (!p.waitFor(5, TimeUnit.SECONDS)) { p.destroyForcibly(); false }
+          else p.exitValue() == 0
+        }
+
+    (fromConfig.toList ++ List("python3", "python", 
"py")).distinct.find(isRunnable)
+  }
+
+  private lazy val probe: Path = {
+    val script =
+      scala.io.Source.fromInputStream(
+        getClass.getResourceAsStream("/python/hf_parse_probe.py"),
+        "UTF-8"
+      )
+    val body =
+      try script.mkString
+      finally script.close()
+    val file = Files.createTempFile("hf_parse_probe", ".py")
+    Files.write(file, body.getBytes(StandardCharsets.UTF_8))
+    file.toFile.deleteOnExit()
+    file
+  }
+
+  /** Runs the generated parser for `task` over `bodies`; returns one outcome 
each. */
+  private def parseAll(codegen: TaskCodegen, task: String, bodies: 
Seq[String]): Seq[String] = {
+    val exe = pythonExe.getOrElse(cancel("no python interpreter available to 
run the probe"))
+    val request = mapper.createObjectNode()
+    request.put("source", HuggingFaceCodegenBase.render(makeCtx(task), 
codegen))
+    request.put("task", task)
+    val arr = request.putArray("bodies")
+    bodies.foreach(b => arr.add(mapper.readTree(b)))
+
+    // Redirect to files rather than pipes: reading a pipe blocks until the 
child
+    // closes it, so a probe that hung would never reach the timeout below, and
+    // waiting first with the output still in a pipe would deadlock instead 
once
+    // it filled. Files bound neither side, so the timeout is always reached.
+    val outFile = Files.createTempFile("hf_parse_probe_out", ".json")
+    val errFile = Files.createTempFile("hf_parse_probe_err", ".txt")
+    try {
+      val process = new ProcessBuilder(exe, probe.toString)
+        .redirectOutput(outFile.toFile)
+        .redirectError(errFile.toFile)
+        .start()
+      process.getOutputStream.write(mapper.writeValueAsBytes(request))
+      process.getOutputStream.close()
+      if (!process.waitFor(60, TimeUnit.SECONDS)) {
+        process.destroyForcibly()
+        fail("probe timed out")
+      }
+      val out = new String(Files.readAllBytes(outFile), StandardCharsets.UTF_8)
+      val err = new String(Files.readAllBytes(errFile), StandardCharsets.UTF_8)
+      withClue(s"probe stderr: $err\nprobe stdout: $out\n") { 
process.exitValue() shouldBe 0 }
+      readResults(out)
+    } finally {
+      Files.deleteIfExists(outFile)
+      Files.deleteIfExists(errFile)
+    }
+  }
+
+  /** Turns the probe's JSON into one outcome string per body. */
+  private def readResults(out: String): Seq[String] = {
+
+    mapper
+      .readTree(out)
+      .get("results")
+      .elements()
+      .asScala
+      .map(n =>
+        if (n.has("raised")) s"RAISED:${n.get("raised").asText()}"
+        else if (n.get("value").isNull) "RETURNED_NONE"
+        else if (!n.get("value").isTextual) 
s"NON_STRING:${n.get("value").getNodeType}"
+        else n.get("value").asText()
+      )
+      .toSeq
+  }
+
+  /** Shapes seen from providers: valid, empty, missing, null and wrong-typed. 
*/
+  private val malformed = Seq(
+    """{"choices": []}""",
+    """{"choices": [null]}""",
+    """{"choices": [{"message": null}]}""",
+    """{"choices": "bad"}""",
+    """{"choices": {"a": 1}}""",
+    """{"choices": [42]}""",
+    """{"choices": [{"message": {}}]}""",
+    """{"choices": [{"message": {"content": null}}]}""",
+    """{"choices": [{"message": {"content": 42}}]}""",
+    """{"choices": [{"no_message": true}]}""",
+    """{}""",
+    """[]""",
+    """[{"generated_text": "x"}]"""
+  )
+
+  // Scoped to the codegens this PR touches. ImageTaskCodegen has the same chat
+  // extractions (merged in #7920) and will be routed through the helper and 
added
+  // here in a follow-up, so that change arrives with its own failing-first 
test.
+  private val cases = Seq(
+    (TextGenCodegen: TaskCodegen, "text-generation"),
+    (QaRankingCodegen, "question-answering"),
+    (QaRankingCodegen, "table-question-answering"),
+    (QaRankingCodegen, "zero-shot-classification"),
+    (QaRankingCodegen, "sentence-similarity"),
+    (QaRankingCodegen, "text-ranking")
+  )
+
+  "The generated parser" should "return the raw body, never raise, on 
malformed chat responses" in {
+    cases.foreach {
+      case (codegen, task) =>
+        val outcomes = parseAll(codegen, task, malformed)
+        outcomes.zip(malformed).foreach {
+          case (outcome, input) =>
+            withClue(s"task=$task body=$input ") {
+              outcome should not startWith "RAISED:"
+              outcome should not be "RETURNED_NONE"
+              outcome should not startWith "NON_STRING:"
+              // Pin the actual fallback, not merely "some string": the parser 
must
+              // hand back the body it was given, serialized.
+              mapper.readTree(outcome) shouldBe mapper.readTree(input)
+            }
+        }
+    }
+  }
+
+  it should "return the assistant text for a well-formed chat response" in {
+    val body = """{"choices": [{"message": {"content": "the answer"}}]}"""
+    cases.foreach {
+      case (codegen, task) =>
+        withClue(s"task=$task ") {
+          parseAll(codegen, task, Seq(body)) shouldBe Seq("the answer")
+        }
+    }
+  }
+
+  it should "join a content list returned as parts" in {
+    val body =
+      """{"choices": [{"message": {"content": [{"type": "text", "text": "a "},
+        |{"type": "text", "text": "b"}]}}]}""".stripMargin.replace("\n", "")
+    parseAll(TextGenCodegen, "text-generation", Seq(body)) shouldBe Seq("a b")
+  }
+
+  it should "keep the native hf-inference shapes working" in {
+    parseAll(QaRankingCodegen, "question-answering", Seq("""{"answer": 
"Ada"}""")) shouldBe
+      Seq("Ada")
+    parseAll(QaRankingCodegen, "text-ranking", Seq("""[{"index": 0}]""")) 
shouldBe
+      Seq("""[{"index": 0}]""")
+  }
+}
diff --git 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
index c7c3c18a91..1f03c330bd 100644
--- 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
+++ 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
@@ -117,7 +117,19 @@ class QaRankingCodegenSpec extends AnyFlatSpec with 
Matchers {
     out should include("""body.get("answer"""")
     // #7195: chat-completions responses (third-party providers) are read from
     // choices[0].message.content, not the native {"answer": ...} shape.
-    out should include("""body["choices"][0]["message"]["content"]""")
+    out should include("""content = self._chat_message_content(body)""")
+  }
+
+  it should "degrade instead of raising when a chat response is malformed 
(#8486)" in {
+    // Indexing straight into choices[0]["message"]["content"] raises on a
+    // malformed provider response. _parse_response catches KeyError, 
IndexError
+    // and TypeError, but an AttributeError from indexing a non-dict escapes it
+    // and that row's cell becomes "Request failed" instead. All three chat
+    // extractions now use a truthiness guard plus .get chaining, matching the
+    // native shapes beside them, which already degrade via json.dumps(body).
+    val out = QaRankingCodegen.parsePython(makeCtx())
+    out should not include ("""["message"]["content"]""")
+    out.split("""content = self\._chat_message_content\(body\)""").length - 1 
shouldBe 3
   }
 
   it should "return the raw JSON body for the ranking-style tasks" in {
@@ -197,4 +209,11 @@ class QaRankingCodegenSpec extends AnyFlatSpec with 
Matchers {
     QaRankingCodegen.payloadPython(ctxA) shouldBe 
QaRankingCodegen.payloadPython(ctxB)
     QaRankingCodegen.parsePython(ctxA) shouldBe 
QaRankingCodegen.parsePython(ctxB)
   }
+
+  it should "read chat content through the shared type-checked helper (#8617 
review)" in {
+    val out = QaRankingCodegen.parsePython(makeCtx())
+    out should include("self._chat_message_content(body)")
+    // No direct index/get chaining survives — the helper owns that logic.
+    out should not include ("""body["choices"][0]""")
+  }
 }
diff --git 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegenSpec.scala
 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegenSpec.scala
index d7714af067..b3fc23c01b 100644
--- 
a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegenSpec.scala
+++ 
b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/TextGenCodegenSpec.scala
@@ -67,11 +67,27 @@ class TextGenCodegenSpec extends AnyFlatSpec with Matchers {
     out should include("""payload = {"inputs": prompt_value}""")
   }
 
-  "TextGenCodegen.parsePython" should "pull text out of 
choices[0].message.content" in {
+  "TextGenCodegen.parsePython" should "pull the assistant text out of the chat 
response" in {
+    // The chat shape (choices[0].message.content) is decoded by the shared,
+    // type-checked _chat_message_content helper in HuggingFaceCodegenBase; 
this
+    // snippet's job is to call it and fall back to the raw body.
     val out = TextGenCodegen.parsePython(makeCtx())
-    out should include("choices")
-    out should include("message")
-    out should include("content")
+    out should include("content = self._chat_message_content(body)")
+    out should include("if content is not None:")
+    out should include("return json.dumps(body)")
+  }
+
+  it should "degrade instead of raising when a chat response is malformed 
(#8486)" in {
+    // This extraction was fully unguarded: a non-dict body, an empty "choices"
+    // list, or a choice missing "message"/"content" raised. _parse_response
+    // catches KeyError, IndexError and TypeError, but an AttributeError from
+    // indexing a non-dict escapes it and that row reads "Request failed"
+    // instead. It now falls back to the raw JSON body, as the other codegens 
do.
+    val out = TextGenCodegen.parsePython(makeCtx())
+    out should include("""content = self._chat_message_content(body)""")
+    out should include("if content is not None:")
+    out should include("return json.dumps(body)")
+    out should not include ("""["message"]["content"]""")
   }
 
   "TextGenCodegen snippets" should "never inline raw CodegenContext string 
values" in {
@@ -130,4 +146,11 @@ class TextGenCodegenSpec extends AnyFlatSpec with Matchers 
{
     TextGenCodegen.payloadPython(ctxA) shouldBe 
TextGenCodegen.payloadPython(ctxB)
     TextGenCodegen.parsePython(ctxA) shouldBe TextGenCodegen.parsePython(ctxB)
   }
+
+  it should "read chat content through the shared type-checked helper (#8617 
review)" in {
+    val out = TextGenCodegen.parsePython(makeCtx())
+    out should include("self._chat_message_content(body)")
+    // No direct index/get chaining survives — the helper owns that logic.
+    out should not include ("""body["choices"][0]""")
+  }
 }

Reply via email to