kiln_ai.adapters.chat

1from .chat_formatter import (
2    ChatFormatter,
3    ChatMessage,
4    ChatStrategy,
5    get_chat_formatter,
6)
7
8__all__ = ["ChatFormatter", "ChatMessage", "ChatStrategy", "get_chat_formatter"]
class ChatFormatter(abc.ABC):
31class ChatFormatter(ABC):
32    def __init__(
33        self,
34        system_message: str,
35        user_input: str | Dict,
36        thinking_instructions: str | None = None,
37    ) -> None:
38        self.system_message = system_message
39        self.user_input = user_input
40        self.thinking_instructions = thinking_instructions
41        self._messages: List[ChatMessage] = []
42        self._state = "start"
43        self._intermediate_outputs: Dict[str, str] = {}
44
45    @property
46    def messages(self) -> List[ChatMessage]:
47        return list(self._messages)
48
49    def message_dicts(self) -> List[dict[str, str | None]]:
50        return [{"role": m.role, "content": m.content} for m in self._messages]
51
52    def intermediate_outputs(self) -> Dict[str, str]:
53        """Get the intermediate outputs from the chat formatter."""
54        return self._intermediate_outputs
55
56    @abstractmethod
57    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
58        """Advance the conversation and return the next messages if any."""
59        raise NotImplementedError

Helper class that provides a standard way to create an ABC using inheritance.

system_message
user_input
thinking_instructions
messages: List[ChatMessage]
45    @property
46    def messages(self) -> List[ChatMessage]:
47        return list(self._messages)
def message_dicts(self) -> List[dict[str, str | None]]:
49    def message_dicts(self) -> List[dict[str, str | None]]:
50        return [{"role": m.role, "content": m.content} for m in self._messages]
def intermediate_outputs(self) -> Dict[str, str]:
52    def intermediate_outputs(self) -> Dict[str, str]:
53        """Get the intermediate outputs from the chat formatter."""
54        return self._intermediate_outputs

Get the intermediate outputs from the chat formatter.

@abstractmethod
def next_turn( self, previous_output: str | None = None) -> Optional[kiln_ai.adapters.chat.chat_formatter.ChatTurn]:
56    @abstractmethod
57    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
58        """Advance the conversation and return the next messages if any."""
59        raise NotImplementedError

Advance the conversation and return the next messages if any.

@dataclass
class ChatMessage:
15@dataclass
16class ChatMessage:
17    role: Literal["system", "assistant", "user"]
18    content: Optional[str]
ChatMessage(role: Literal['system', 'assistant', 'user'], content: Optional[str])
role: Literal['system', 'assistant', 'user']
content: Optional[str]
class ChatStrategy(builtins.str, enum.Enum):
61class ChatStrategy(str, Enum):
62    """Strategy for how a chat is structured."""
63
64    # Single turn, immediately return the answer
65    single_turn = "final_only"
66    # Two turn, first turn is the thinking, second turn is the answer. Legacy format - used for old fine tunes but not new trains.
67    two_message_cot_legacy = "final_and_intermediate"
68    # Two turn, first turn is the thinking, second turn is the answer. New format - used for new trains.
69    two_message_cot = "two_message_cot"
70    # Single turn, with both the thinking and the answer in the same message, using R1-style thinking format in <think> tags
71    single_turn_r1_thinking = "final_and_intermediate_r1_compatible"

Strategy for how a chat is structured.

single_turn = <ChatStrategy.single_turn: 'final_only'>
two_message_cot_legacy = <ChatStrategy.two_message_cot_legacy: 'final_and_intermediate'>
two_message_cot = <ChatStrategy.two_message_cot: 'two_message_cot'>
single_turn_r1_thinking = <ChatStrategy.single_turn_r1_thinking: 'final_and_intermediate_r1_compatible'>
def get_chat_formatter( strategy: ChatStrategy, system_message: str, user_input: Union[str, Dict], thinking_instructions: str | None = None) -> ChatFormatter:
195def get_chat_formatter(
196    strategy: ChatStrategy,
197    system_message: str,
198    user_input: str | Dict,
199    thinking_instructions: str | None = None,
200) -> ChatFormatter:
201    match strategy:
202        case ChatStrategy.single_turn:
203            return SingleTurnFormatter(system_message, user_input)
204        case ChatStrategy.two_message_cot_legacy:
205            return TwoMessageCotLegacyFormatter(
206                system_message, user_input, thinking_instructions
207            )
208        case ChatStrategy.two_message_cot:
209            return TwoMessageCotFormatter(
210                system_message, user_input, thinking_instructions
211            )
212        case ChatStrategy.single_turn_r1_thinking:
213            return SingleTurnR1ThinkingFormatter(system_message, user_input)
214        case _:
215            raise_exhaustive_enum_error(strategy)