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

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

system_message
user_input
thinking_instructions
messages: List[ChatMessage]
46    @property
47    def messages(self) -> List[ChatMessage]:
48        return list(self._messages)
def message_dicts(self) -> List[dict[str, str | None]]:
50    def message_dicts(self) -> List[dict[str, str | None]]:
51        return [{"role": m.role, "content": m.content} for m in self._messages]
def intermediate_outputs(self) -> Dict[str, str]:
53    def intermediate_outputs(self) -> Dict[str, str]:
54        """Get the intermediate outputs from the chat formatter."""
55        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]:
57    @abstractmethod
58    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
59        """Advance the conversation and return the next messages if any."""
60        raise NotImplementedError

Advance the conversation and return the next messages if any.

@dataclass
class ChatMessage:
16@dataclass
17class ChatMessage:
18    role: Literal["system", "assistant", "user"]
19    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:
200def get_chat_formatter(
201    strategy: ChatStrategy,
202    system_message: str,
203    user_input: str | Dict,
204    thinking_instructions: str | None = None,
205) -> ChatFormatter:
206    match strategy:
207        case ChatStrategy.single_turn:
208            return SingleTurnFormatter(system_message, user_input)
209        case ChatStrategy.two_message_cot_legacy:
210            return TwoMessageCotLegacyFormatter(
211                system_message, user_input, thinking_instructions
212            )
213        case ChatStrategy.two_message_cot:
214            return TwoMessageCotFormatter(
215                system_message, user_input, thinking_instructions
216            )
217        case ChatStrategy.single_turn_r1_thinking:
218            return SingleTurnR1ThinkingFormatter(system_message, user_input)
219        case _:
220            raise_exhaustive_enum_error(strategy)