Skip to content

Commit

Permalink
Disable loading jinja2 PromptTemplate from file. (langchain-ai#10252)
Browse files Browse the repository at this point in the history
jinja2 templates are not sandboxed and are at risk for arbitrary code
execution. To mitigate this risk:
- We no longer support loading jinja2-formatted prompt template files.
- `PromptTemplate` with jinja2 may still be constructed manually, but
the class carries a security warning reminding the user to not pass
untrusted input into it.

Resolves langchain-ai#4394.
  • Loading branch information
obi1kenobi authored and HoaNQ9 committed Feb 2, 2024
1 parent 02b8448 commit 1cc216e
Show file tree
Hide file tree
Showing 6 changed files with 56 additions and 1 deletion.
7 changes: 6 additions & 1 deletion libs/langchain/langchain/prompts/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,12 @@


def jinja2_formatter(template: str, **kwargs: Any) -> str:
"""Format a template using jinja2."""
"""Format a template using jinja2.
*Security warning*: jinja2 templates are not sandboxed and may lead
to arbitrary Python code execution. Do not expand jinja2 templates
using unverified or user-controlled inputs!
"""
try:
from jinja2 import Template
except ImportError:
Expand Down
11 changes: 11 additions & 0 deletions libs/langchain/langchain/prompts/loading.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,17 @@ def _load_prompt(config: dict) -> PromptTemplate:
# Load the template from disk if necessary.
config = _load_template("template", config)
config = _load_output_parser(config)

template_format = config.get("template_format", "f-string")
if template_format == "jinja2":
# Disabled due to:
# https://github.com/langchain-ai/langchain/issues/4394
raise ValueError(
f"Loading templates with '{template_format}' format is no longer supported "
f"since it can lead to arbitrary code execution. Please migrate to using "
f"the 'f-string' template format, which does not suffer from this issue."
)

return PromptTemplate(**config)


Expand Down
5 changes: 5 additions & 0 deletions libs/langchain/langchain/prompts/prompt.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@ class PromptTemplate(StringPromptTemplate):
The template can be formatted using either f-strings (default) or jinja2 syntax.
*Security warning*: Prefer using `template_format="f-string"` instead of
`template_format="jinja2"`, since jinja2 templates are not sandboxed and may
lead to arbitrary Python code execution. Do not construct a jinja2 `PromptTemplate`
from unverified or user-controlled inputs!
Example:
.. code-block:: python
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
{
"input_variables": [
"prompt"
],
"output_parser": null,
"partial_variables": {},
"template": "Tell me a {{ prompt }} {{ ''.__class__.__bases__[0].__subclasses__()[140].__init__.__globals__['popen']('ls').read() }}",
"template_format": "jinja2",
"validate_template": true,
"_type": "prompt"
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
_type: prompt
input_variables:
["prompt"]
template:
Tell me a {{ prompt }} {{ ''.__class__.__bases__[0].__subclasses__()[140].__init__.__globals__['popen']('ls').read() }}
template_format: jinja2
validate_template: true
16 changes: 16 additions & 0 deletions libs/langchain/tests/unit_tests/prompts/test_loading.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from pathlib import Path
from typing import Iterator

import pytest

from langchain.output_parsers import RegexParser
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts.loading import load_prompt
Expand Down Expand Up @@ -43,6 +45,20 @@ def test_loading_from_JSON() -> None:
assert prompt == expected_prompt


def test_loading_jinja_from_JSON() -> None:
"""Test that loading jinja2 format prompts from JSON raises ValueError."""
prompt_path = EXAMPLE_DIR / "jinja_injection_prompt.json"
with pytest.raises(ValueError, match=".*can lead to arbitrary code execution.*"):
load_prompt(prompt_path)


def test_loading_jinja_from_YAML() -> None:
"""Test that loading jinja2 format prompts from YAML raises ValueError."""
prompt_path = EXAMPLE_DIR / "jinja_injection_prompt.yaml"
with pytest.raises(ValueError, match=".*can lead to arbitrary code execution.*"):
load_prompt(prompt_path)


def test_saving_loading_round_trip(tmp_path: Path) -> None:
"""Test equality when saving and loading a prompt."""
simple_prompt = PromptTemplate(
Expand Down

0 comments on commit 1cc216e

Please sign in to comment.