agAdvisor / src /utils /input_guard.py
tirtho149's picture
Deploy AgAdvisor
b30f068 verified
Raw
History Blame Contribute Delete
2 kB
"""Lightweight user-input hygiene for the public chat surface.
This is intentionally conservative: the app's tools are all read-only (RAG lookup,
weather, soil, curated web search) with no code-execution or arbitrary-URL-fetch
sink, so prompt injection cannot escalate to real actions here. The guard's job is
just to (a) bound length, (b) strip control characters, and (c) collapse the kind
of whitespace/newline flooding used to push a system prompt out of context — WITHOUT
mangling legitimate agronomy questions. It does not try to detect or rewrite
"ignore previous instructions"-style text; the LLM prompts already isolate the
user block and instruct the model to answer only from retrieved data.
"""
from __future__ import annotations
import re
import unicodedata
# Keep in sync with MAX_QUERY_CHARS in the Streamlit app.
DEFAULT_MAX_CHARS = 2000
# Control characters except tab/newline/carriage-return.
_CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
# 3+ consecutive blank lines -> at most 2 (defeats newline flooding).
_EXCESS_NEWLINES = re.compile(r"\n{3,}")
# Runs of the same non-word char repeated absurdly (e.g. 500 dashes).
_REPEAT_RUN = re.compile(r"([^\w\s])\1{40,}")
def sanitize_user_query(text: str, max_chars: int = DEFAULT_MAX_CHARS) -> str:
"""Normalize and bound a raw user query. Returns a safe-to-process string.
Never raises; on falsy input returns "".
"""
if not text:
return ""
# Normalize unicode (defeats homoglyph-based obfuscation, NFKC folds compatibility forms).
text = unicodedata.normalize("NFKC", text)
# Drop control chars that could corrupt logs/prompts.
text = _CONTROL_CHARS.sub("", text)
# Collapse flooding.
text = _EXCESS_NEWLINES.sub("\n\n", text)
text = _REPEAT_RUN.sub(lambda m: m.group(1) * 40, text)
# Trim leading/trailing whitespace and hard-cap length.
text = text.strip()
if len(text) > max_chars:
text = text[:max_chars].rstrip()
return text