From 05e02f829ce6b473f78ccbf6e83bb74b595b3da4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Apr 2024 15:42:17 -0700 Subject: [PATCH] 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 2abb544095..b60307fefa 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 hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = 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}")