fix(backend): Enable Jinja SandboxedEnvironment for TextFormatter (#8891)

We still use plain Jinja objects for text formatting in our block codes.

### Changes 🏗️

Introduced a `TextFormatter` utility class that uses jina
SandboxedEnvironment for safer text formatting.

### Checklist 📋

#### For code changes:
- [ ] I have clearly listed my changes in the PR description
- [ ] I have made a test plan
- [ ] I have tested my changes according to the test plan:
  <!-- Put your test plan here: -->
  - [ ] ...

<details>
  <summary>Example test plan</summary>
  
  - [ ] Create from scratch and execute an agent with at least 3 blocks
- [ ] Import an agent from file upload, and confirm it executes
correctly
  - [ ] Upload agent to marketplace
- [ ] Import an agent from marketplace and confirm it executes correctly
  - [ ] Edit an agent from monitor, and confirm it executes correctly
</details>

#### For configuration changes:
- [ ] `.env.example` is updated or already compatible with my changes
- [ ] `docker-compose.yml` is updated or already compatible with my
changes
- [ ] I have included a list of my configuration changes in the PR
description (under **Changes**)

<details>
  <summary>Examples of configuration changes</summary>

  - Changing ports
  - Adding new services that need to communicate with each other
  - Secrets or environment variable changes
  - New or infrastructure changes such as databases
</details>
This commit is contained in:
Zamil Majdy 2024-12-06 11:21:30 +07:00 committed by GitHub
parent ffc3eff7e2
commit 6dba31e021
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 35 additions and 16 deletions

View File

@ -1,13 +1,11 @@
import re
from typing import Any, List from typing import Any, List
from jinja2 import BaseLoader, Environment
from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema, BlockType from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema, BlockType
from backend.data.model import SchemaField from backend.data.model import SchemaField
from backend.util.mock import MockObject from backend.util.mock import MockObject
from backend.util.text import TextFormatter
jinja = Environment(loader=BaseLoader()) formatter = TextFormatter()
class StoreValueBlock(Block): class StoreValueBlock(Block):
@ -304,9 +302,9 @@ class AgentOutputBlock(Block):
""" """
if input_data.format: if input_data.format:
try: try:
fmt = re.sub(r"(?<!{){[ a-zA-Z0-9_]+}", r"{\g<0>}", input_data.format) yield "output", formatter.format_string(
template = jinja.from_string(fmt) input_data.format, {input_data.name: input_data.value}
yield "output", template.render({input_data.name: input_data.value}) )
except Exception as e: except Exception as e:
yield "output", f"Error: {e}, {input_data.value}" yield "output", f"Error: {e}, {input_data.value}"
else: else:

View File

@ -1,13 +1,11 @@
import re import re
from typing import Any from typing import Any
from jinja2 import BaseLoader, Environment
from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema
from backend.data.model import SchemaField from backend.data.model import SchemaField
from backend.util import json from backend.util import json, text
jinja = Environment(loader=BaseLoader()) formatter = text.TextFormatter()
class MatchTextPatternBlock(Block): class MatchTextPatternBlock(Block):
@ -146,19 +144,20 @@ class FillTextTemplateBlock(Block):
"values": {"list": ["Hello", " World!"]}, "values": {"list": ["Hello", " World!"]},
"format": "{% for item in list %}{{ item }}{% endfor %}", "format": "{% for item in list %}{{ item }}{% endfor %}",
}, },
{
"values": {},
"format": "{% set name = 'Alice' %}Hello, World! {{ name }}",
},
], ],
test_output=[ test_output=[
("output", "Hello, World! Alice"), ("output", "Hello, World! Alice"),
("output", "Hello World!"), ("output", "Hello World!"),
("output", "Hello, World! Alice"),
], ],
) )
def run(self, input_data: Input, **kwargs) -> BlockOutput: def run(self, input_data: Input, **kwargs) -> BlockOutput:
# For python.format compatibility: replace all {...} with {{..}}. yield "output", formatter.format_string(input_data.format, input_data.values)
# But avoid replacing {{...}} to {{{...}}}.
fmt = re.sub(r"(?<!{){[ a-zA-Z0-9_]+}", r"{\g<0>}", input_data.format)
template = jinja.from_string(fmt)
yield "output", template.render(**input_data.values)
class CombineTextsBlock(Block): class CombineTextsBlock(Block):

View File

@ -0,0 +1,22 @@
import re
from jinja2 import BaseLoader
from jinja2.sandbox import SandboxedEnvironment
class TextFormatter:
def __init__(self):
# Create a sandboxed environment
self.env = SandboxedEnvironment(loader=BaseLoader(), autoescape=True)
# Clear any registered filters, tests, and globals to minimize attack surface
self.env.filters.clear()
self.env.tests.clear()
self.env.globals.clear()
def format_string(self, template_str: str, values=None, **kwargs) -> str:
# For python.format compatibility: replace all {...} with {{..}}.
# But avoid replacing {{...}} to {{{...}}}.
template_str = re.sub(r"(?<!{){[ a-zA-Z0-9_]+}", r"{\g<0>}", template_str)
template = self.env.from_string(template_str)
return template.render(values or {}, **kwargs)