Coverage for src / harnessutils / models / conversation.py: 100%

1"""Conversation model.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6from typing import TYPE_CHECKING, Any 

7 

8from harnessutils.models.velocity import ConversationVelocity 

9 

10if TYPE_CHECKING: 

11 from harnessutils.quality import QualityHistory, QualitySnapshot 

12 

13 

14@dataclass 

15class Conversation: 

16 """A conversation containing multiple messages. 

17 

18 Conversations track the overall context and metadata for a series 

19 of messages between user and assistant. 

20 """ 

21 

22 id: str 

23 project_id: str | None = None 

24 created: int | None = None # Unix timestamp in milliseconds 

25 updated: int | None = None # Unix timestamp in milliseconds 

26 pending_summarization: bool = False 

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

28 

29 @property 

30 def excluded_outputs(self) -> list[str]: 

31 """Call IDs of tool outputs excluded from API payloads.""" 

32 return self.metadata.get("excluded_outputs", []) 

33 

34 def add_excluded_outputs(self, call_ids: list[str]) -> None: 

35 """Accumulate excluded call IDs (deduplicates).""" 

36 existing = set(self.excluded_outputs) 

37 existing.update(call_ids) 

38 self.metadata["excluded_outputs"] = list(existing) 

39 

40 def get_velocity(self) -> ConversationVelocity | None: 

41 """Get velocity tracker from metadata. 

42 

43 Returns: 

44 ConversationVelocity instance or None if not tracked 

45 """ 

46 velocity_data = self.metadata.get("velocity") 

47 if velocity_data is None: 

48 return None 

49 

50 return ConversationVelocity.from_dict(velocity_data) 

51 

52 def update_velocity(self, tokens_added: int) -> None: 

53 """Update velocity tracker with new token delta. 

54 

55 Args: 

56 tokens_added: Tokens added in this turn 

57 """ 

58 velocity = self.get_velocity() 

59 if velocity is None: 

60 velocity = ConversationVelocity() 

61 

62 velocity.add_delta(tokens_added) 

63 self.metadata["velocity"] = velocity.to_dict() 

64 

65 def get_quality_history(self) -> QualityHistory | None: 

66 """Get quality history from metadata. 

67 

68 Returns: 

69 QualityHistory instance or None if not tracked 

70 """ 

71 from harnessutils.quality import QualityHistory 

72 

73 history_data = self.metadata.get("quality_history") 

74 if history_data is None: 

75 return None 

76 

77 return QualityHistory.from_dict(history_data) 

78 

79 def update_quality_history(self, snapshot: QualitySnapshot) -> None: 

80 """Add quality snapshot to history. 

81 

82 Args: 

83 snapshot: Quality snapshot to add 

84 """ 

85 from harnessutils.quality import QualityHistory 

86 

87 history = self.get_quality_history() 

88 if history is None: 

89 history = QualityHistory() 

90 

91 history.add_snapshot(snapshot) 

92 self.metadata["quality_history"] = history.to_dict() 

93 

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

95 """Convert conversation to dictionary for storage. 

96 

97 Returns: 

98 Dictionary representation 

99 """ 

100 return { 

101 "id": self.id, 

102 "project_id": self.project_id, 

103 "created": self.created, 

104 "updated": self.updated, 

105 "pending_summarization": self.pending_summarization, 

106 "metadata": self.metadata, 

107 } 

108 

109 @classmethod 

110 def from_dict(cls, data: dict[str, Any]) -> Conversation: 

111 """Create conversation from dictionary. 

112 

113 Args: 

114 data: Dictionary representation 

115 

116 Returns: 

117 Conversation instance 

118 """ 

119 return cls( 

120 id=data["id"], 

121 project_id=data.get("project_id"), 

122 created=data.get("created"), 

123 updated=data.get("updated"), 

124 pending_summarization=data.get("pending_summarization", False), 

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

126 )