sql_data_analyst / env /utils.py
YashashMathur's picture
SQL Data Analyst OpenEnv - Initial commit
d103a0f verified
raw
history blame
727 Bytes
import re
def normalize_answer(raw: str) -> str:
"""Remove common LLM answer preambles and formatting."""
text = raw.strip().lower()
text = re.sub(r"the (answer|result) is:?\s*", "", text)
text = re.sub(r"\*+", "", text)
text = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
text = re.sub(r"`[^`]+`", lambda m: m.group().strip("`"), text)
text = re.sub(r"\s+", " ", text)
return text.strip()
FORBIDDEN_KEYWORDS = [
"DROP",
"DELETE",
"INSERT",
"UPDATE",
"ALTER",
"CREATE",
"TRUNCATE",
]
def is_safe_query(query: str) -> bool:
"""Check if query is safe (SELECT-only)."""
upper = query.upper()
return not any(kw in upper for kw in FORBIDDEN_KEYWORDS)