Coverage for src / harnessutils / models / message.py: 79%

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

2 

3from dataclasses import dataclass, field 

4from typing import Any, Literal 

5 

6from harnessutils.models.parts import Part 

7from harnessutils.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 harnessutils.models.usage import CacheUsage, Usage 

90 

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

92 tokens = Usage( 

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

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

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

96 cache=CacheUsage( 

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

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

99 ), 

100 ) 

101 

102 return cls( 

103 id=data["id"], 

104 role=data["role"], 

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

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

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

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

109 tokens=tokens, 

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

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

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

113 )