github/MemGPT
7
0
Fork 0
mirror of https://github.com/cpacker/MemGPT.git synced 2025-06-03 04:30:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support system message for chatml wrapper (#2187)

Browse Source
This commit is contained in:
Sarah Wooders 2024-12-07 11:50:18 -08:00 committed by GitHub
parent d2f60b804a
commit ccf15b3b8b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
letta/local_llm/llm_chat_completion_wrappers/chatml.py
View File

@ -3,6 +3,7 @@ from letta.local_llm.json_parser import clean_json
from letta.local_llm.llm_chat_completion_wrappers.wrapper_base import ( from letta.local_llm.llm_chat_completion_wrappers.wrapper_base import (
LLMChatCompletionWrapper, LLMChatCompletionWrapper,
) )
from letta.schemas.enums import MessageRole
from letta.utils import json_dumps, json_loads from letta.utils import json_dumps, json_loads
PREFIX_HINT = """# Reminders: PREFIX_HINT = """# Reminders:
@ -208,7 +209,9 @@ class ChatMLInnerMonologueWrapper(LLMChatCompletionWrapper):
# Last are the user/assistant messages # Last are the user/assistant messages
for message in messages[1:]: for message in messages[1:]:
assert message["role"] in ["user", "assistant", "function", "tool"], message # check that message["role"] is a valid option for MessageRole
# TODO: this shouldn't be necessary if we use pydantic in the future
assert message["role"] in [role.value for role in MessageRole]
if message["role"] == "user": if message["role"] == "user":
# Support for AutoGen naming of agents # Support for AutoGen naming of agents
@ -231,6 +234,15 @@ class ChatMLInnerMonologueWrapper(LLMChatCompletionWrapper):
prompt += f"\n<|im_start|>{role_str}\n{msg_str.strip()}<|im_end|>" prompt += f"\n<|im_start|>{role_str}\n{msg_str.strip()}<|im_end|>"
elif message["role"] == "system":
role_str = "system"
msg_str = self._compile_system_message(
system_message=message["content"], functions=functions, function_documentation=function_documentation
)
prompt += f"\n<|im_start|>{role_str}\n{msg_str.strip()}<|im_end|>"
elif message["role"] in ["tool", "function"]: elif message["role"] in ["tool", "function"]:
if self.allow_function_role: if self.allow_function_role:
role_str = message["role"] role_str = message["role"]