ColtenOuO commented on code in PR #70137:
URL: https://github.com/apache/airflow/pull/70137#discussion_r3616562263
##########
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()
Review Comment:
I think this defensive extension is worth handling — updated
`_strip_llm_output` to also strip a same-line language tag on the single-line
path — but only when the leading word is literally `sql` or the resolved
dialect (the same value used in the system prompt's `dialect_label`), so a real
query starting with a SQL keyword is never mistaken for a tag.
Added test cases for `sql`, uppercase `SQL`, dialect match, and dialect
mismatch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
