From 05e02f829ce6b473f78ccbf6e83bb74b595b3da4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Apr 2024 15:42:17 -0700 Subject: [PATCH 1/4] fix - security report / vuln --- litellm/llms/prompt_templates/factory.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index 2abb544095d7..b60307fefa4e 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -1,7 +1,7 @@ from enum import Enum import requests, traceback import json, re, xml.etree.ElementTree as ET -from jinja2 import Template, exceptions, Environment, meta +from jinja2 import Template, exceptions, Environment, meta, BaseLoader from typing import Optional, Any import imghdr, base64 from typing import List @@ -219,6 +219,9 @@ def phind_codellama_pt(messages): def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None): + # Define Jinja2 environment with autoescaping disabled + env = Environment(loader=BaseLoader(), autoescape=False) + ## get the tokenizer config from huggingface bos_token = "" eos_token = "" @@ -249,6 +252,13 @@ def _get_tokenizer_config(hf_model_name): eos_token = tokenizer_config["eos_token"] chat_template = tokenizer_config["chat_template"] + # Render the chat_template safely + rendered_template = env.from_string(chat_template).render(messages=messages) + + # Now you can use the rendered_template in your application + # For example, return it or process further + return rendered_template + def raise_exception(message): raise Exception(f"Error message - {message}") From f9cabf09ff49544d3e677f34e6da26de06af1596 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Apr 2024 20:30:13 -0700 Subject: [PATCH 2/4] fix jinja2 use ImmutableSandboxedEnvironment --- litellm/llms/prompt_templates/factory.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index b60307fefa4e..26b86a085e33 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -1,7 +1,8 @@ from enum import Enum import requests, traceback import json, re, xml.etree.ElementTree as ET -from jinja2 import Template, exceptions, Environment, meta, BaseLoader +from jinja2 import Template, exceptions, meta, BaseLoader +from jinja2.sandbox import ImmutableSandboxedEnvironment from typing import Optional, Any import imghdr, base64 from typing import List @@ -220,7 +221,13 @@ def phind_codellama_pt(messages): def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None): # Define Jinja2 environment with autoescaping disabled - env = Environment(loader=BaseLoader(), autoescape=False) + env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True) + + def raise_exception(message): + raise Exception(f"Error message - {message}") + + # Create a template object from the template text + env.globals["raise_exception"] = raise_exception ## get the tokenizer config from huggingface bos_token = "" @@ -259,12 +266,6 @@ def _get_tokenizer_config(hf_model_name): # For example, return it or process further return rendered_template - def raise_exception(message): - raise Exception(f"Error message - {message}") - - # Create a template object from the template text - env = Environment() - env.globals["raise_exception"] = raise_exception try: template = env.from_string(chat_template) except Exception as e: From 6a5d5bcd47539242f5e689355dc4d57ad3973224 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Apr 2024 21:26:50 -0700 Subject: [PATCH 3/4] fix use ImmutableSandboxedEnvironment --- litellm/llms/prompt_templates/factory.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index 26b86a085e33..f137c90f65d0 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -221,7 +221,7 @@ def phind_codellama_pt(messages): def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None): # Define Jinja2 environment with autoescaping disabled - env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True) + env = ImmutableSandboxedEnvironment() def raise_exception(message): raise Exception(f"Error message - {message}") @@ -259,13 +259,6 @@ def _get_tokenizer_config(hf_model_name): eos_token = tokenizer_config["eos_token"] chat_template = tokenizer_config["chat_template"] - # Render the chat_template safely - rendered_template = env.from_string(chat_template).render(messages=messages) - - # Now you can use the rendered_template in your application - # For example, return it or process further - return rendered_template - try: template = env.from_string(chat_template) except Exception as e: From b3f62b7ce35c05de7228c6ecf73189de1469a8db Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Apr 2024 21:28:04 -0700 Subject: [PATCH 4/4] fix cleanup --- litellm/llms/prompt_templates/factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index f137c90f65d0..d7f2272c4f58 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -220,7 +220,7 @@ def phind_codellama_pt(messages): def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None): - # Define Jinja2 environment with autoescaping disabled + # Define Jinja2 environment env = ImmutableSandboxedEnvironment() def raise_exception(message):