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 def get_velocity(self) -> ConversationVelocity | None: 

30 """Get velocity tracker from metadata. 

31 

32 Returns: 

33 ConversationVelocity instance or None if not tracked 

34 """ 

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

36 if velocity_data is None: 

37 return None 

38 

39 return ConversationVelocity.from_dict(velocity_data) 

40 

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

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

43 

44 Args: 

45 tokens_added: Tokens added in this turn 

46 """ 

47 velocity = self.get_velocity() 

48 if velocity is None: 

49 velocity = ConversationVelocity() 

50 

51 velocity.add_delta(tokens_added) 

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

53 

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

55 """Get quality history from metadata. 

56 

57 Returns: 

58 QualityHistory instance or None if not tracked 

59 """ 

60 from harnessutils.quality import QualityHistory 

61 

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

63 if history_data is None: 

64 return None 

65 

66 return QualityHistory.from_dict(history_data) 

67 

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

69 """Add quality snapshot to history. 

70 

71 Args: 

72 snapshot: Quality snapshot to add 

73 """ 

74 from harnessutils.quality import QualityHistory 

75 

76 history = self.get_quality_history() 

77 if history is None: 

78 history = QualityHistory() 

79 

80 history.add_snapshot(snapshot) 

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

82 

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

84 """Convert conversation to dictionary for storage. 

85 

86 Returns: 

87 Dictionary representation 

88 """ 

89 return { 

90 "id": self.id, 

91 "project_id": self.project_id, 

92 "created": self.created, 

93 "updated": self.updated, 

94 "pending_summarization": self.pending_summarization, 

95 "metadata": self.metadata, 

96 } 

97 

98 @classmethod 

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

100 """Create conversation from dictionary. 

101 

102 Args: 

103 data: Dictionary representation 

104 

105 Returns: 

106 Conversation instance 

107 """ 

108 return cls( 

109 id=data["id"], 

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

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

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

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

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

115 )