brief-extractor/backend/venv/lib/python3.10/site-packages/banks/env.py
2026-03-06 18:42:46 +00:00

43 lines
1.2 KiB
Python
Executable file

# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
#
# SPDX-License-Identifier: MIT
from jinja2 import Environment, select_autoescape
from .config import config
from .filters import audio, cache_control, image, lemmatize, tool, xml
def _add_extensions(_env):
"""
We lazily add extensions so that we can use the env in the extensions themselves if needed.
For example, we use banks to manage the system prompt in `GenerateExtension`
"""
from .extensions.chat import ChatExtension # pylint: disable=import-outside-toplevel
from .extensions.completion import CompletionExtension # pylint: disable=import-outside-toplevel
_env.add_extension(ChatExtension)
_env.add_extension(CompletionExtension)
# Init the Jinja env
env = Environment(
autoescape=select_autoescape(
enabled_extensions=("html", "xml"),
default_for_string=False,
),
trim_blocks=True,
lstrip_blocks=True,
enable_async=bool(config.ASYNC_ENABLED),
)
# Setup custom filters and defaults
env.filters["cache_control"] = cache_control
env.filters["image"] = image
env.filters["lemmatize"] = lemmatize
env.filters["tool"] = tool
env.filters["audio"] = audio
env.filters["to_xml"] = xml
_add_extensions(env)