"""Tests for ReAct engine compression integration (U3)""" from unittest.mock import AsyncMock, MagicMock, patch import pytest from agentkit.core.compressor import CompressionStrategy, ContextCompressor from agentkit.core.react import ReActEngine from agentkit.llm.protocol import LLMResponse, StreamChunk, TokenUsage, ToolCall # ── Helpers ────────────────────────────────────────── def make_mock_gateway() -> MagicMock: """创建一个 mock LLMGateway""" from agentkit.llm.gateway import LLMGateway gateway = MagicMock(spec=LLMGateway) response = LLMResponse( content="Final answer", model="test-model", usage=TokenUsage(prompt_tokens=10, completion_tokens=10), ) gateway.chat = AsyncMock(return_value=response) return gateway def make_mock_gateway_with_tool_call() -> MagicMock: """创建一个返回 tool_call 的 mock LLMGateway，第二次调用返回最终答案 同时设置 chat 和 chat_stream，使 execute 和 execute_stream 路径都能正常工作。 """ from agentkit.llm.gateway import LLMGateway gateway = MagicMock(spec=LLMGateway) # 第一次调用返回 tool_call，第二次返回最终答案 tool_response = LLMResponse( content="", model="test-model", usage=TokenUsage(prompt_tokens=10, completion_tokens=10), tool_calls=[ ToolCall(id="call_1", name="search", arguments={"query": "test"}), ], ) final_response = LLMResponse( content="Final answer after tool", model="test-model", usage=TokenUsage(prompt_tokens=10, completion_tokens=10), ) gateway.chat = AsyncMock(side_effect=[tool_response, final_response]) # ponytail: chat_stream yields StreamChunk equivalents of the chat responses # so execute_stream (which uses chat_stream) exercises the same tool path. tool_chunk = StreamChunk( content="", model="test-model", tool_calls=[ToolCall(id="call_1", name="search", arguments={"query": "test"})], usage=TokenUsage(prompt_tokens=10, completion_tokens=10), is_final=True, ) final_chunk = StreamChunk( content="Final answer after tool", model="test-model", usage=TokenUsage(prompt_tokens=10, completion_tokens=10), is_final=True, ) async def _stream(**kwargs): # Closure state tracks which response to yield (1st call=tool, 2nd=final) _stream._call_count = getattr(_stream, "_call_count", 0) + 1 if _stream._call_count == 1: yield tool_chunk else: yield final_chunk gateway.chat_stream = _stream return gateway def make_long_messages(count: int = 10, content_length: int = 2000) -> list[dict]: """生成长消息列表用于测试压缩""" messages = [{"role": "system", "content": "You are a helpful assistant."}] for i in range(count): messages.append({ "role": "user", "content": "x" * content_length + f" message {i}", }) messages.append({ "role": "assistant", "content": "y" * content_length + f" reply {i}", }) return messages def make_mock_compressor() -> MagicMock: """创建一个 mock CompressionStrategy""" compressor = MagicMock(spec=CompressionStrategy) compressor.compress = AsyncMock(return_value=[{"role": "user", "content": "compressed"}]) compressor.compress_tool_result = AsyncMock(return_value="compressed tool result") compressor.is_available = MagicMock(return_value=True) return compressor # ── TestBuildToolResultMessage ──────────────────────── class TestBuildToolResultMessage: """_build_tool_result_message 方法测试""" async def test_no_compressor_returns_original(self): engine = ReActEngine(llm_gateway=make_mock_gateway()) result = await engine._build_tool_result_message("tc_1", {"key": "value"}) assert result == { "role": "tool", "tool_call_id": "tc_1", "content": "{'key': 'value'}", } async def test_with_compressor_calls_compress_tool_result(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) result = await engine._build_tool_result_message( "tc_1", {"key": "value"}, compressor=compressor, tool_name="search" ) compressor.compress_tool_result.assert_called_once_with("search", {"key": "value"}) assert result["content"] == "compressed tool result" assert result["role"] == "tool" assert result["tool_call_id"] == "tc_1" async def test_compressor_failure_falls_back(self): compressor = make_mock_compressor() compressor.compress_tool_result = AsyncMock(side_effect=RuntimeError("compression error")) engine = ReActEngine(llm_gateway=make_mock_gateway()) result = await engine._build_tool_result_message( "tc_1", {"key": "value"}, compressor=compressor, tool_name="search" ) # 应该回退到 str(result) assert result["content"] == "{'key': 'value'}" assert result["role"] == "tool" async def test_compressor_receives_tool_name(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) await engine._build_tool_result_message( "tc_1", "some result", compressor=compressor, tool_name="web_crawl" ) compressor.compress_tool_result.assert_called_once_with("web_crawl", "some result") async def test_compressor_without_tool_name_skips_compression(self): """compressor 存在但 tool_name 为 None 时不压缩""" compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) result = await engine._build_tool_result_message( "tc_1", "some result", compressor=compressor, tool_name=None ) compressor.compress_tool_result.assert_not_called() assert result["content"] == "some result" # ── TestShouldCompress ─────────────────────────────── class TestShouldCompress: """_should_compress 辅助方法测试""" def test_no_compressor_returns_false(self): engine = ReActEngine(llm_gateway=make_mock_gateway()) conversation = [{"role": "user", "content": "x" * 100000}] assert engine._should_compress(conversation, None) is False def test_short_conversation_returns_false(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) conversation = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there"}, ] assert engine._should_compress(conversation, compressor) is False def test_long_conversation_returns_true(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) # 8000 tokens * 4 chars/token = 32000 chars needed conversation = [{"role": "user", "content": "x" * 40000}] assert engine._should_compress(conversation, compressor) is True def test_boundary_at_threshold(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) # Exactly 8000 tokens = 32000 chars → should NOT trigger (> not >=) conversation = [{"role": "user", "content": "x" * 32000}] assert engine._should_compress(conversation, compressor) is False def test_just_above_threshold(self): compressor = make_mock_compressor() engine = ReActEngine(llm_gateway=make_mock_gateway()) # 32001 chars → 32001//4 = 8000 tokens, still not > 8000 # 32004 chars → 32004//4 = 8001 tokens, > 8000 conversation = [{"role": "user", "content": "x" * 32004}] assert engine._should_compress(conversation, compressor) is True # ── TestReActLoopCompression ───────────────────────── class TestReActLoopCompression: """ReAct 循环内压缩集成测试""" async def test_tool_results_compressed_in_loop(self): """工具结果在循环中被压缩后拼入 conversation""" compressor = make_mock_compressor() compressor.compress_tool_result = AsyncMock(return_value="COMPRESSED:search_result") gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) # 注册一个 mock tool mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="original search result") result = await engine.execute( messages=[{"role": "user", "content": "Search for test"}], tools=[mock_tool], compressor=compressor, ) # 验证 compress_tool_result 被调用 compressor.compress_tool_result.assert_called_once_with("search", "original search result") async def test_incremental_compression_triggered(self): """长对话触发增量压缩 compressor.compress()""" compressor = make_mock_compressor() # 让 compress 返回压缩后的短对话 compressor.compress = AsyncMock(return_value=[ {"role": "system", "content": "Summary"}, {"role": "user", "content": "recent"}, ]) gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) # 构造一个很长的初始消息，使 conversation 超过阈值 long_content = "x" * 40000 mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="result") result = await engine.execute( messages=[{"role": "user", "content": long_content}], tools=[mock_tool], compressor=compressor, ) # 验证增量压缩被触发（compress 被调用了至少一次） # 初始 compress 在 L218-222，增量 compress 在工具结果后 assert compressor.compress.call_count >= 1 async def test_incremental_compression_failure_handled(self): """compressor.compress() 异常时循环继续""" compressor = make_mock_compressor() # 第一次 compress 调用成功（初始压缩），增量压缩时失败 call_count = 0 async def compress_side_effect(messages): nonlocal call_count call_count += 1 if call_count > 1: raise RuntimeError("Incremental compression failed") return messages compressor.compress = AsyncMock(side_effect=compress_side_effect) gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) # 构造长消息触发增量压缩 long_content = "x" * 40000 mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="result") # 不应该抛出异常 result = await engine.execute( messages=[{"role": "user", "content": long_content}], tools=[mock_tool], compressor=compressor, ) # 应该正常返回结果 assert result.output == "Final answer after tool" async def test_no_compressor_backward_compatible(self): """compressor=None 时行为与之前完全一致""" gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="search result data") result = await engine.execute( messages=[{"role": "user", "content": "Search for test"}], tools=[mock_tool], compressor=None, ) assert result.output == "Final answer after tool" assert result.status == "success" assert len(result.trajectory) == 2 # 1 tool_call + 1 final_answer async def test_execute_stream_with_compressor(self): """execute_stream 模式下压缩正常工作""" compressor = make_mock_compressor() compressor.compress_tool_result = AsyncMock(return_value="COMPRESSED:result") gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="original result") events = [] async for event in engine.execute_stream( messages=[{"role": "user", "content": "Search for test"}], tools=[mock_tool], compressor=compressor, ): events.append(event) # 验证 compress_tool_result 被调用 compressor.compress_tool_result.assert_called_once_with("search", "original result") # 验证事件流完整 event_types = [e.event_type for e in events] assert "thinking" in event_types assert "tool_call" in event_types assert "tool_result" in event_types assert "final_answer" in event_types async def test_execute_stream_incremental_compression(self): """execute_stream 模式下增量压缩触发""" compressor = make_mock_compressor() compressor.compress = AsyncMock(return_value=[ {"role": "system", "content": "Summary"}, {"role": "user", "content": "recent"}, ]) gateway = make_mock_gateway_with_tool_call() engine = ReActEngine(llm_gateway=gateway) long_content = "x" * 40000 mock_tool = MagicMock() mock_tool.name = "search" mock_tool.safe_execute = AsyncMock(return_value="result") events = [] async for event in engine.execute_stream( messages=[{"role": "user", "content": long_content}], tools=[mock_tool], compressor=compressor, ): events.append(event) # 验证增量压缩被触发 assert compressor.compress.call_count >= 1 async def test_context_compressor_satisfies_protocol(self): """ContextCompressor 满足 CompressionStrategy Protocol""" compressor = ContextCompressor() assert isinstance(compressor, CompressionStrategy) async def test_context_compressor_compress_tool_result_default(self): """ContextCompressor.compress_tool_result 默认返回 str(result)""" compressor = ContextCompressor() result = await compressor.compress_tool_result("search", {"key": "value"}) assert result == "{'key': 'value'}" # ── TestOTelSpanLifecycle (P0 fix: span leak) ────────────────── class TestOTelSpanLifecycle: """测试 OTel span 生命周期 — 异常时 span 必须正确关闭""" async def test_span_closed_on_success(self): """正常执行时 span 被正确关闭""" gateway = make_mock_gateway() engine = ReActEngine(llm_gateway=gateway) mock_span = MagicMock() mock_span_cm = MagicMock() mock_span_cm.__enter__ = MagicMock(return_value=mock_span) mock_span_cm.__exit__ = MagicMock(return_value=False) with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \ patch("agentkit.core.react._OTEL_AVAILABLE", True): await engine.execute(messages=[{"role": "user", "content": "hello"}]) # __exit__ should have been called mock_span_cm.__exit__.assert_called_once() async def test_span_closed_on_exception(self): """LLM 抛出异常时 span 仍被正确关闭""" gateway = make_mock_gateway() gateway.chat = AsyncMock(side_effect=RuntimeError("LLM error")) engine = ReActEngine(llm_gateway=gateway) mock_span = MagicMock() mock_span_cm = MagicMock() mock_span_cm.__enter__ = MagicMock(return_value=mock_span) mock_span_cm.__exit__ = MagicMock(return_value=False) with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \ patch("agentkit.core.react._OTEL_AVAILABLE", True): with pytest.raises(RuntimeError, match="LLM error"): await engine.execute(messages=[{"role": "user", "content": "hello"}]) # __exit__ must have been called even though exception was raised mock_span_cm.__exit__.assert_called_once() async def test_span_attributes_set_on_success(self): """正常执行时 span 属性被设置""" gateway = make_mock_gateway() engine = ReActEngine(llm_gateway=gateway) mock_span = MagicMock() mock_span_cm = MagicMock() mock_span_cm.__enter__ = MagicMock(return_value=mock_span) mock_span_cm.__exit__ = MagicMock(return_value=False) with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \ patch("agentkit.core.react._OTEL_AVAILABLE", True): await engine.execute(messages=[{"role": "user", "content": "hello"}]) # Verify span attributes were set mock_span.set_attribute.assert_any_call("agent.total_steps", 1) mock_span.set_attribute.assert_any_call("agent.total_tokens", 20) mock_span.set_attribute.assert_any_call("agent.outcome", "success") async def test_no_span_when_otel_unavailable(self): """_OTEL_AVAILABLE=False 时不创建 span""" gateway = make_mock_gateway() engine = ReActEngine(llm_gateway=gateway) with patch("agentkit.core.react._OTEL_AVAILABLE", False), \ patch("agentkit.core.react.start_span") as mock_start_span: await engine.execute(messages=[{"role": "user", "content": "hello"}]) # start_span should not be called when OTel is unavailable mock_start_span.assert_not_called()