fix: use SandboxedEnvironment in templating (#1168)

Author: dylanjcastillo

docs/concepts/templating.md

@@ -5,7 +5,6 @@ description: Learn to dynamically create prompts using Jinja templating and vali
 
 # Prompt Templating
 
-
 With Instructor's Jinja templating, you can:
 
 - Dynamically adapt prompts to any context
@@ -239,3 +238,9 @@ print(address)
 ```
 
 This allows you to preserve your sensitive information while still using it in your prompts.
+
+## Security
+
+We use the `jinja2.sandbox.SandboxedEnvironment` to prevent security issues with the templating engine. This means that you can't use arbitrary python code in your prompts. But this doesn't mean that you should pass untrusted input to the templating engine, as this could still be abused for things like Denial of Service attacks.
+
+You should [always sanitize](https://jinja.palletsprojects.com/en/stable/sandbox/#security-considerations) any input that you pass to the templating engine.

instructor/templating.py

@@ -1,13 +1,14 @@
 # type: ignore[all]
 from __future__ import annotations
 from typing import Any
-from jinja2 import Template
 from textwrap import dedent
 
+from jinja2.sandbox import SandboxedEnvironment
+
 
 def apply_template(text: str, context: dict[str, Any]) -> str:
     """Apply Jinja2 template to the given text."""
-    return dedent(Template(text).render(**context))
+    return dedent(SandboxedEnvironment().from_string(text).render(**context))
 
 
 def process_message(message: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:

tests/test_formatting.py

@@ -1,6 +1,23 @@
+import pytest
+from jinja2.exceptions import SecurityError
+
 from instructor.templating import handle_templating
 
 
+def test_handle_insecure_template():
+    with pytest.raises(SecurityError):
+        kwargs = {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "{{ self.__init__.__globals__.__builtins__.__import__('os').system('ls') }} {{ variable }}",
+                }
+            ]
+        }
+        context = {"variable": "test"}
+        handle_templating(kwargs, context)
+
+
 def test_handle_templating_with_context():
     kwargs = {"messages": [{"role": "user", "content": "Hello {{ name }}!"}]}
     context = {"name": "Alice"}