codekingpro/portable-devtools
114k
1from typing import List2 3from langchain_core.messages import (4 AIMessage,5 BaseMessage,6 ChatMessage,7 HumanMessage,8 SystemMessage,9)10 11 12def _convert_one_message_to_text_llama(message: BaseMessage) -> str:13 if isinstance(message, ChatMessage):14 message_text = f"\n\n{message.role.capitalize()}: {message.content}"15 elif isinstance(message, HumanMessage):16 message_text = f"[INST] {message.content} [/INST]"17 elif isinstance(message, AIMessage):18 message_text = f"{message.content}"19 elif isinstance(message, SystemMessage):20 message_text = f"<<SYS>> {message.content} <</SYS>>"21 else:22 raise ValueError(f"Got unknown type {message}")23 return message_text24 25 26def convert_messages_to_prompt_llama(messages: List[BaseMessage]) -> str:27 """Convert a list of messages to a prompt for llama."""28 29 return "\n".join(30 [_convert_one_message_to_text_llama(message) for message in messages]31 )32 