diff --git a/mlflow/genai/prompts/utils.py b/mlflow/genai/prompts/utils.py index cb0debe49f58d..6a8c52dfb9962 100644 --- a/mlflow/genai/prompts/utils.py +++ b/mlflow/genai/prompts/utils.py @@ -1,4 +1,5 @@ import re +from functools import lru_cache from typing import Any @@ -8,5 +9,11 @@ def format_prompt(prompt: str, **values: Any) -> str: # Escape backslashes in the replacement string to prevent re.sub from interpreting # them as escape sequences (e.g. \u being treated as Unicode escape) replacement = str(value).replace("\\", "\\\\") - prompt = re.sub(r"\{\{\s*" + key + r"\s*\}\}", replacement, prompt) + prompt = _compiled_key_pattern(key).sub(replacement, prompt) return prompt + + +@lru_cache(maxsize=64) +def _compiled_key_pattern(key: str) -> re.Pattern: + # Compile and cache the regex pattern for a template variable key + return re.compile(r"\{\{\s*" + re.escape(key) + r"\s*\}\}")