kiln_ai.adapters.chat

26@dataclass
27class BasicChatMessage:
28    role: Literal["system", "assistant", "user"]
29    content: Optional[str]

 88class ChatFormatter(ABC):
 89    def __init__(
 90        self,
 91        system_message: str,
 92        user_input: InputType,
 93        thinking_instructions: str | None = None,
 94    ) -> None:
 95        self.system_message = system_message
 96        self.user_input = user_input
 97        self.thinking_instructions = thinking_instructions
 98        self._messages: List[ChatMessage] = []
 99        self._state = "start"
100        self._intermediate_outputs: Dict[str, str] = {}
101
102    @property
103    def messages(self) -> List[ChatMessage]:
104        return list(self._messages)
105
106    def append_messages(self, messages: Sequence[ChatMessage]) -> None:
107        """Append messages to the internal messages list."""
108        self._messages.extend(messages)
109
110    def message_dicts(self) -> List[dict]:
111        result = []
112        for m in self._messages:
113            msg_dict = {"role": m.role, "content": m.content}
114            if isinstance(m, ToolCallMessage):
115                msg_dict["tool_calls"] = m.tool_calls
116            elif isinstance(m, ToolResponseMessage):
117                msg_dict["tool_call_id"] = m.tool_call_id
118            result.append(msg_dict)
119        return result
120
121    def intermediate_outputs(self) -> Dict[str, str]:
122        """Get the intermediate outputs from the chat formatter."""
123        return self._intermediate_outputs
124
125    def initial_messages(self) -> list[ChatCompletionMessageIncludingLiteLLM]:
126        """Messages to seed the conversation. Empty for fresh runs; prior trace for continuation."""
127        return []
128
129    @abstractmethod
130    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
131        """Advance the conversation and return the next messages if any."""
132        raise NotImplementedError

81class ChatStrategy(str, Enum):
82    """Strategy for how a chat is structured."""
83
84    # Single turn, immediately return the answer
85    single_turn = "final_only"
86    # Two turn, first turn is the thinking, second turn is the answer. Legacy format - used for old fine tunes but not new trains.
87    two_message_cot_legacy = "final_and_intermediate"
88    # Two turn, first turn is the thinking, second turn is the answer. New format - used for new trains.
89    two_message_cot = "two_message_cot"
90    # Single turn, with both the thinking and the answer in the same message, using R1-style thinking format in <think> tags
91    single_turn_r1_thinking = "final_and_intermediate_r1_compatible"

279class MultiturnFormatter(ChatFormatter):
280    """
281    Formatter for continuing a multi-turn conversation with prior trace.
282    Takes prior_trace (existing conversation) and appends the new user message.
283    Produces a single turn: the new user message. Tool calls and multi-turn
284    model responses are handled by _run_model_turn's internal loop.
285
286    When user_input is a dict or list with tool_call_id keys, the input is
287    treated as tool call results (role "tool") rather than a user message.
288    This supports resuming after a return_on_tool_call interrupt.
289    """
290
291    def __init__(
292        self,
293        prior_trace: list[ChatCompletionMessageParam],
294        user_input: InputType,
295    ) -> None:
296        super().__init__(
297            system_message="",
298            user_input=user_input,
299            thinking_instructions=None,
300        )
301        self._prior_trace = prior_trace
302
303    def initial_messages(self) -> list[ChatCompletionMessageIncludingLiteLLM]:
304        """Messages to seed the conversation (prior trace)."""
305        return list(self._prior_trace)
306
307    @property
308    def _is_tool_result(self) -> bool:
309        """Return True if user_input looks like one or more tool call results."""
310        input = self.user_input
311        if isinstance(input, dict):
312            return "tool_call_id" in input
313        if isinstance(input, list):
314            return bool(input) and all(
315                isinstance(item, dict) and "tool_call_id" in item for item in input
316            )
317        return False
318
319    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
320        if self._state == "start":
321            self._state = "awaiting_final"
322            if self._is_tool_result:
323                if isinstance(self.user_input, dict):
324                    raw_items: list[dict] = [self.user_input]
325                else:
326                    raw_items = list(self.user_input)  # type: ignore[arg-type]
327                msgs: list[ChatMessage] = [
328                    ToolResponseMessage(
329                        role="tool",
330                        content=str(item.get("content", "")),
331                        tool_call_id=item["tool_call_id"],
332                        is_error=item.get("is_error"),
333                        error_message=item.get("error_message"),
334                        kiln_task_tool_data=item.get("kiln_task_tool_data"),
335                    )
336                    for item in raw_items
337                ]
338                self._messages.extend(msgs)
339                return ChatTurn(messages=msgs, final_call=True)
340            else:
341                # prior trace is already in the messages list and contains system and so on, we only need
342                # to append the latest new user message
343                user_msg = BasicChatMessage(
344                    "user", format_user_message(self.user_input)
345                )
346                self._messages.append(user_msg)
347                return ChatTurn(messages=[user_msg], final_call=True)
348
349        if self._state == "awaiting_final":
350            if previous_output is None:
351                raise ValueError("previous_output required for final step")
352            self._messages.append(BasicChatMessage("assistant", previous_output))
353            self._state = "done"
354            return None
355
356        return None

