add filters to prompt function

outlines/prompts.py

@@ -6,7 +6,7 @@
 import textwrap
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional, Type, cast
+from typing import Any, Callable, Dict, List, Optional, Type, Union, cast
 
 import jinja2
 import pydantic
@@ -40,7 +40,9 @@ def __call__(self, *args, **kwargs) -> str:
             return self.template.render(**kwargs)
 
     @classmethod
-    def from_str(cls, content: str):
+    def from_str(
+        cls, content: str, filters: Union[List[Callable], Dict[str, Callable]] = []
+    ):
         """
         Create an instance of the class from a string.
 
@@ -53,10 +55,12 @@ def from_str(cls, content: str):
         -------
         An instance of the class with the provided content as a template.
         """
-        return cls(cls._template_from_str(content), None)
+        return cls(cls._template_from_str(content, filters), None)
 
     @classmethod
-    def from_file(cls, path: Path):
+    def from_file(
+        cls, path: Path, filters: Union[List[Callable], Dict[str, Callable]] = []
+    ):
         """
         Create a Prompt instance from a file containing a Jinja template.
 
@@ -75,10 +79,12 @@ def from_file(cls, path: Path):
         """
         # We don't use a `Signature` here because it seems not feasible to infer one from a Jinja2 environment that is
         # split across multiple files (since e.g. we support features like Jinja2 includes and template inheritance)
-        return cls(cls._template_from_file(path), None)
+        return cls(cls._template_from_file(path, filters), None)
 
     @classmethod
-    def _template_from_str(_, content: str) -> jinja2.Template:
+    def _template_from_str(
+        _, content: str, filters: Union[List[Callable], Dict[str, Callable]] = []
+    ) -> jinja2.Template:
         # Dedent, and remove extra linebreak
         cleaned_template = inspect.cleandoc(content)
 
@@ -106,10 +112,14 @@ def _template_from_str(_, content: str) -> jinja2.Template:
         env.filters["schema"] = get_schema
         env.filters["args"] = get_fn_args
 
+        _add_filters(env, filters)
+
         return env.from_string(cleaned_template)
 
     @classmethod
-    def _template_from_file(_, path: Path) -> jinja2.Template:
+    def _template_from_file(
+        _, path: Path, filters: Union[List[Callable], Dict[str, Callable]] = []
+    ) -> jinja2.Template:
         file_directory = os.path.dirname(os.path.abspath(path))
         env = jinja2.Environment(
             loader=jinja2.FileSystemLoader(file_directory),
@@ -118,10 +128,16 @@ def _template_from_file(_, path: Path) -> jinja2.Template:
             keep_trailing_newline=True,
             undefined=jinja2.StrictUndefined,
         )
+
+        _add_filters(env, filters)
+
         return env.get_template(os.path.basename(path))
 
 
-def prompt(fn: Callable) -> Prompt:
+def prompt(
+    fn: Optional[Callable] = None,
+    filters: Union[List[Callable], Dict[str, Callable]] = [],
+) -> Callable:
     """Decorate a function that contains a prompt template.
 
     This allows to define prompts in the docstring of a function and simplify their
@@ -152,11 +168,28 @@ def prompt(fn: Callable) -> Prompt:
     ...
     >>> hal = ft.partial(solve_task, "HAL", "Travel to Jupiter")
 
+    Additional Jinja2 filters can be provided as keyword arguments to the decorator.
+
+    >>> def reverse(s: str) -> str:
+    ...     return s[::-1]
+    ...
+    >>> @outlines.prompt(filters=[reverse])
+    ... def reverse_prompt(text):
+    ...     '''{{ text | reverse }}'''
+    ...
+    >>> prompt = reverse_prompt("Hello")
+    >>> print(prompt)
+    ... "olleH"
+
     Returns
     -------
     A `Prompt` callable class which will render the template when called.
 
     """
+    if fn is None:
+        return lambda fn: prompt(
+            fn, cast(Union[List[Callable], Dict[str, Callable]], filters)
+        )
 
     signature = inspect.signature(fn)
 
@@ -166,11 +199,22 @@ def prompt(fn: Callable) -> Prompt:
     if docstring is None:
         raise TypeError("Could not find a template in the function's docstring.")
 
-    template = Prompt._template_from_str(cast(str, docstring))
+    template = Prompt._template_from_str(cast(str, docstring), filters)
 
     return Prompt(template, signature)
 
 
+def _add_filters(
+    env: jinja2.Environment, filters: Union[List[Callable], Dict[str, Callable]]
+):
+    if isinstance(filters, list):
+        for filter_fn in filters:
+            env.filters[filter_fn.__name__] = filter_fn
+    elif isinstance(filters, dict):
+        for name, filter_fn in filters.items():
+            env.filters[name] = filter_fn
+
+
 def get_fn_name(fn: Callable):
     """Returns the name of a callable."""
     if not callable(fn):

tests/test_prompts.py

@@ -321,6 +321,40 @@ def args_prompt(fn):
     )
 
 
+def test_prompt_with_additional_filters_as_dict():
+    def reverse(s: str) -> str:
+        return s[::-1]
+
+    @outlines.prompt(filters=dict(reverse=reverse))
+    def test_tpl(variable):
+        """{{ variable | reverse }} test"""
+
+    assert list(test_tpl.signature.parameters) == ["variable"]
+
+    p = test_tpl("test")
+    assert p == "tset test"
+
+    p = test_tpl(variable="example")
+    assert p == "elpmaxe test"
+
+
+def test_prompt_with_additional_filters_as_list():
+    def reverse(s: str) -> str:
+        return s[::-1]
+
+    @outlines.prompt(filters=[reverse])
+    def test_tpl(variable):
+        """{{ variable | reverse }} test"""
+
+    assert list(test_tpl.signature.parameters) == ["variable"]
+
+    p = test_tpl("test")
+    assert p == "tset test"
+
+    p = test_tpl(variable="example")
+    assert p == "elpmaxe test"
+
+
 @pytest.fixture
 def temp_prompt_file():
     test_dir = tempfile.mkdtemp()