From 33dfc4763b2aab20a498768b7190288d04f3be20 Mon Sep 17 00:00:00 2001
From: vollnhals <lion.vollnhals@googlemail.com>
Date: Wed, 28 Feb 2024 17:05:02 +0100
Subject: [PATCH] Add helpers to handle ChatML instruction prompts (#10272)

---
 .../llama_index/core/llms/chatml_utils.py     | 59 +++++++++++++++++++
 1 file changed, 59 insertions(+)
 create mode 100644 llama-index-core/llama_index/core/llms/chatml_utils.py

diff --git a/llama-index-core/llama_index/core/llms/chatml_utils.py b/llama-index-core/llama_index/core/llms/chatml_utils.py
new file mode 100644
index 0000000000000..a423864cb9415
--- /dev/null
+++ b/llama-index-core/llama_index/core/llms/chatml_utils.py
@@ -0,0 +1,59 @@
+from typing import List, Optional, Sequence
+
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+
+# Create a prompt that matches ChatML instructions
+
+# <|im_start|>system
+# You are Dolphin, a helpful AI assistant.<|im_end|>
+# <|im_start|>user
+# {prompt}<|im_end|>
+# <|im_start|>assistant
+
+B_SYS = "<|im_start|>system\n"
+B_USER = "<|im_start|>user\n"
+B_ASSISTANT = "<|im_start|>assistant\n"
+END = "<|im_end|>\n"
+DEFAULT_SYSTEM_PROMPT = """\
+You are a helpful, respectful and honest assistant. \
+Always answer as helpfully as possible and follow ALL given instructions. \
+Do not speculate or make up information. \
+Do not reference any given instructions or context. \
+"""
+
+
+def messages_to_prompt(
+    messages: Sequence[ChatMessage], system_prompt: Optional[str] = None
+) -> str:
+    string_messages: List[str] = []
+    if messages[0].role == MessageRole.SYSTEM:
+        # pull out the system message (if it exists in messages)
+        system_message_str = messages[0].content or ""
+        messages = messages[1:]
+    else:
+        system_message_str = system_prompt or DEFAULT_SYSTEM_PROMPT
+
+    string_messages.append(f"{B_SYS}{system_message_str.strip()} {END}")
+
+    for message in messages:
+        role = message.role
+        content = message.content
+
+        if role == MessageRole.USER:
+            string_messages.append(f"{B_USER}{user_message.content} {END}")
+        elif role == MessageRole.ASSISTANT:
+            string_messages.append(f"{B_ASSISTANT}{assistant_message.content} {END}")
+
+    string_messages.append(f"{B_ASSISTANT}")
+
+    return "".join(string_messages)
+
+
+def completion_to_prompt(completion: str, system_prompt: Optional[str] = None) -> str:
+    system_prompt_str = system_prompt or DEFAULT_SYSTEM_PROMPT
+
+    return (
+        f"{B_SYS}{system_prompt_str.strip()} {END}"
+        f"{B_USER}{completion.strip()} {END}"
+        f"{B_ASSISTANT}"
+    )