Coverage for src/lilbee/providers/fleet/normalize.py: 100%

1"""Reshape an OpenAI tool conversation into strict user/assistant alternation. 

2 

3Some GGUF chat templates (Mistral-Nemo, Cohere command-r) reject a standard 

4OpenAI tool exchange: they require plain user/assistant turns to alternate and 

5raise a Jinja exception on the ``tool`` role or on two same-role turns in a row. 

6:func:`to_alternating` rewrites the conversation into the shape those templates 

7accept. The fleet client learns which models need this by probing the live 

8template once, then applies it proactively before every such request. 

9""" 

10 

11from __future__ import annotations 

12 

13import json 

14from collections.abc import Mapping, Sequence 

15from typing import Any, Final, Literal, TypedDict 

16 

17# Roles this reshaper understands; an unrecognized role falls through to a plain 

18# user turn, so the Literal documents the set without rejecting foreign input. 

19ChatRole = Literal["system", "user", "assistant", "tool"] 

20_SYSTEM_ROLE: Final = "system" 

21_USER_ROLE: Final = "user" 

22_ASSISTANT_ROLE: Final = "assistant" 

23_TOOL_ROLE: Final = "tool" 

24# A tool result carried as a user turn is labelled so the model reads it as a 

25# result rather than a fresh user instruction. 

26_TOOL_RESULT_PREFIX = "Tool result:" 

27# An assistant turn whose only payload was tool_calls becomes a short text note 

28# so the turn is non-empty (the templates reject empty content). 

29_TOOL_CALL_NOTE_PREFIX = "Calling tool" 

30 

31 

32class ChatToolCallFunction(TypedDict, total=False): 

33 """The ``function`` payload of an OpenAI tool call (wire shape).""" 

34 

35 name: str 

36 arguments: str 

37 

38 

39class ChatToolCall(TypedDict, total=False): 

40 """One entry of an assistant message's ``tool_calls`` (wire shape).""" 

41 

42 id: str 

43 type: str 

44 function: ChatToolCallFunction 

45 

46 

47class ChatMessage(TypedDict, total=False): 

48 """One OpenAI chat message in transit. Keys are partial: which are present 

49 depends on the role (a tool result carries ``tool_call_id``, an assistant 

50 tool call carries ``tool_calls``).""" 

51 

52 role: ChatRole 

53 content: Any 

54 tool_calls: list[ChatToolCall] 

55 tool_call_id: str 

56 

57 

58def _content_text(content: Any) -> str: 

59 """Flatten a message ``content`` (string or OpenAI multipart list) to text.""" 

60 if isinstance(content, str): 

61 return content 

62 if isinstance(content, list): 

63 parts = [ 

64 str(part.get("text", "")) 

65 for part in content 

66 if isinstance(part, dict) and part.get("type") == "text" 

67 ] 

68 return "\n".join(p for p in parts if p) 

69 if content is None: 

70 return "" 

71 return str(content) 

72 

73 

74def _one_tool_call_note(call: Any) -> str: 

75 """Render one tool call as ``Calling name(args)``; empty when malformed.""" 

76 fn = call.get("function") if isinstance(call, dict) else None 

77 if not isinstance(fn, dict): 

78 return "" 

79 name = fn.get("name") 

80 if not isinstance(name, str) or not name: 

81 return "" 

82 arguments = fn.get("arguments") 

83 rendered_args = arguments if isinstance(arguments, str) else json.dumps(arguments) 

84 return f"{_TOOL_CALL_NOTE_PREFIX} {name}({rendered_args})" 

85 

86 

87def _tool_calls_note(tool_calls: Any) -> str: 

88 """Render an assistant message's ``tool_calls`` as a short text note.""" 

89 if not isinstance(tool_calls, list): 

90 return "" 

91 return "\n".join(note for call in tool_calls if (note := _one_tool_call_note(call))) 

92 

93 

94def _assistant_text(message: Mapping[str, Any]) -> str: 

95 """Assistant turn text: its content plus a note for any tool_calls it made.""" 

96 pieces = [_content_text(message.get("content")), _tool_calls_note(message.get("tool_calls"))] 

97 return "\n".join(piece for piece in pieces if piece) 

98 

99 

100def _to_turn(message: Mapping[str, Any]) -> tuple[ChatRole, str]: 

101 """Map one OpenAI message to a ``(user|assistant, text)`` pair. 

102 

103 An assistant turn keeps its content and gains a note for any tool calls; a 

104 ``tool`` result becomes a labelled user turn; everything else, including a 

105 stray non-leading ``system`` message, becomes a plain user turn (so it can't 

106 re-open a system block or break alternation mid-conversation). 

107 """ 

108 role = message.get("role") 

109 if role == _ASSISTANT_ROLE: 

110 return _ASSISTANT_ROLE, _assistant_text(message) 

111 if role == _TOOL_ROLE: 

112 result = _content_text(message.get("content")) 

113 return _USER_ROLE, f"{_TOOL_RESULT_PREFIX} {result}".strip() 

114 return _USER_ROLE, _content_text(message.get("content")) 

115 

116 

117def _append_or_merge(turns: list[ChatMessage], turn: ChatMessage) -> None: 

118 """Append a turn, or fold it into the previous turn when the role repeats.""" 

119 if turns and turns[-1]["role"] == turn["role"]: 

120 turns[-1]["content"] = f"{turns[-1]['content']}\n{turn['content']}" 

121 else: 

122 turns.append(turn) 

123 

124 

125def to_alternating(messages: Sequence[Mapping[str, Any]]) -> list[ChatMessage]: 

126 """Rewrite an OpenAI tool conversation into strict user/assistant alternation. 

127 

128 Keeps any leading system messages verbatim, maps each remaining message to a 

129 plain user or assistant turn (``tool`` results become labelled user turns; 

130 assistant tool calls become a text note), then merges consecutive same-role 

131 turns so the result alternates user/assistant after the system block. Every 

132 emitted turn has non-empty content, so a strict-alternation template accepts 

133 it. 

134 """ 

135 leading_system: list[ChatMessage] = [] 

136 index = 0 

137 while index < len(messages) and messages[index].get("role") == _SYSTEM_ROLE: 

138 leading_system.append( 

139 {"role": _SYSTEM_ROLE, "content": _content_text(messages[index].get("content"))} 

140 ) 

141 index += 1 

142 

143 turns: list[ChatMessage] = [] 

144 for message in messages[index:]: 

145 role, text = _to_turn(message) 

146 if text: 

147 _append_or_merge(turns, {"role": role, "content": text}) 

148 return leading_system + turns