From e6a9e6c92492b0ac74c860fb5dbc4c2fc1d0dec6 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Mon, 20 Jul 2026 15:27:57 +0000 Subject: [PATCH 1/2] Fix LLMSQLQueryOperator not stripping single-line markdown code fences _strip_llm_output only removed ```sql fences when the LLM's response spanned multiple lines. LLMs that wrap short queries in a fence on a single line (e.g. "```SELECT 1```") passed the fenced text straight into SQL validation/execution with the backticks still attached. Add a fallback branch for the single-line case, alongside the existing multi-line handling. --- .../airflow/providers/common/ai/operators/llm_sql.py | 5 ++++- .../ai/tests/unit/common/ai/operators/test_llm_sql.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index 7342be2b7e34f..9c80c0b11b618 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -177,10 +177,13 @@ def _strip_llm_output(raw: str) -> str: text = raw.strip() if text.startswith("```"): lines = text.split("\n") - # Remove opening fence (```sql, ```, etc.) and closing fence if len(lines) >= 2: + # Remove opening fence (```sql, ```, etc.) and closing fence end = -1 if lines[-1].strip().startswith("```") else len(lines) text = "\n".join(lines[1:end]).strip() + elif text.endswith("```") and len(text) > 6: + # Whole fenced block on one line, e.g. "```SELECT 1```" -> "SELECT 1" + text = text[3:-3].strip() return text def _get_schema_context(self) -> str: diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index 1862971c9539d..ec6b360a704ed 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py @@ -77,6 +77,16 @@ class TestStripLLMOutput: "SELECT 1", id="missing_closing_fence", ), + pytest.param( + "```SELECT 1```", + "SELECT 1", + id="single_line_fence_no_language_tag", + ), + pytest.param( + "```SELECT * FROM users LIMIT 10```", + "SELECT * FROM users LIMIT 10", + id="single_line_fence_with_query", + ), ), ) def test_strip_llm_output(self, raw, expected): From a79b2d47082df1a36e204480b258ee7a201ec3bf Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Mon, 20 Jul 2026 18:13:31 +0000 Subject: [PATCH 2/2] Also strip a same-line language tag on single-line fences Address review feedback: a single-line fence like "```sql SELECT 1```" left "sql " stuck to the front of the query, since the earlier fix only stripped the outer backticks. Drop the leading word too, but only when it's "sql" or the resolved dialect, so a real leading SQL keyword is never mistaken for a tag. --- .../providers/common/ai/operators/llm_sql.py | 15 +++++++--- .../unit/common/ai/operators/test_llm_sql.py | 28 +++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index 9c80c0b11b618..d850bf1137216 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -152,7 +152,7 @@ def execute(self, context: Context) -> str: ) result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) log_run_summary(self.log, result) - sql = self._strip_llm_output(result.output) + sql = self._strip_llm_output(result.output, dialect=self._resolved_dialect) if self.validate_sql: _validate_sql(sql, allowed_types=self.allowed_sql_types, dialect=self._resolved_dialect) @@ -172,7 +172,7 @@ def execute_complete(self, context: Context, generated_output: str, event: dict[ return output @staticmethod - def _strip_llm_output(raw: str) -> str: + def _strip_llm_output(raw: str, *, dialect: str | None = None) -> str: """Strip whitespace and markdown code fences from LLM output.""" text = raw.strip() if text.startswith("```"): @@ -182,8 +182,15 @@ def _strip_llm_output(raw: str) -> str: end = -1 if lines[-1].strip().startswith("```") else len(lines) text = "\n".join(lines[1:end]).strip() elif text.endswith("```") and len(text) > 6: - # Whole fenced block on one line, e.g. "```SELECT 1```" -> "SELECT 1" - text = text[3:-3].strip() + # Whole fenced block on one line, e.g. "```sql SELECT 1```" -> "SELECT 1". + # Only drop the leading word if it's a known tag, so a real keyword + # like "SELECT" is never mistaken for one. + inner = text[3:-3].strip() + tags = {"sql", *([dialect.lower()] if dialect else [])} + first_word, sep, rest = inner.partition(" ") + if sep and first_word.lower() in tags: + inner = rest.strip() + text = inner return text def _get_schema_context(self) -> str: diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index ec6b360a704ed..812c7f58a362f 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py @@ -87,11 +87,39 @@ class TestStripLLMOutput: "SELECT * FROM users LIMIT 10", id="single_line_fence_with_query", ), + pytest.param( + "```sql SELECT 1```", + "SELECT 1", + id="single_line_fence_with_language_tag", + ), + pytest.param( + "```SQL SELECT 1```", + "SELECT 1", + id="single_line_fence_with_uppercase_language_tag", + ), ), ) def test_strip_llm_output(self, raw, expected): assert LLMSQLQueryOperator._strip_llm_output(raw) == expected + @pytest.mark.parametrize( + ("raw", "dialect", "expected"), + ( + pytest.param("```postgres SELECT 1```", "postgres", "SELECT 1", id="dialect_tag_matches"), + pytest.param( + "```POSTGRES SELECT 1```", "postgres", "SELECT 1", id="dialect_tag_case_insensitive" + ), + pytest.param( + "```mysql SELECT 1```", + "postgres", + "mysql SELECT 1", + id="dialect_tag_mismatch_left_alone", + ), + ), + ) + def test_strip_llm_output_with_dialect(self, raw, dialect, expected): + assert LLMSQLQueryOperator._strip_llm_output(raw, dialect=dialect) == expected + class TestLLMSQLQueryOperator: def test_inherits_from_llm_operator(self):