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 

110 

111@dataclass(frozen=True) 

112class CanonicalUsage: 

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

114 

115 input_tokens: int 

116 output_tokens: int 

117 

118 

119@dataclass(frozen=True) 

120class CanonicalResponse: 

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

122 

123 id: str 

124 model: str 

125 content: list[ContentBlock] 

126 stop_reason: StopReason 

127 usage: CanonicalUsage 

128 

129 

130@dataclass(frozen=True) 

131class MessageStart: 

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

133 

134 id: str 

135 model: str 

136 

137 

138@dataclass(frozen=True) 

139class ContentBlockStart: 

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

141 

142 index: int 

143 block: ContentBlock 

144 

145 

146@dataclass(frozen=True) 

147class TextDelta: 

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

149 

150 text: str 

151 

152 

153@dataclass(frozen=True) 

154class ToolUseDelta: 

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

156 

157 partial_json: str 

158 

159 

160@dataclass(frozen=True) 

161class ContentBlockDelta: 

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

163 

164 index: int 

165 delta: TextDelta | ToolUseDelta 

166 

167 

168@dataclass(frozen=True) 

169class ContentBlockStop: 

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

171 

172 index: int 

173 

174 

175@dataclass(frozen=True) 

176class MessageDelta: 

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

178 

179 stop_reason: StopReason | None = None 

180 usage: CanonicalUsage | None = None 

181 

182 

183@dataclass(frozen=True) 

184class MessageStop: 

185 """Stream terminator.""" 

186 

187 

188CanonicalStreamEvent = ( 

189 MessageStart 

190 | ContentBlockStart 

191 | ContentBlockDelta 

192 | ContentBlockStop 

193 | MessageDelta 

194 | MessageStop 

195)