377 lines
13 KiB
Python
377 lines
13 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from langchain_core.messages import ToolMessage
|
|
from artemis.agents.history_analyzer.history_analyzer import HistoryAnalyzer
|
|
from artemis.context import ArtemisContext
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_analyzer_no_history():
|
|
# Mock context with a data engine that has no steps
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.get_agent_friendly_steps.return_value = []
|
|
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
result = await analyzer.run("What happened?")
|
|
assert result == "No history recorded for this session yet."
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_analyzer_simple_query_no_tool_call():
|
|
# Mock context and data engine
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/fake_traces"
|
|
|
|
steps = [
|
|
{
|
|
"step_number": 1,
|
|
"relative_time": "1.2s",
|
|
"summary": "Opened the settings app",
|
|
},
|
|
{
|
|
"step_number": 2,
|
|
"relative_time": "4.5s",
|
|
"summary": "Toggled the wifi switch",
|
|
},
|
|
]
|
|
mock_ctx.data_engine.get_agent_friendly_steps.return_value = steps
|
|
|
|
# Mock LLM to return a direct text response (no tool calls)
|
|
mock_response = MagicMock()
|
|
mock_response.content = "First you opened settings, then you toggled wifi."
|
|
mock_response.tool_calls = []
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.bind_tools.return_value = mock_llm
|
|
|
|
async def mock_astream(*args, **kwargs):
|
|
yield mock_response
|
|
|
|
mock_llm.astream.side_effect = mock_astream
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.history_analyzer.history_analyzer.get_llm",
|
|
return_value=mock_llm,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=False), # Force system prompt fallback
|
|
):
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
result = await analyzer.run("Summarize the session")
|
|
|
|
assert result == "First you opened settings, then you toggled wifi."
|
|
assert mock_llm.astream.called
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_analyzer_detailed_query_with_tool_call():
|
|
# Mock context and data engine with full details
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/fake_traces"
|
|
|
|
steps = [
|
|
{
|
|
"step_number": 1,
|
|
"relative_time": "1.2s",
|
|
"summary": "Opened the settings app",
|
|
"action_taken": {"action": "tap", "coordinates": [500, 600]},
|
|
"operator_raw_thinking": "Need to open settings to configure wifi",
|
|
"last_execution_result": {"status": "success"},
|
|
},
|
|
{
|
|
"step_number": 2,
|
|
"relative_time": "4.5s",
|
|
"summary": "Toggled the wifi switch",
|
|
"action_taken": {"action": "tap", "coordinates": [100, 200]},
|
|
"operator_raw_thinking": "Toggle the switch to turn wifi on",
|
|
"last_execution_result": {"status": "success"},
|
|
},
|
|
]
|
|
mock_ctx.data_engine.get_agent_friendly_steps.return_value = steps
|
|
mock_ctx.data_engine.get_agent_friendly_steps_in_range.side_effect = lambda s, e: [
|
|
st for st in steps if s <= st["step_number"] <= e
|
|
]
|
|
|
|
# Turn 1 LLM response: Tool call to replay_steps for step 2
|
|
mock_response_turn_1 = MagicMock()
|
|
mock_response_turn_1.content = ""
|
|
mock_tool_call = {
|
|
"name": "replay_steps",
|
|
"args": {"start_step": 2, "end_step": 2},
|
|
"id": "call_123456",
|
|
}
|
|
mock_response_turn_1.tool_calls = [mock_tool_call]
|
|
|
|
# Turn 2 LLM response: Final answer using the details
|
|
mock_response_turn_2 = MagicMock()
|
|
mock_response_turn_2.content = (
|
|
"In step 2, the operator was thinking: 'Toggle the switch to turn wifi on'."
|
|
)
|
|
mock_response_turn_2.tool_calls = []
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.bind_tools.return_value = mock_llm
|
|
responses = [mock_response_turn_1, mock_response_turn_2]
|
|
call_count = 0
|
|
|
|
async def mock_astream(*args, **kwargs):
|
|
nonlocal call_count
|
|
if call_count < len(responses):
|
|
yield responses[call_count]
|
|
call_count += 1
|
|
|
|
mock_llm.astream.side_effect = mock_astream
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.history_analyzer.history_analyzer.get_llm",
|
|
return_value=mock_llm,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=False), # Force fallback prompt
|
|
):
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
result = await analyzer.run("What was the operator thinking in step 2?")
|
|
|
|
assert (
|
|
result == "In step 2, the operator was thinking: 'Toggle the switch to turn wifi on'."
|
|
)
|
|
assert mock_llm.astream.call_count == 2
|
|
|
|
# Verify the tool was invoked with correct parameters internally
|
|
tool_msg = mock_llm.astream.call_args_list[1][0][0][
|
|
-2
|
|
] # ToolMessage is second to last after final response is appended
|
|
assert isinstance(tool_msg, ToolMessage)
|
|
assert tool_msg.tool_call_id == "call_123456"
|
|
assert tool_msg.name == "replay_steps"
|
|
|
|
# The shared replay renderer: step header, screen description, thinking.
|
|
mock_ctx.data_engine.get_agent_friendly_steps_in_range.assert_called_once_with(2, 2)
|
|
assert "- **Step 2 (Start: 4.5s)**" in tool_msg.content
|
|
assert "[Screen]: Toggled the wifi switch" in tool_msg.content
|
|
assert "Toggle the switch to turn wifi on" in tool_msg.content
|
|
assert "**Step 1" not in tool_msg.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_analyzer_read_note_tool_call():
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/fake_traces"
|
|
|
|
steps = [{"step_number": 1, "relative_time": "1.0s", "summary": "Initiated"}]
|
|
mock_ctx.data_engine.get_agent_friendly_steps.return_value = steps
|
|
|
|
# Turn 1 LLM response: Tool call to list_notes
|
|
mock_response_turn_1 = MagicMock()
|
|
mock_response_turn_1.content = ""
|
|
mock_tool_call_list = {"name": "list_notes", "args": {}, "id": "call_list"}
|
|
mock_response_turn_1.tool_calls = [mock_tool_call_list]
|
|
|
|
# Turn 2 LLM response: Tool call to read_note for tactical_plan
|
|
mock_response_turn_2 = MagicMock()
|
|
mock_response_turn_2.content = ""
|
|
mock_tool_call_read = {
|
|
"name": "read_note",
|
|
"args": {"key": "tactical_plan"},
|
|
"id": "call_read",
|
|
}
|
|
mock_response_turn_2.tool_calls = [mock_tool_call_read]
|
|
|
|
# Turn 3 LLM response: Final answer
|
|
mock_response_turn_3 = MagicMock()
|
|
mock_response_turn_3.content = "The plan was to toggle wifi."
|
|
mock_response_turn_3.tool_calls = []
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.bind_tools.return_value = mock_llm
|
|
responses = [
|
|
mock_response_turn_1,
|
|
mock_response_turn_2,
|
|
mock_response_turn_3,
|
|
]
|
|
call_count = 0
|
|
|
|
async def mock_astream(*args, **kwargs):
|
|
nonlocal call_count
|
|
if call_count < len(responses):
|
|
yield responses[call_count]
|
|
call_count += 1
|
|
|
|
mock_llm.astream.side_effect = mock_astream
|
|
|
|
# Mock the file system read using patch
|
|
fake_plan_content = "Goal: Toggle wifi"
|
|
|
|
def mock_exists(path):
|
|
# Let Path("/tmp/fake_traces/notes").exists() and notes/tactical_plan.md exist
|
|
return True
|
|
|
|
def mock_read_text(self, encoding="utf-8"):
|
|
return fake_plan_content
|
|
|
|
def mock_glob(self, pattern):
|
|
return [Path("/tmp/fake_traces/notes/tactical_plan.md")]
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.history_analyzer.history_analyzer.get_llm",
|
|
return_value=mock_llm,
|
|
),
|
|
patch("pathlib.Path.exists", mock_exists),
|
|
patch("pathlib.Path.read_text", mock_read_text),
|
|
patch("pathlib.Path.glob", mock_glob),
|
|
):
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
result = await analyzer.run("What is the saved tactical plan?")
|
|
|
|
assert result == "The plan was to toggle wifi."
|
|
assert mock_llm.astream.call_count == 3
|
|
|
|
# Verify list_notes response
|
|
list_msg = mock_llm.astream.call_args_list[1][0][0][
|
|
3
|
|
] # Index 3 contains the ToolMessage for list_notes
|
|
assert isinstance(list_msg, ToolMessage)
|
|
assert list_msg.tool_call_id == "call_list"
|
|
assert list_msg.content == "Here are all the notes:\n- tactical_plan (1 lines)"
|
|
|
|
# Verify read_note response
|
|
read_msg = mock_llm.astream.call_args_list[2][0][0][
|
|
5
|
|
] # Index 5 contains the ToolMessage for read_note
|
|
assert isinstance(read_msg, ToolMessage)
|
|
assert read_msg.tool_call_id == "call_read"
|
|
assert (
|
|
read_msg.content == "Successfully read note 'tactical_plan'. 'tactical_plan' note"
|
|
f" content:\n{fake_plan_content}"
|
|
)
|
|
|
|
|
|
def test_history_analyzer_robust_tools_behavior():
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/fake_traces"
|
|
|
|
steps = [
|
|
{
|
|
"step_number": 1,
|
|
"relative_time": "1.0s",
|
|
"summary": "Action 1",
|
|
"action_taken": {"action": "tap"},
|
|
}
|
|
]
|
|
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
mock_ctx.data_engine.get_agent_friendly_steps_in_range.return_value = steps
|
|
|
|
# 1. The shared replay tool coerces string step numbers
|
|
replay_tool = next(t for t in analyzer._build_tools() if t.name == "replay_steps")
|
|
result_details = replay_tool.invoke({"start_step": "1", "end_step": "1"})
|
|
mock_ctx.data_engine.get_agent_friendly_steps_in_range.assert_called_once_with(1, 1)
|
|
assert "- **Step 1 (Start: 1.0s)**" in result_details
|
|
assert "[Screen]: Action 1" in result_details
|
|
|
|
# 2. Test read_note with .md suffix
|
|
from artemis.tools.scratchpad import get_read_note_tool_pure
|
|
|
|
read_tool = get_read_note_tool_pure(analyzer.ctx)
|
|
|
|
def mock_exists(path):
|
|
return "tactical_plan.md" in str(path)
|
|
|
|
with (
|
|
patch("pathlib.Path.exists", mock_exists),
|
|
patch("pathlib.Path.read_text", return_value="Tactical plan content"),
|
|
):
|
|
result_read = read_tool.invoke({"key": "tactical_plan.md"})
|
|
assert (
|
|
result_read == "Successfully read note 'tactical_plan.md'. 'tactical_plan.md'"
|
|
" note content:\nTactical plan content"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_analyzer_integration_with_task_tree():
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/fake_traces"
|
|
|
|
steps = [
|
|
{
|
|
"step_id": "step_1",
|
|
"step_number": 1,
|
|
"relative_time": "1.0s",
|
|
"summary": "Step 1",
|
|
},
|
|
]
|
|
mock_ctx.data_engine.get_agent_friendly_steps.return_value = steps
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.content = "Answer"
|
|
mock_response.tool_calls = []
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.bind_tools.return_value = mock_llm
|
|
|
|
async def mock_astream(*args, **kwargs):
|
|
yield mock_response
|
|
|
|
mock_llm.astream.side_effect = mock_astream
|
|
|
|
fake_plan = """- [/] Active Subgoal
|
|
- [ ] Pending Subgoal"""
|
|
|
|
mock_path = MagicMock()
|
|
mock_path.exists.return_value = True
|
|
mock_path.read_text.return_value = fake_plan
|
|
|
|
# Mock the policy-table compiled-history renderer (M4)
|
|
with (
|
|
patch(
|
|
"artemis.agents.history_analyzer.history_analyzer.get_llm",
|
|
return_value=mock_llm,
|
|
),
|
|
patch(
|
|
"artemis.agents.history_analyzer.history_analyzer.get_note_file_path",
|
|
return_value=mock_path,
|
|
),
|
|
patch("artemis.agents.history_analyzer.history_analyzer.build_history_for") as mock_build,
|
|
):
|
|
mock_build.return_value = "Mocked operation history"
|
|
analyzer = HistoryAnalyzer(mock_ctx)
|
|
await analyzer.run("query")
|
|
|
|
# Verify build_history_for call args
|
|
mock_build.assert_called_once()
|
|
args, kwargs = mock_build.call_args
|
|
assert args[0] == "history_analyzer"
|
|
assert args[1] == fake_plan
|
|
assert args[2] == steps
|
|
|
|
import hashlib
|
|
|
|
expected_hash = hashlib.md5(b"Active Subgoal").hexdigest()
|
|
assert args[3] == expected_hash
|
|
|
|
assert kwargs.get("engine") is mock_ctx.data_engine
|