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..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,15 +172,25 @@ 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("```"): 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. "```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 1862971c9539d..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 @@ -77,11 +77,49 @@ 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", + ), + 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):