32@dataclass
33class ToolCallMessage:
34    """Assistant message with tool calls for chat formatting"""
35
36    role: Literal["assistant"]
37    tool_calls: List[ChatCompletionMessageToolCallParam]
38    content: Optional[str] = None

41@dataclass
42class ToolResponseMessage:
43    """Tool response message for chat formatting"""
44
45    role: Literal["tool"]
46    content: str
47    tool_call_id: str
48    is_error: Optional[bool] = None
49    error_message: Optional[str] = None
50    kiln_task_tool_data: Optional[str] = None

11def build_tool_call_messages(
12    trace: list[ChatCompletionMessageParam] | None,
13) -> list[Union[ToolCallMessage, ToolResponseMessage]]:
14    """
15    Extract tool call and tool response messages from a trace. It's based off the OpenAI schema.
16
17    Args:
18        trace: The trace of the task run in OpenAI format
19
20    Returns:
21        List of ToolCallMessage and ToolResponseMessage objects extracted from the trace
22    """
23    if trace is None:
24        return []
25
26    tool_messages: list[Union[ToolCallMessage, ToolResponseMessage]] = []
27
28    for message in trace:
29        role = message.get("role")
30
31        if role == "assistant" and "tool_calls" in message:
32            tool_calls = message.get("tool_calls")
33            if tool_calls:
34                content = message.get("content")
35                tool_messages.append(
36                    ToolCallMessage(
37                        role="assistant",
38                        tool_calls=tool_calls,
39                        content=extract_text_from_content(content),
40                    )
41                )
42        elif role == "tool":
43            content = message.get("content")
44            tool_call_id = message.get("tool_call_id")
45
46            if tool_call_id is None:
47                raise ValueError("Tool call ID is required for tool response messages")
48            if content is None:
49                raise ValueError("Content is required for tool response messages")
50
51            if not isinstance(content, str):
52                content = str(content)
53
54            tool_messages.append(
55                ToolResponseMessage(
56                    role="tool",
57                    content=content,
58                    tool_call_id=tool_call_id,
59                )
60            )
61
62    return tool_messages

359def get_chat_formatter(
360    strategy: ChatStrategy,
361    system_message: str,
362    user_input: InputType,
363    thinking_instructions: str | None = None,
364) -> ChatFormatter:
365    match strategy:
366        case ChatStrategy.single_turn:
367            return SingleTurnFormatter(system_message, user_input)
368        case ChatStrategy.two_message_cot_legacy:
369            return TwoMessageCotLegacyFormatter(
370                system_message, user_input, thinking_instructions
371            )
372        case ChatStrategy.two_message_cot:
373            return TwoMessageCotFormatter(
374                system_message, user_input, thinking_instructions
375            )
376        case ChatStrategy.single_turn_r1_thinking:
377            return SingleTurnR1ThinkingFormatter(system_message, user_input)
378        case _:
379            raise_exhaustive_enum_error(strategy)