From df7db2b870d2e1201888bb625c446e4473759ffb Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 24 Apr 2024 20:34:54 -0700
Subject: [PATCH] fix(factory.py): support llama3 instuct chat template

allows automatic templating for llama3 instruct requests
---
 litellm/llms/prompt_templates/factory.py | 31 +++++++++++++++++++++---
 litellm/tests/test_prompt_factory.py     | 18 +++++++++++++-
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py
index 305400f4aa..9ae6a2d238 100644
--- a/litellm/llms/prompt_templates/factory.py
+++ b/litellm/llms/prompt_templates/factory.py
@@ -232,7 +232,15 @@ known_tokenizer_config = {
             "eos_token": "</s>",
         },
         "status": "success",
-    }
+    },
+    "meta-llama/Meta-Llama-3-8B-Instruct": {
+        "tokenizer": {
+            "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}",
+            "bos_token": "<|begin_of_text|>",
+            "eos_token": "",
+        },
+        "status": "success",
+    },
 }
 
 
@@ -640,7 +648,7 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str:
         if get_attribute_or_key(tool, "type") != "function":
             continue
 
-        tool_function =  get_attribute_or_key(tool,"function")
+        tool_function = get_attribute_or_key(tool, "function")
         tool_name = get_attribute_or_key(tool_function, "name")
         tool_arguments = get_attribute_or_key(tool_function, "arguments")
         parameters = "".join(
@@ -833,8 +841,14 @@ def convert_to_anthropic_tool_invoke(tool_calls: list) -> list:
         {
             "type": "tool_use",
             "id": get_attribute_or_key(tool, "id"),
-            "name": get_attribute_or_key(get_attribute_or_key(tool, "function"), "name"),
-            "input": json.loads(get_attribute_or_key(get_attribute_or_key(tool, "function"), "arguments")),
+            "name": get_attribute_or_key(
+                get_attribute_or_key(tool, "function"), "name"
+            ),
+            "input": json.loads(
+                get_attribute_or_key(
+                    get_attribute_or_key(tool, "function"), "arguments"
+                )
+            ),
         }
         for tool in tool_calls
         if get_attribute_or_key(tool, "type") == "function"
@@ -1341,6 +1355,14 @@ def prompt_factory(
     try:
         if "meta-llama/llama-2" in model and "chat" in model:
             return llama_2_chat_pt(messages=messages)
+        elif "meta-llama/llama-3" in model and "instruct" in model:
+            return hf_chat_template(
+                model=model,
+                messages=messages,
+                chat_template=known_tokenizer_config[  # type: ignore
+                    "meta-llama/Meta-Llama-3-8B-Instruct"
+                ]["tokenizer"]["chat_template"],
+            )
         elif (
             "tiiuae/falcon" in model
         ):  # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.
@@ -1382,6 +1404,7 @@ def prompt_factory(
             messages=messages
         )  # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2)
 
+
 def get_attribute_or_key(tool_or_function, attribute, default=None):
     if hasattr(tool_or_function, attribute):
         return getattr(tool_or_function, attribute)
diff --git a/litellm/tests/test_prompt_factory.py b/litellm/tests/test_prompt_factory.py
index 8bf66e84d9..2fc04ec528 100644
--- a/litellm/tests/test_prompt_factory.py
+++ b/litellm/tests/test_prompt_factory.py
@@ -14,9 +14,24 @@ from litellm.llms.prompt_templates.factory import (
     anthropic_messages_pt,
     claude_2_1_pt,
     llama_2_chat_pt,
+    prompt_factory,
 )
 
 
+def test_llama_3_prompt():
+    messages = [
+        {"role": "system", "content": "You are a good bot"},
+        {"role": "user", "content": "Hey, how's it going?"},
+    ]
+    received_prompt = prompt_factory(
+        model="meta-llama/Meta-Llama-3-8B-Instruct", messages=messages
+    )
+    print(f"received_prompt: {received_prompt}")
+
+    expected_prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a good bot<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nHey, how's it going?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""
+    assert received_prompt == expected_prompt
+
+
 def test_codellama_prompt_format():
     messages = [
         {"role": "system", "content": "You are a good bot"},
@@ -109,6 +124,7 @@ def test_anthropic_messages_pt():
     messages = []
     with pytest.raises(Exception) as err:
         anthropic_messages_pt(messages)
-    assert("Invalid first message." in str(err.value))
+    assert "Invalid first message." in str(err.value)
+
 
 # codellama_prompt_format()