phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 19:54:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix - security report / vuln

Browse source
This commit is contained in:
Ishaan Jaff 2024-04-10 15:42:17 -07:00
parent c86213ac98
commit 05e02f829c
1 changed files with 11 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/prompt_templates/factory.py
View file

@ -1,7 +1,7 @@
from enum import Enum from enum import Enum
import requests, traceback import requests, traceback
import json, re, xml.etree.ElementTree as ET 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 from typing import Optional, Any
import imghdr, base64 import imghdr, base64
from typing import List 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): 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 ## get the tokenizer config from huggingface
bos_token = "" bos_token = ""
eos_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"] eos_token = tokenizer_config["eos_token"]
chat_template = tokenizer_config["chat_template"] 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): def raise_exception(message):
raise Exception(f"Error message - {message}") raise Exception(f"Error message - {message}")