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

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

315class MultiturnFormatter(ChatFormatter):
316    """
317    Formatter for continuing a multi-turn conversation with prior trace.
318    Takes prior_trace (existing conversation) and appends the new user message.
319    Produces a single turn: the new user message. Tool calls and multi-turn
320    model responses are handled by _run_model_turn's internal loop.
321
322    When user_input is a dict or list with tool_call_id keys, the input is
323    treated as tool call results (role "tool") rather than a user message.
324    This supports resuming after a return_on_tool_call interrupt.
325    """
326
327    def __init__(
328        self,
329        prior_trace: list[ChatCompletionMessageParam],
330        user_input: InputType,
331    ) -> None:
332        super().__init__(
333            system_message="",
334            user_input=user_input,
335            thinking_instructions=None,
336        )
337        self._prior_trace = prior_trace
338
339    def initial_messages(self) -> list[ChatCompletionMessageIncludingLiteLLM]:
340        """Messages to seed the conversation (prior trace)."""
341        return list(self._prior_trace)
342
343    @property
344    def _is_tool_result(self) -> bool:
345        """Return True if user_input looks like one or more tool call results."""
346        input = self.user_input
347        if isinstance(input, dict):
348            return "tool_call_id" in input
349        if isinstance(input, list):
350            return bool(input) and all(
351                isinstance(item, dict) and "tool_call_id" in item for item in input
352            )
353        return False
354
355    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
356        if self._state == "start":
357            self._state = "awaiting_final"
358            if self._is_tool_result:
359                if isinstance(self.user_input, dict):
360                    raw_items: list[dict] = [self.user_input]
361                else:
362                    raw_items = list(self.user_input)  # type: ignore[arg-type]
363                msgs: list[ChatMessage] = [
364                    ToolResponseMessage(
365                        role="tool",
366                        content=str(item.get("content", "")),
367                        tool_call_id=item["tool_call_id"],
368                        is_error=item.get("is_error"),
369                        error_message=item.get("error_message"),
370                        kiln_task_tool_data=item.get("kiln_task_tool_data"),
371                    )
372                    for item in raw_items
373                ]
374                self._messages.extend(msgs)
375                return ChatTurn(messages=msgs, final_call=True)
376            else:
377                # prior trace is already in the messages list and contains system and so on, we only need
378                # to append the latest new user message
379                user_msg = BasicChatMessage(
380                    "user", format_user_message(self.user_input)
381                )
382                self._messages.append(user_msg)
383                return ChatTurn(messages=[user_msg], final_call=True)
384
385        if self._state == "awaiting_final":
386            if previous_output is None:
387                raise ValueError("previous_output required for final step")
388            self._messages.append(BasicChatMessage("assistant", previous_output))
389            self._state = "done"
390            return None
391
392        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

395def get_chat_formatter(
396    strategy: ChatStrategy,
397    system_message: str,
398    user_input: InputType,
399    thinking_instructions: str | None = None,
400    forward_thinking_instructions: bool = False,
401) -> ChatFormatter:
402    match strategy:
403        case ChatStrategy.single_turn:
404            return SingleTurnFormatter(system_message, user_input)
405        case ChatStrategy.two_message_cot_legacy:
406            return TwoMessageCotLegacyFormatter(
407                system_message, user_input, thinking_instructions
408            )
409        case ChatStrategy.two_message_cot:
410            return TwoMessageCotFormatter(
411                system_message, user_input, thinking_instructions
412            )
413        case ChatStrategy.single_turn_r1_thinking:
414            return SingleTurnR1ThinkingFormatter(
415                system_message,
416                user_input,
417                thinking_instructions,
418                forward_thinking_instructions=forward_thinking_instructions,
419            )
420        case _:
421            raise_exhaustive_enum_error(strategy)