add tests

Author: sainivedh

clarifai_datautils/text/prompt_factory.py

@@ -1,9 +1,9 @@
-# This was taken from litellm/litellm_core_utils/prompt_templates/factory.py
+# This was taken from litellm
 
 from enum import Enum
 from typing import Any, Optional
 import json
-from jinja2 import Environment
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 
 import requests
 
@@ -415,7 +415,7 @@ def raise_exception(message):
         raise Exception(f"Error message - {message}")
 
     # Create a template object from the template text
-    env = Environment()
+    env = ImmutableSandboxedEnvironment()
     env.globals["raise_exception"] = raise_exception
     try:
         template = env.from_string(chat_template)
@@ -555,6 +555,11 @@ def prompt_factory(
     try:
         if "meta-llama/llama-2" in model and "chat" in model:
             return llama_2_chat_pt(messages=messages)
+        elif "llama3" in model and "instruct" in model:
+            return hf_chat_template(
+                model="meta-llama/Meta-Llama-3-8B-Instruct",
+                messages=messages,
+            )
         elif (
             "tiiuae/falcon" in model
         ):  # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.

tests/test_prompt_factory.py

@@ -0,0 +1,47 @@
+# This was taken from litellm
+
+from clarifai_datautils.text.prompt_factory import claude_2_1_pt, llama_2_chat_pt
+
+
+def test_codellama_prompt_format():
+    messages = [
+        {"role": "system", "content": "You are a good bot"},
+        {"role": "user", "content": "Hey, how's it going?"},
+    ]
+    expected_prompt = "<s>[INST] <<SYS>>\nYou are a good bot\n<</SYS>>\n [/INST]\n[INST] Hey, how's it going? [/INST]\n"
+    assert llama_2_chat_pt(messages) == expected_prompt
+
+def test_claude_2_1_pt_formatting():
+    # Test case: User only, should add Assistant
+    messages = [{"role": "user", "content": "Hello"}]
+    expected_prompt = "\n\nHuman: Hello\n\nAssistant: "
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, User, and Assistant "pre-fill" sequence,
+    #            Should return pre-fill
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": 'Please return "Hello World" as a JSON object.'},
+        {"role": "assistant", "content": "{"},
+    ]
+    expected_prompt = 'You are a helpful assistant.\n\nHuman: Please return "Hello World" as a JSON object.\n\nAssistant: {'
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, Assistant sequence, should insert blank Human message
+    #            before Assistant pre-fill
+    messages = [
+        {"role": "system", "content": "You are a storyteller."},
+        {"role": "assistant", "content": "Once upon a time, there "},
+    ]
+    expected_prompt = (
+        "You are a storyteller.\n\nHuman: \n\nAssistant: Once upon a time, there "
+    )
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, User sequence
+    messages = [
+        {"role": "system", "content": "System reboot"},
+        {"role": "user", "content": "Is everything okay?"},
+    ]
+    expected_prompt = "System reboot\n\nHuman: Is everything okay?\n\nAssistant: "
+    assert claude_2_1_pt(messages) == expected_prompt