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

1"""Conversation model.""" 

2 

3from dataclasses import dataclass, field 

4from typing import Any 

5 

6from harnessutils.models.velocity import ConversationVelocity 

7 

8 

9@dataclass 

10class Conversation: 

11 """A conversation containing multiple messages. 

12 

13 Conversations track the overall context and metadata for a series 

14 of messages between user and assistant. 

15 """ 

16 

17 id: str 

18 project_id: str | None = None 

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

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

21 pending_summarization: bool = False 

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

23 

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

25 """Get velocity tracker from metadata. 

26 

27 Returns: 

28 ConversationVelocity instance or None if not tracked 

29 """ 

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

31 if velocity_data is None: 

32 return None 

33 

34 return ConversationVelocity.from_dict(velocity_data) 

35 

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

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

38 

39 Args: 

40 tokens_added: Tokens added in this turn 

41 """ 

42 velocity = self.get_velocity() 

43 if velocity is None: 

44 velocity = ConversationVelocity() 

45 

46 velocity.add_delta(tokens_added) 

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

48 

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

50 """Convert conversation to dictionary for storage. 

51 

52 Returns: 

53 Dictionary representation 

54 """ 

55 return { 

56 "id": self.id, 

57 "project_id": self.project_id, 

58 "created": self.created, 

59 "updated": self.updated, 

60 "pending_summarization": self.pending_summarization, 

61 "metadata": self.metadata, 

62 } 

63 

64 @classmethod 

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

66 """Create conversation from dictionary. 

67 

68 Args: 

69 data: Dictionary representation 

70 

71 Returns: 

72 Conversation instance 

73 """ 

74 return cls( 

75 id=data["id"], 

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

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

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

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

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

81 )