1"""Tests for agentos.conversation.conversation."""

2

3 from __future__ import annotations

4

5 import pytest

6 from agentos.conversation.conversation import (

7 ConversationManager,

8 ConversationConfig,

9 Message,

10 MessageRole,

11 TrimStrategy,

12)

13

14

15@pytest.fixture

16 def conv():

17 return ConversationManager(ConversationConfig(max_messages=10, max_tokens=100000))

18

19

20 def test_add_message(conv):

21 """添加单条消息。"""

22 msg = conv.add("user", "Hello")

23 assert msg.role == MessageRole.USER

24 assert msg.content == "Hello"

25 assert conv.message_count == 1

26

27

28 def test_add_many(conv):

29 """批量添加消息。"""

30 msgs = conv.add_many([("user", "Hi"), ("assistant", "Hello there")])

31 assert len(msgs) == 2

32 assert conv.message_count == 2

33

34

35 def test_get_context(conv):

36 """获取 OpenAI 兼容上下文。"""

37 conv.add("system", "You are helpful.")

38 conv.add("user", "Question?")

39 ctx = conv.get_context()

40 assert len(ctx) == 2

41 assert ctx[0]["role"] == "system"

42 assert ctx[1]["role"] == "user"

43

44

45 def test_fifo_trim(conv):

46 """FIFO 裁剪：超出 max_messages 时移除最旧消息。"""

47 conv.config.max_messages = 5

48 conv.config.preserve_last_n = 2

49 for i in range(10):

50 conv.add("user", f"msg{i}")

51 assert conv.message_count <= 5

52

53

54 def test_preserve_system(conv):

55 """裁剪时保留 system 消息。"""

56 conv.config.max_messages = 4

57 conv.config.preserve_last_n = 1

58 conv.config.preserve_system = True

59 conv.add("system", "System prompt")

60 for i in range(8):

61 conv.add("user", f"msg{i}")

62 system_msgs = [m for m in conv._messages if m.role == MessageRole.SYSTEM]

63 assert len(system_msgs) >= 1

64 assert system_msgs[0].content == "System prompt"

65

66

67 def test_token_tracking(conv):

68 """token 统计。"""

69 conv.add("user", "A" * 300)

70 assert conv.token_count > 50

71

72

73 def test_fork_and_switch(conv):

74 """对话分支：fork -> switch -> 验证。"""

75 conv.add("user", "msg1")

76 conv.add("assistant", "reply1")

77 snapshot = conv.fork("test-branch")

78 assert snapshot.label == "test-branch"

79

80 conv.add("user", "msg2")

81 assert conv.message_count == 3

82

83 conv.switch_branch(snapshot.snapshot_id)

84 assert conv.message_count == 2

85

86

87 def test_fork_branch_not_found(conv):

88 """切换不存在分支抛出 KeyError。"""

89 with pytest.raises(KeyError):

90 conv.switch_branch("nonexistent")

91

92

93 def test_merge_branch_append(conv):

94 """合并分支（追加模式）。"""

95 conv.add("user", "msg1")

96 snapshot = conv.fork("side")

97 conv.add("user", "msg2")

98 conv.add("user", "msg3")

99 conv.merge_branch(snapshot.snapshot_id, strategy="append")

100 assert conv.message_count >= 3

101

102

103 def test_clear(conv):

104 """清空对话。"""

105 conv.add_many([("user", "a"), ("assistant", "b"), ("user", "c")])

106 conv.clear()

107 assert conv.message_count == 0

108 assert conv.token_count == 0

109

110

111 def test_clear_keep_system(conv):

112 """清空但保留 system 消息。"""

113 conv.add("system", "sys")

114 conv.add("user", "q")

115 conv.add("assistant", "a")

116 conv.clear(keep_system=True)

117 assert conv.message_count == 1

118 assert conv._messages[0].content == "sys"

119

120

121 def test_stats_tracking(conv):

122 """统计数据正确累加。"""

123 conv.add("user", "hello world")

124 conv.add("assistant", "hi there")

125 assert conv.stats.total_messages == 2

126 assert conv.stats.total_tokens > 0

127 assert conv.stats.oldest_timestamp > 0

128

129

130 def test_message_id_unique(conv):

131 """每条消息 ID 唯一。"""

132 msgs = conv.add_many([("user", f"msg{i}") for i in range(5)])

133 ids = {m.message_id for m in msgs}

134 assert len(ids) == 5

135

136

137 def test_empty_context(conv):

138 """空对话上下文。"""

139 ctx = conv.get_context()

140 assert ctx == []

141

142

143 def test_get_system_prompt(conv):

144 """提取 system prompt。"""

145 conv.add("system", "Be concise.")

146 conv.add("user", "Q")

147 assert conv.get_system_prompt() == "Be concise."

148

149

150 def test_importance_weighted_trim():

151 """重要性加权裁剪。"""

152 cfg = ConversationConfig(max_messages=5, trim_strategy=TrimStrategy.IMPORTANCE_WEIGHTED, preserve_last_n=2)

153 c = ConversationManager(cfg)

154 for i in range(10):

155 m = c.add("user", f"msg{i}")

156 m.importance = float(i % 3)

157 assert c.message_count <= 5

158

159

160 def test_trim_stats_increment(conv):

161 """裁剪后 trim_count 递增。"""

162 conv.config.max_messages = 3

163 conv.config.preserve_last_n = 1

164 for i in range(8):

165 conv.add("user", f"msg{i}")

166 assert conv.stats.trim_count > 0