From b23fbfdec4ea2d2a077f5ab3742dab6cb44b1cb4 Mon Sep 17 00:00:00 2001 From: PG1204 Date: Sun, 20 Sep 2026 17:04:26 -0700 Subject: [PATCH 1/3] fix(workflow-operator): degrade instead of raising on malformed chat responses --- .../huggingFace/codegen/QaRankingCodegen.scala | 12 ++++++------ .../huggingFace/codegen/TextGenCodegen.scala | 8 ++++++-- .../HuggingFaceInferenceOpDescSpec.scala | 16 ++++++++++++---- .../codegen/QaRankingCodegenSpec.scala | 14 +++++++++++++- .../huggingFace/codegen/TextGenCodegenSpec.scala | 12 ++++++++++++ 5 files changed, 49 insertions(+), 13 deletions(-) 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 bdc266080af..433d0558351 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 @@ -79,18 +79,18 @@ object QaRankingCodegen extends TaskCodegen { | 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"] + | if body.get("choices"): + | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) | return body.get("answer", json.dumps(body)) | return json.dumps(body) | elif task == "table-question-answering": | if isinstance(body, dict): - | if "choices" in body: - | return body["choices"][0]["message"]["content"] + | if body.get("choices"): + | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) | 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"] + | if isinstance(body, dict) and body.get("choices"): + | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) | 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 b836de9e121..ac41d2fb583 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 the raw JSON body when a provider returns a + * shape that does not carry it — parsing runs per row, so raising here + * would abort the whole run over one malformed response. */ object TextGenCodegen extends TaskCodegen { @@ -50,5 +52,7 @@ object TextGenCodegen extends TaskCodegen { override def parsePython(ctx: CodegenContext): String = """ if task == "text-generation": - | return body["choices"][0]["message"]["content"]""".stripMargin + | if isinstance(body, dict) and body.get("choices"): + | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) + | return json.dumps(body)""".stripMargin } 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 56b4b65ebf8..8d9e267f915 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( + """body["choices"][0].get("message", {}).get("content", json.dumps(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( + """body["choices"][0].get("message", {}).get("content", json.dumps(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( + """body["choices"][0].get("message", {}).get("content", json.dumps(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( + """body["choices"][0].get("message", {}).get("content", json.dumps(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/QaRankingCodegenSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala index c7c3c18a91a..210f2a06ed7 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("""body["choices"][0].get("message", {}).get("content", json.dumps(body))""") + } + + it should "degrade instead of raising when a chat response is malformed (#8486)" in { + // parsePython runs per row, so indexing straight into + // choices[0]["message"]["content"] turned one malformed provider response + // into an aborted run: an empty "choices" list raises IndexError and a + // choice missing "message"/"content" raises KeyError. 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("""body\.get\("choices"\)""").length - 1 shouldBe 3 } it should "return the raw JSON body for the ranking-style tasks" in { 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 d7714af0672..525dedaa057 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 @@ -74,6 +74,18 @@ class TextGenCodegenSpec extends AnyFlatSpec with Matchers { out should include("content") } + 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, and since parsing + // runs per row that aborted the whole run over one bad response. It now + // falls back to the raw JSON body, as the other codegens do. + val out = TextGenCodegen.parsePython(makeCtx()) + out should include("""if isinstance(body, dict) and body.get("choices"):""") + out should include("""body["choices"][0].get("message", {}).get("content", json.dumps(body))""") + out should include("return json.dumps(body)") + out should not include ("""["message"]["content"]""") + } + "TextGenCodegen snippets" should "never inline raw CodegenContext string values" in { // The snippets must reference self.* attributes — the base class decodes // user-supplied strings safely at runtime. Sentinel values chosen to be From 54ffe68fccbaba5407f9537675aeebb3848b6fca Mon Sep 17 00:00:00 2001 From: PG1204 Date: Mon, 21 Sep 2026 20:25:26 -0700 Subject: [PATCH 2/3] fix(workflow-operator): type-check chat responses before extracting content --- .../codegen/HuggingFaceCodegenBase.scala | 31 +++ .../codegen/QaRankingCodegen.scala | 19 +- .../huggingFace/codegen/TextGenCodegen.scala | 5 +- .../test/resources/python/hf_parse_probe.py | 73 +++++++ .../HuggingFaceInferenceOpDescSpec.scala | 8 +- .../codegen/HuggingFaceCodegenBaseSpec.scala | 15 ++ .../HuggingFaceParseBehaviorSpec.scala | 196 ++++++++++++++++++ .../codegen/QaRankingCodegenSpec.scala | 11 +- .../codegen/TextGenCodegenSpec.scala | 22 +- 9 files changed, 358 insertions(+), 22 deletions(-) create mode 100644 common/workflow-operator/src/test/resources/python/hf_parse_probe.py create mode 100644 common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceParseBehaviorSpec.scala 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 a40d4aeac70..31ac69a81dc 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,37 @@ 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: parsing + | runs once per row, and an exception here aborts the whole run. 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 433d0558351..d4b4da95f17 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 body.get("choices"): - | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) | 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 body.get("choices"): - | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) | 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 body.get("choices"): - | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) + | 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 ac41d2fb583..b1bf5bb061e 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 @@ -52,7 +52,8 @@ object TextGenCodegen extends TaskCodegen { override def parsePython(ctx: CodegenContext): String = """ if task == "text-generation": - | if isinstance(body, dict) and body.get("choices"): - | return body["choices"][0].get("message", {}).get("content", json.dumps(body)) + | 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 00000000000..3635d86fb3c --- /dev/null +++ b/common/workflow-operator/src/test/resources/python/hf_parse_probe.py @@ -0,0 +1,73 @@ +# +# 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": , "task": , +"bodies": [, ...]} on stdin and writes +{"results": [{"value": } | {"raised": ""}, ...]} 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 + +WANTED = ("_parse_response", "_chat_message_content", "_url_to_data_url", "_format_error") + + +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, "", "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 8d9e267f915..5592944018b 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 @@ -112,7 +112,7 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { code should include("self.TEMPERATURE") // Parse — text-gen pulls choices[0].message.content out of the response. code should include( - """body["choices"][0].get("message", {}).get("content", json.dumps(body))""" + """content = self._chat_message_content(body)""" ) } @@ -231,7 +231,7 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { ) TextGenCodegen.payloadPython(ctx) should include("self.MODEL_ID") TextGenCodegen.parsePython(ctx) should include( - """body["choices"][0].get("message", {}).get("content", json.dumps(body))""" + """content = self._chat_message_content(body)""" ) } @@ -573,7 +573,7 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { 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].get("message", {}).get("content", json.dumps(body))""" + """content = self._chat_message_content(body)""" ) } @@ -584,7 +584,7 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { 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].get("message", {}).get("content", json.dumps(body))""" + """content = self._chat_message_content(body)""" ) } 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 0a5e96f3307..4520b68275c 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 @@ -144,4 +144,19 @@ class HuggingFaceCodegenBaseSpec extends AnyFlatSpec with Matchers { out should not include "MARKER_TASK_zXyq42" out should not include "MARKER_SYSTEM_zXyq42" } + + // 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 00000000000..61c9a39a093 --- /dev/null +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceParseBehaviorSpec.scala @@ -0,0 +1,196 @@ +/* + * 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))) + + val process = new ProcessBuilder(exe, probe.toString).redirectErrorStream(false).start() + process.getOutputStream.write(mapper.writeValueAsBytes(request)) + process.getOutputStream.close() + val out = new String(process.getInputStream.readAllBytes(), StandardCharsets.UTF_8) + val err = new String(process.getErrorStream.readAllBytes(), StandardCharsets.UTF_8) + if (!process.waitFor(60, TimeUnit.SECONDS)) { + process.destroyForcibly(); fail("probe timed out") + } + withClue(s"probe stderr: $err\nprobe stdout: $out\n") { process.exitValue() shouldBe 0 } + + 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) + withClue(s"task=$task ") { + outcomes.foreach { outcome => + outcome should not startWith "RAISED:" + outcome should not be "RETURNED_NONE" + outcome should not startWith "NON_STRING:" + } + } + } + } + + 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 210f2a06ed7..7c59e64159c 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,7 @@ 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].get("message", {}).get("content", json.dumps(body))""") + out should include("""content = self._chat_message_content(body)""") } it should "degrade instead of raising when a chat response is malformed (#8486)" in { @@ -129,7 +129,7 @@ class QaRankingCodegenSpec extends AnyFlatSpec with Matchers { // native shapes beside them, which already degrade via json.dumps(body). val out = QaRankingCodegen.parsePython(makeCtx()) out should not include ("""["message"]["content"]""") - out.split("""body\.get\("choices"\)""").length - 1 shouldBe 3 + 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 { @@ -209,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 525dedaa057..92d3cf3eabc 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,14 @@ 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 { @@ -80,8 +83,8 @@ class TextGenCodegenSpec extends AnyFlatSpec with Matchers { // runs per row that aborted the whole run over one bad response. It now // falls back to the raw JSON body, as the other codegens do. val out = TextGenCodegen.parsePython(makeCtx()) - out should include("""if isinstance(body, dict) and body.get("choices"):""") - out should include("""body["choices"][0].get("message", {}).get("content", json.dumps(body))""") + 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"]""") } @@ -142,4 +145,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]""") + } } From 1fe91a67d5487d52aafa29acb6a07e6300c1d268 Mon Sep 17 00:00:00 2001 From: PG1204 Date: Mon, 21 Sep 2026 20:41:43 -0700 Subject: [PATCH 3/3] test(workflow-operator): lift only the methods the parse probe can reach --- .../src/test/resources/python/hf_parse_probe.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 index 3635d86fb3c..a15fefc2995 100644 --- a/common/workflow-operator/src/test/resources/python/hf_parse_probe.py +++ b/common/workflow-operator/src/test/resources/python/hf_parse_probe.py @@ -32,7 +32,10 @@ import re import sys -WANTED = ("_parse_response", "_chat_message_content", "_url_to_data_url", "_format_error") +# 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):