Coverage for src / harnessutils / models / parts.py: 90%

1"""Part types for message decomposition. 

2 

3Parts are the granular units that make up messages. This enables 

4selective compaction where tool outputs can be cleared while 

5preserving text and metadata. 

6""" 

7 

8from dataclasses import dataclass, field 

9from typing import Any, Literal 

10 

11 

12@dataclass 

13class TimeInfo: 

14 """Timing information for parts.""" 

15 

16 start: int # Unix timestamp in milliseconds 

17 end: int | None = None 

18 compacted: int | None = None # When output was compacted 

19 

20 

21@dataclass 

22class ToolState: 

23 """State information for tool execution.""" 

24 

25 status: Literal["pending", "running", "completed", "error"] 

26 input: dict[str, Any] = field(default_factory=dict) 

27 output: str = "" 

28 summary_text: str | None = None # One-sentence summary generated before output was cleared 

29 title: str = "" 

30 metadata: dict[str, Any] = field(default_factory=dict) 

31 error: str | None = None 

32 time: TimeInfo | None = None 

33 attachments: list[dict[str, Any]] = field(default_factory=list) 

34 

35 

36@dataclass 

37class Part: 

38 """Base class for message parts.""" 

39 

40 type: str = field(init=False) 

41 time: TimeInfo | None = None 

42 metadata: dict[str, Any] = field(default_factory=dict) 

43 

44 def to_dict(self) -> dict[str, Any]: 

45 """Serialize part to dict for storage.""" 

46 d: dict[str, Any] = {"type": self.type, "metadata": self.metadata} 

47 if self.time: 

48 d["time"] = { 

49 "start": self.time.start, 

50 "end": self.time.end, 

51 "compacted": self.time.compacted, 

52 } 

53 return d 

54 

55 

56@dataclass 

57class TextPart(Part): 

58 """Text content part.""" 

59 

60 text: str = "" 

61 ignored: bool = False # If true, skip when converting to model format 

62 

63 def __post_init__(self) -> None: 

64 self.type = "text" 

65 

66 def to_dict(self) -> dict[str, Any]: 

67 d = super().to_dict() 

68 d["text"] = self.text 

69 d["ignored"] = self.ignored 

70 return d 

71 

72 

73@dataclass 

74class ReasoningPart(Part): 

75 """Extended thinking/reasoning content part.""" 

76 

77 text: str = "" 

78 

79 def __post_init__(self) -> None: 

80 self.type = "reasoning" 

81 

82 def to_dict(self) -> dict[str, Any]: 

83 d = super().to_dict() 

84 d["text"] = self.text 

85 return d 

86 

87 

88@dataclass 

89class ToolPart(Part): 

90 """Tool execution part.""" 

91 

92 tool: str = "" 

93 call_id: str = "" 

94 state: ToolState = field(default_factory=lambda: ToolState(status="pending")) 

95 

96 def __post_init__(self) -> None: 

97 self.type = "tool" 

98 

99 def to_dict(self) -> dict[str, Any]: 

100 d = super().to_dict() 

101 d["tool"] = self.tool 

102 d["call_id"] = self.call_id 

103 state = self.state 

104 state_dict: dict[str, Any] = { 

105 "status": state.status, 

106 "input": state.input, 

107 "output": state.output, 

108 "summary_text": state.summary_text, 

109 "title": state.title, 

110 "metadata": state.metadata, 

111 "attachments": state.attachments, 

112 } 

113 if state.error is not None: 

114 state_dict["error"] = state.error 

115 if state.time: 

116 state_dict["time"] = { 

117 "start": state.time.start, 

118 "end": state.time.end, 

119 "compacted": state.time.compacted, 

120 } 

121 d["state"] = state_dict 

122 return d 

123 

124 

125@dataclass 

126class StepStartPart(Part): 

127 """Step boundary marker - start of LLM turn.""" 

128 

129 snapshot: str = "" # State snapshot ID 

130 

131 def __post_init__(self) -> None: 

132 self.type = "step-start" 

133 

134 def to_dict(self) -> dict[str, Any]: 

135 d = super().to_dict() 

136 d["snapshot"] = self.snapshot 

137 return d 

138 

139 

140@dataclass 

141class StepFinishPart(Part): 

142 """Step boundary marker - end of LLM turn.""" 

143 

144 reason: Literal["stop", "tool-calls", "length"] = "stop" 

