Coverage for agentos/tests/test_conversation.py: 0%

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

2 

3from __future__ import annotations 

4 

5import pytest 

6from agentos.conversation.conversation import ( 

7 ConversationManager, 

8 ConversationConfig, 

9 MessageRole, 

10 TrimStrategy, 

11) 

12 

13 

14@pytest.fixture 

15def conv(): 

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

17 

18 

19def test_add_message(conv): 

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

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

22 assert msg.role == MessageRole.USER 

23 assert msg.content == "Hello" 

24 assert conv.message_count == 1 

25 

26 

27def test_add_many(conv): 

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

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

30 assert len(msgs) == 2 

31 assert conv.message_count == 2 

32 

33 

34def test_get_context(conv): 

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

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

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

38 ctx = conv.get_context() 

39 assert len(ctx) == 2 

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

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

42 

43 

44def test_fifo_trim(conv): 

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

46 conv.config.max_messages = 5 

47 conv.config.preserve_last_n = 2 

48 for i in range(10): 

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

50 assert conv.message_count <= 5 

51 

52 

53def test_preserve_system(conv): 

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

55 conv.config.max_messages = 4 

56 conv.config.preserve_last_n = 1 

57 conv.config.preserve_system = True 

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

59 for i in range(8): 

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

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

62 assert len(system_msgs) >= 1 

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

64 

65 

66def test_token_tracking(conv): 

67 """token 统计。""" 

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

69 assert conv.token_count > 50 

70 

71 

72def test_fork_and_switch(conv): 

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

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

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

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

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

78 

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

80 assert conv.message_count == 3 

81 

82 conv.switch_branch(snapshot.snapshot_id) 

83 assert conv.message_count == 2 

84 

85 

86def test_fork_branch_not_found(conv): 

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

88 with pytest.raises(KeyError): 

89 conv.switch_branch("nonexistent") 

90 

91 

92def test_merge_branch_append(conv): 

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

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

95 snapshot = conv.fork("side") 

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

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

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

99 assert conv.message_count >= 3 

100 

101 

102def test_clear(conv): 

103 """清空对话。""" 

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

105 conv.clear() 

106 assert conv.message_count == 0 

107 assert conv.token_count == 0 

108 

109 

110def test_clear_keep_system(conv): 

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

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

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

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

115 conv.clear(keep_system=True) 

116 assert conv.message_count == 1 

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

118 

119 

120def test_stats_tracking(conv): 

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

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

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

124 assert conv.stats.total_messages == 2 

125 assert conv.stats.total_tokens > 0 

126 assert conv.stats.oldest_timestamp > 0 

127 

128 

129def test_message_id_unique(conv): 

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

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

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

133 assert len(ids) == 5 

134 

135 

136def test_empty_context(conv): 

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

138 ctx = conv.get_context() 

139 assert ctx == [] 

140 

141 

142def test_get_system_prompt(conv): 

143 """提取 system prompt。""" 

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

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

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

147 

148 

149def test_importance_weighted_trim(): 

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

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

152 c = ConversationManager(cfg) 

153 for i in range(10): 

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

155 m.importance = float(i % 3) 

156 assert c.message_count <= 5 

157 

158 

159def test_trim_stats_increment(conv): 

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

161 conv.config.max_messages = 3 

162 conv.config.preserve_last_n = 1 

163 for i in range(8): 

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

165 assert conv.stats.trim_count > 0