diff --git a/agent/component/message.py b/agent/component/message.py index bf393f541..621240bf6 100644 --- a/agent/component/message.py +++ b/agent/component/message.py @@ -27,7 +27,9 @@ from functools import partial from typing import Any from agent.component.base import ComponentBase, ComponentParamBase -from jinja2 import Template as Jinja2Template +from jinja2.sandbox import SandboxedEnvironment + +_jinja2_sandbox = SandboxedEnvironment() from common.connection_utils import timeout from common.misc_utils import get_uuid @@ -187,7 +189,7 @@ class Message(ComponentBase): return rand_cnt, kwargs = self.get_kwargs(rand_cnt, kwargs) - template = Jinja2Template(rand_cnt) + template = _jinja2_sandbox.from_string(rand_cnt) try: content = template.render(kwargs) except Exception: diff --git a/agent/component/string_transform.py b/agent/component/string_transform.py index 444161f72..d298e5a1b 100644 --- a/agent/component/string_transform.py +++ b/agent/component/string_transform.py @@ -18,7 +18,9 @@ import re from abc import ABC from typing import Any -from jinja2 import Template as Jinja2Template +from jinja2.sandbox import SandboxedEnvironment + +_jinja2_sandbox = SandboxedEnvironment() from agent.component.base import ComponentParamBase from common.connection_utils import timeout from .message import Message @@ -96,7 +98,7 @@ class StringTransform(Message, ABC): script, kwargs = self.get_kwargs(script, kwargs, self._param.delimiters[0]) if self._is_jinjia2(script): - template = Jinja2Template(script) + template = _jinja2_sandbox.from_string(script) try: script = template.render(kwargs) except Exception: