Coverage for src/lilbee/server/chat_dispatch/canonical.py: 100%

1"""Protocol-neutral chat request, response, and stream-event types.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from enum import StrEnum 

7from typing import Any, Literal 

8 

9 

10class StopReason(StrEnum): 

11 """Why a canonical chat response ended.""" 

12 

13 END_TURN = "end_turn" 

14 MAX_TOKENS = "max_tokens" 

15 TOOL_USE = "tool_use" 

16 

17 

18@dataclass(frozen=True) 

19class TextBlock: 

20 """Plain-text content block.""" 

21 

22 text: str 

23 type: Literal["text"] = "text" 

24 

25 

26@dataclass(frozen=True) 

27class ToolUseBlock: 

28 """Assistant-emitted tool invocation with parsed JSON arguments.""" 

29 

30 id: str 

31 name: str 

32 input: dict[str, Any] 

33 type: Literal["tool_use"] = "tool_use" 

34 

35 

36@dataclass(frozen=True) 

37class ToolResultBlock: 

38 """Caller-supplied tool result paired to a prior ToolUseBlock by id.""" 

39 

40 tool_use_id: str 

41 content: list[ContentBlock] 

42 is_error: bool = False 

43 type: Literal["tool_result"] = "tool_result" 

44 

45 

46ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock 

47 

48 

49@dataclass(frozen=True) 

50class CanonicalMessage: 

51 """One chat turn; content is always a typed-block list.""" 

52 

53 role: Literal["user", "assistant", "tool"] 

54 content: list[ContentBlock] 

55 

56 @classmethod 

57 def from_string( 

58 cls, 

59 *, 

60 role: Literal["user", "assistant", "tool"], 

61 text: str, 

62 ) -> CanonicalMessage: 

63 """Build a single-text-block message from a raw string.""" 

64 return cls(role=role, content=[TextBlock(text=text)]) 

65 

66 

67@dataclass(frozen=True) 

68class CanonicalTool: 

69 """Tool definition (JSON-Schema input shape).""" 

70 

71 name: str 

72 description: str 

73 input_schema: dict[str, Any] 

74 

75 

76@dataclass(frozen=True) 

77class CanonicalToolChoice: 

78 """Tool-choice mode; ``tool_name`` is required only when ``mode == "tool"``.""" 

79 

80 mode: Literal["auto", "any", "none", "tool"] 

81 tool_name: str | None = None 

82 

83 def __post_init__(self) -> None: 

84 # Without this the None reaches the provider as 

85 # {"function": {"name": None}}, a malformed tool choice rather than a 

86 # rejected request. 

87 if self.mode == "tool" and not self.tool_name: 

88 raise ValueError('CanonicalToolChoice(mode="tool") requires a tool_name') 

89 

90 

91@dataclass(frozen=True) 

92class CanonicalChatRequest: 

93 """Canonical chat request consumed by the dispatch layer.""" 

94 

95 model: str 

96 messages: list[CanonicalMessage] 

97 system: str | None = None 

98 tools: list[CanonicalTool] | None = None 

99 tool_choice: CanonicalToolChoice | None = None 

100 temperature: float | None = None 

101 top_p: float | None = None 

102 top_k: int | None = None 

103 max_tokens: int | None = None 

104 seed: int | None = None 

105 frequency_penalty: float | None = None 

106 presence_penalty: float | None = None 

107 stop: list[str] | None = None 

108 stream: bool = False 

109 # Thinking-template control (chat_template_kwargs.enable_thinking); None 

110 # leaves the model's template default in place. 

111 think: bool | None = None 

112 

113 

114@dataclass(frozen=True) 

115class CanonicalUsage: 

116 """Token-count summary for one chat response.""" 

117 

118 input_tokens: int 

119 output_tokens: int 

120 

121 

122@dataclass(frozen=True) 

123class CanonicalResponse: 

124 """Canonical non-streaming chat response.""" 

125 

126 id: str 

127 model: str 

128 content: list[ContentBlock] 

129 stop_reason: StopReason 

130 usage: CanonicalUsage 

131 

132 

133@dataclass(frozen=True) 

134class MessageStart: 

135 """Stream prelude carrying the message id and model ref.""" 

136 

137 id: str 

138 model: str 

139 

140 

141@dataclass(frozen=True) 

142class ContentBlockStart: 

143 """Opens a fresh content block at ``index`` with an initial shell.""" 

144 

145 index: int 

146 block: ContentBlock 

147 

148 

149@dataclass(frozen=True) 

150class TextDelta: 

151 """One text-token delta within an open text block.""" 

152 

153 text: str 

154 

155 

156@dataclass(frozen=True) 

157class ToolUseDelta: 

158 """Accumulating JSON fragment within an open tool-use block.""" 

159 

160 partial_json: str 

161 

162 

163@dataclass(frozen=True) 

164class ContentBlockDelta: 

165 """Delta payload routed to the content block at ``index``.""" 

166 

167 index: int 

168 delta: TextDelta | ToolUseDelta 

169 

170 

171@dataclass(frozen=True) 

172class ContentBlockStop: 

173 """Closes the content block at ``index``.""" 

174 

175 index: int 

176 

177 

178@dataclass(frozen=True) 

179class MessageDelta: 

180 """Trailing metadata; either field may carry a value, never both required.""" 

181 

182 stop_reason: StopReason | None = None 

183 usage: CanonicalUsage | None = None 

184 

185 

186@dataclass(frozen=True) 

187class MessageStop: 

188 """Stream terminator.""" 

189 

190 

191CanonicalStreamEvent = ( 

192 MessageStart 

193 | ContentBlockStart 

194 | ContentBlockDelta 

195 | ContentBlockStop 

196 | MessageDelta 

197 | MessageStop 

198)