145 snapshot: str = "" # State snapshot ID 

146 tokens: dict[str, Any] = field(default_factory=dict) 

147 cost: float = 0.0 

148 

149 def __post_init__(self) -> None: 

150 self.type = "step-finish" 

151 

152 def to_dict(self) -> dict[str, Any]: 

153 d = super().to_dict() 

154 d["reason"] = self.reason 

155 d["snapshot"] = self.snapshot 

156 d["tokens"] = self.tokens 

157 d["cost"] = self.cost 

158 return d 

159 

160 

161@dataclass 

162class CompactionPart(Part): 

163 """Marker for compaction request.""" 

164 

165 auto: bool = False # Was this auto-triggered? 

166 

167 def __post_init__(self) -> None: 

168 self.type = "compaction" 

169 

170 def to_dict(self) -> dict[str, Any]: 

171 d = super().to_dict() 

172 d["auto"] = self.auto 

173 return d 

174 

175 

176@dataclass 

177class PatchPart(Part): 

178 """Code change marker.""" 

179 

180 hash: str = "" # Diff hash 

181 files: list[str] = field(default_factory=list) 

182 

183 def __post_init__(self) -> None: 

184 self.type = "patch" 

185 

186 def to_dict(self) -> dict[str, Any]: 

187 d = super().to_dict() 

188 d["hash"] = self.hash 

189 d["files"] = self.files 

190 return d 

191 

192 

193@dataclass 

194class SubtaskPart(Part): 

195 """Subtask invocation marker.""" 

196 

197 prompt: str = "" 

198 description: str = "" 

199 agent: str = "" 

200 model: dict[str, str] = field(default_factory=dict) 

201 

202 def __post_init__(self) -> None: 

203 self.type = "subtask" 

204 

205 def to_dict(self) -> dict[str, Any]: 

206 d = super().to_dict() 

207 d["prompt"] = self.prompt 

208 d["description"] = self.description 

209 d["agent"] = self.agent 

210 d["model"] = self.model 

211 return d 

212 

213 

214def part_from_dict(data: dict[str, Any]) -> Part: 

215 """Deserialize a part from its stored dict representation.""" 

216 part_type = data.get("type", "") 

217 

218 def _time(d: dict[str, Any] | None) -> TimeInfo | None: 

219 if not d: 

220 return None 

221 return TimeInfo(start=d["start"], end=d.get("end"), compacted=d.get("compacted")) 

222 

223 time = _time(data.get("time")) 

224 metadata = data.get("metadata", {}) 

225 

226 if part_type == "text": 

227 p = TextPart(text=data.get("text", ""), ignored=data.get("ignored", False)) 

228 elif part_type == "reasoning": 

229 p = ReasoningPart(text=data.get("text", "")) 

230 elif part_type == "tool": 

231 sd = data.get("state", {}) 

232 state = ToolState( 

233 status=sd.get("status", "pending"), 

234 input=sd.get("input", {}), 

235 output=sd.get("output", ""), 

236 summary_text=sd.get("summary_text"), 

237 title=sd.get("title", ""), 

238 metadata=sd.get("metadata", {}), 

239 error=sd.get("error"), 

240 time=_time(sd.get("time")), 

241 attachments=sd.get("attachments", []), 

242 ) 

243 p = ToolPart(tool=data.get("tool", ""), call_id=data.get("call_id", ""), state=state) 

244 elif part_type == "step-start": 

245 p = StepStartPart(snapshot=data.get("snapshot", "")) 

246 elif part_type == "step-finish": 

247 p = StepFinishPart( 

248 reason=data.get("reason", "stop"), 

249 snapshot=data.get("snapshot", ""), 

250 tokens=data.get("tokens", {}), 

251 cost=data.get("cost", 0.0), 

252 ) 

253 elif part_type == "compaction": 

254 p = CompactionPart(auto=data.get("auto", False)) 

255 elif part_type == "patch": 

256 p = PatchPart(hash=data.get("hash", ""), files=data.get("files", [])) 

257 elif part_type == "subtask": 

258 p = SubtaskPart( 

259 prompt=data.get("prompt", ""), 

260 description=data.get("description", ""), 

261 agent=data.get("agent", ""), 

262 model=data.get("model", {}), 

263 ) 

264 else: 

265 # Unknown type — create a generic part so we don't crash 

266 p = TextPart(text="") 

267 p.time = time 

268 p.metadata = metadata 

269 return p