Coverage for src / harness_utils / models / message.py: 76%

1"""Message model with part-based decomposition.""" 

2 

3from dataclasses import dataclass, field 

4from typing import Any, Literal 

5 

6from harness_utils.models.parts import Part 

7from harness_utils.models.usage import Usage 

8 

9 

10@dataclass 

11class Message: 

12 """A message in the conversation. 

13 

14 Messages are decomposed into parts for granular compaction. 

15 Each message can contain multiple parts (text, tool calls, reasoning, etc.). 

16 """ 

17 

18 id: str 

19 role: Literal["user", "assistant"] 

20 parts: list[Part] = field(default_factory=list) 

21 parent_id: str | None = None 

22 summary: bool = False # Is this a summary message? 

23 agent: str | None = None 

24 model: dict[str, str] | None = None 

25 tokens: Usage | None = None 

26 cost: float = 0.0 

27 error: str | None = None 

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

29 

30 def add_part(self, part: Part) -> None: 

31 """Add a part to this message. 

32 

33 Args: 

34 part: The part to add 

35 """ 

36 self.parts.append(part) 

37 

38 def has_partial_output(self) -> bool: 

39 """Check if message has any partial output despite errors. 

40 

41 Returns: 

42 True if there are text parts even with errors 

43 """ 

44 return any(p.type == "text" for p in self.parts) 

45 

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

47 """Convert message to dictionary for storage. 

48 

49 Returns: 

50 Dictionary representation 

51 """ 

52 data: dict[str, Any] = { 

53 "id": self.id, 

54 "role": self.role, 

55 "parent_id": self.parent_id, 

56 "summary": self.summary, 

57 "agent": self.agent, 

58 "model": self.model, 

59 "cost": self.cost, 

60 "error": self.error, 

61 "metadata": self.metadata, 

62 } 

63 

64 if self.tokens: 

65 data["tokens"] = { 

66 "input": self.tokens.input, 

67 "output": self.tokens.output, 

68 "reasoning": self.tokens.reasoning, 

69 "cache": { 

70 "read": self.tokens.cache.read, 

71 "write": self.tokens.cache.write, 

72 }, 

73 } 

74 

75 return data 

76 

77 @classmethod 

78 def from_dict(cls, data: dict[str, Any]) -> "Message": 

79 """Create message from dictionary. 

80 

81 Args: 

82 data: Dictionary representation 

83 

84 Returns: 

85 Message instance 

86 """ 

87 tokens = None 

88 if "tokens" in data and data["tokens"]: 

89 from harness_utils.models.usage import CacheUsage, Usage 

90 cache_data = data["tokens"].get("cache", {}) 

91 tokens = Usage( 

92 input=data["tokens"].get("input", 0), 

93 output=data["tokens"].get("output", 0), 

94 reasoning=data["tokens"].get("reasoning", 0), 

95 cache=CacheUsage( 

96 read=cache_data.get("read", 0), 

97 write=cache_data.get("write", 0), 

98 ), 

99 ) 

100 

101 return cls( 

102 id=data["id"], 

103 role=data["role"], 

104 parent_id=data.get("parent_id"), 

105 summary=data.get("summary", False), 

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

107 model=data.get("model"), 

108 tokens=tokens, 

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

110 error=data.get("error"), 

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

112 )