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 jinja2 import BaseLoader, Environment
from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema, BlockType
from backend.data.model import SchemaField
from backend.util.mock import MockObject
from backend.util.text import TextFormatter
jinja = Environment(loader=BaseLoader())
formatter = TextFormatter()
class StoreValueBlock(Block):
@ -304,9 +302,9 @@ class AgentOutputBlock(Block):
"""
if input_data.format:
try:
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.name: input_data.value})
yield "output", formatter.format_string(
input_data.format, {input_data.name: input_data.value}
)
except Exception as e:
yield "output", f"Error: {e}, {input_data.value}"
else:

View File

@ -1,13 +1,11 @@
import re
from typing import Any
from jinja2 import BaseLoader, Environment
from backend.data.block import Block, BlockCategory, BlockOutput, BlockSchema
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):
@ -146,19 +144,20 @@ class FillTextTemplateBlock(Block):
"values": {"list": ["Hello", " World!"]},
"format": "{% for item in list %}{{ item }}{% endfor %}",
},
{
"values": {},
"format": "{% set name = 'Alice' %}Hello, World! {{ name }}",
},
],
test_output=[
("output", "Hello, World! Alice"),
("output", "Hello World!"),
("output", "Hello, World! Alice"),
],
)
def run(self, input_data: Input, **kwargs) -> BlockOutput:
# For python.format compatibility: replace all {...} with {{..}}.
# 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)
yield "output", formatter.format_string(input_data.format, input_data.values)
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)