335 lines
12 KiB
Python
335 lines
12 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.
|
||
|
||
"""search_history boundaries and search correctness.
|
||
|
||
Hard boundaries under test: bounded result count, bounded response tokens,
|
||
every result carrying a step number. Plus: keyword/step-range/notes/chunk
|
||
filter correctness, the step-range ledger re-entry (full-width
|
||
``build_action_ledger`` rows), and the screen-text surface: the step's own
|
||
pre AND post screenshots plus screenshots its tool results embedded.
|
||
"""
|
||
|
||
import hashlib
|
||
import io
|
||
import re
|
||
from types import SimpleNamespace
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
from PIL import Image
|
||
|
||
from artemis.context import ArtemisContext
|
||
from artemis.data_engine.engine import DataEngine
|
||
from artemis.tools.history import (
|
||
SearchHistoryTool,
|
||
search_history,
|
||
search_history_available,
|
||
search_history_text,
|
||
)
|
||
from artemis.tools.history.search import result_search_text
|
||
from artemis.utils.notes import save_note_content
|
||
|
||
|
||
def _jpeg(color: str) -> bytes:
|
||
buf = io.BytesIO()
|
||
Image.new("RGB", (64, 128), color).save(buf, format="JPEG")
|
||
return buf.getvalue()
|
||
|
||
|
||
def _make_engine(tmp_path):
|
||
mock_ctx = MagicMock(spec=ArtemisContext)
|
||
mock_execution_setup = MagicMock()
|
||
mock_execution_setup.traces_path = str(tmp_path)
|
||
mock_ctx.execution_setup = mock_execution_setup
|
||
mock_ctx.device = None
|
||
engine = DataEngine(mock_ctx)
|
||
engine.start_session("search test session")
|
||
return engine
|
||
|
||
|
||
def _flush(engine):
|
||
for t in list(engine._pending_threads):
|
||
t.join()
|
||
|
||
|
||
def _cfg(**overrides):
|
||
values = {"enabled": True, "max_results": 5, "max_text_tokens": 2000, "screen_scan_steps": 150}
|
||
values.update(overrides)
|
||
return SimpleNamespace(**values)
|
||
|
||
|
||
@pytest.fixture
|
||
def engine(tmp_path):
|
||
engine = _make_engine(tmp_path)
|
||
engine.record_step(
|
||
pre_screenshot_bytes=_jpeg("red"),
|
||
summary="Opened the login page.",
|
||
action_taken={"action": "click", "target_text": "Login entry"},
|
||
last_execution_result={"status": "success"},
|
||
)
|
||
engine.record_step(
|
||
pre_screenshot_bytes=_jpeg("green"),
|
||
post_screenshot_bytes=_jpeg("blue"),
|
||
summary="Typed the verification code 4711 into the input.",
|
||
action_taken={"action": "input_text", "text": "4711"},
|
||
last_execution_result={"status": "success"},
|
||
operator_raw_thinking="The code arrived via SMS; entering 4711 now.",
|
||
)
|
||
engine.record_step(
|
||
summary="A promo popup appeared and was dismissed.",
|
||
action_taken={"action": "click", "target_text": "Close popup"},
|
||
last_execution_result={"status": "failed", "error": "element vanished"},
|
||
)
|
||
_flush(engine)
|
||
return engine
|
||
|
||
|
||
def _result_blocks(text: str) -> list[str]:
|
||
return [b for b in text.split("\n") if b.startswith("[")]
|
||
|
||
|
||
def test_keyword_search_finds_step_and_carries_step_number(engine):
|
||
out = search_history_text(engine, query="verification code", recall_config=_cfg())
|
||
assert isinstance(out, str)
|
||
assert "Step 2" in out
|
||
assert "4711" in out
|
||
# Every result header names a step / step range / step anchor.
|
||
for header in _result_blocks(out):
|
||
assert re.search(r"Step[s]? \d", header), header
|
||
|
||
|
||
def test_step_range_filters_matches(engine):
|
||
out = search_history_text(engine, query="popup", step_range=[1, 2], recall_config=_cfg())
|
||
# Step 3 (the popup step) is outside the range: only the range ledger and
|
||
# a no-match line may appear.
|
||
assert "A promo popup appeared" not in out
|
||
assert "No matches" in out
|
||
|
||
|
||
def test_step_range_returns_full_width_ledger_rows(engine):
|
||
"""The compressed-history marker line's re-entry point: a step range
|
||
returns that range's full per-step action ledger."""
|
||
out = search_history_text(engine, query="", step_range=[3, 1], recall_config=_cfg())
|
||
assert "Action ledger for Steps 1–3" in out
|
||
# Full-width build_action_ledger rows: step number + T+mm:ss offset +
|
||
# semantic action + result phrase.
|
||
assert re.search(r"- Step 1 \(T\+\d{2}:\d{2}\): .*Login entry", out)
|
||
assert re.search(r"- Step 2 \(T\+\d{2}:\d{2}\): ", out)
|
||
assert re.search(r"- Step 3 \(T\+\d{2}:\d{2}\): ", out)
|
||
|
||
|
||
def test_result_count_is_clamped_by_config(engine):
|
||
for i in range(8):
|
||
engine.record_step(
|
||
summary=f"Scrolled the anchorterm feed page {i}.",
|
||
action_taken={"action": "swipe"},
|
||
last_execution_result={"status": "success"},
|
||
)
|
||
_flush(engine)
|
||
out = search_history_text(
|
||
engine, query="anchorterm", max_results=50, recall_config=_cfg(max_results=3)
|
||
)
|
||
assert out.count("[Step ") == 3
|
||
assert "more not shown" in out
|
||
|
||
|
||
def test_response_is_truncated_at_token_budget(engine):
|
||
for i in range(6):
|
||
engine.record_step(
|
||
summary=f"needleword page {i} " + "filler content " * 60,
|
||
action_taken={"action": "click"},
|
||
last_execution_result={"status": "success"},
|
||
)
|
||
_flush(engine)
|
||
out = search_history_text(
|
||
engine, query="needleword", max_results=6, recall_config=_cfg(max_text_tokens=150)
|
||
)
|
||
assert len(out) <= 150 * 4 + 200 # budget + truncation notice
|
||
assert "truncated at the search token budget" in out
|
||
|
||
|
||
def test_notes_are_searchable_and_anchored_to_a_step(engine):
|
||
save_note_content(engine.base_dir, "login_flow", "The OTP entry lives behind the SMS tab.")
|
||
out = search_history_text(engine, query="OTP entry", recall_config=_cfg())
|
||
assert "[Note 'login_flow'" in out
|
||
assert re.search(r"as of Step \d|last written at Step \d", out)
|
||
|
||
|
||
def test_chunk_rows_are_searchable(engine):
|
||
engine.record_history_chunk(
|
||
start_step_number=1,
|
||
end_step_number=2,
|
||
version=1,
|
||
status="ready",
|
||
band2=" - Steps 1–2: walked the chunkneedle flow",
|
||
band3="- Step 1 (T+00:01): click -> executed",
|
||
rendered_text=(
|
||
"[Chunk 1 | Steps 1–2]\n② Compressed step summary\n"
|
||
" - Steps 1–2: walked the chunkneedle flow\n"
|
||
"③ Step action ledger\n- Step 1 (T+00:01): click -> executed"
|
||
),
|
||
)
|
||
_flush(engine)
|
||
out = search_history_text(engine, query="chunkneedle", recall_config=_cfg())
|
||
assert "[History chunk | Steps 1–2" in out
|
||
|
||
|
||
def test_foreground_app_stamp_joins_search_surface(tmp_path):
|
||
"""A package-name query hits the record_step foreground_app stamp."""
|
||
engine = _make_engine(tmp_path)
|
||
engine.record_step(
|
||
summary="Opened the alarms list.",
|
||
action_taken={"action": "click", "target_text": "Alarm"},
|
||
ui_tree=[{"package": "com.google.android.deskclock", "bounds": "[0,0][9,9]"}],
|
||
)
|
||
_flush(engine)
|
||
|
||
out = search_history_text(engine, query="deskclock", recall_config=_cfg())
|
||
assert "Step 1" in out
|
||
assert "No matches" not in out
|
||
|
||
|
||
def test_tool_calls_of_a_step_are_searchable(engine):
|
||
engine.record_step(
|
||
summary="Asked the explorer.",
|
||
action_taken={"action": "click", "target_text": "Save"},
|
||
last_execution_result={"status": "success"},
|
||
)
|
||
engine.record_trace(
|
||
type="tool",
|
||
name="ask_explorer",
|
||
payload={"args": {"question": "toolneedle?"}, "result": "The resultneedle is ON."},
|
||
step_id=engine.last_recorded_step_id,
|
||
)
|
||
_flush(engine)
|
||
assert "[Step 4 " in search_history_text(engine, query="toolneedle", recall_config=_cfg())
|
||
assert "[Step 4 " in search_history_text(engine, query="resultneedle", recall_config=_cfg())
|
||
|
||
|
||
# --- Screen text: pre + post screenshots and referenced screenshots ----------------------
|
||
|
||
|
||
def test_post_screenshot_text_joins_the_search_surface(engine):
|
||
"""OCR text that only exists on the step's post-action screenshot is found."""
|
||
record = engine.get_step_record(2)
|
||
engine.storage.update_image_data(
|
||
record.post_image_name,
|
||
ocr_result=[{"text": "Welcome back postneedle", "bounds": [0, 0, 10, 10]}],
|
||
ui_tree=None,
|
||
)
|
||
out = search_history_text(engine, query="postneedle", recall_config=_cfg())
|
||
assert "[Step 2 " in out
|
||
|
||
|
||
def test_screen_scan_cap_limits_the_screenshot_sweep(engine):
|
||
record = engine.get_step_record(1)
|
||
engine.storage.update_image_data(
|
||
record.pre_image_name,
|
||
ocr_result=[{"text": "capneedle", "bounds": [0, 0, 10, 10]}],
|
||
ui_tree=None,
|
||
)
|
||
# Only the most recent step is scanned: Step 1's OCR is out of reach ...
|
||
out = search_history_text(engine, query="capneedle", recall_config=_cfg(screen_scan_steps=1))
|
||
assert "No matches" in out
|
||
# ... unless the range narrows the sweep onto it.
|
||
out = search_history_text(
|
||
engine, query="capneedle", step_range=[1, 1], recall_config=_cfg(screen_scan_steps=1)
|
||
)
|
||
assert "[Step 1 " in out
|
||
|
||
|
||
def test_embedded_screenshot_description_joins_the_search_surface(tmp_path):
|
||
"""A step whose tool result embedded an earlier screenshot is found by
|
||
that screenshot's OCR text: the image is described, not stripped."""
|
||
engine = _make_engine(tmp_path)
|
||
pre_bytes = _jpeg("red")
|
||
engine.record_step(
|
||
pre_screenshot_bytes=pre_bytes,
|
||
summary="Opened the alarms list.",
|
||
action_taken={"action": "click", "target_text": "Alarm"},
|
||
ocr_result=[{"text": "Alarm 07:30 ocrneedle", "bounds": [0, 0, 10, 10]}],
|
||
)
|
||
engine.record_step(
|
||
summary="Looked back at the alarm screen.",
|
||
action_taken={"action": "click", "target_text": "Save"},
|
||
)
|
||
engine.record_trace(
|
||
type="tool",
|
||
name="get_step_screenshot",
|
||
payload={
|
||
"args": {"step_number": 1, "which": "pre"},
|
||
"result": [
|
||
{"type": "text", "text": "Screenshot of step 1 (pre-action) is attached."},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"<ImageRef: sha256={hashlib.sha256(pre_bytes).hexdigest()} length=1>"
|
||
},
|
||
},
|
||
],
|
||
},
|
||
step_id=engine.last_recorded_step_id,
|
||
)
|
||
_flush(engine)
|
||
|
||
out = search_history_text(engine, query="ocrneedle", recall_config=_cfg())
|
||
assert "[Step 1 " in out
|
||
assert "[Step 2 " in out
|
||
assert "base64" not in out
|
||
|
||
|
||
def test_result_search_text_strips_image_blocks():
|
||
blocks = [
|
||
{"type": "text", "text": "hello needle"},
|
||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,QUJDRA=="}},
|
||
]
|
||
text = result_search_text(blocks)
|
||
assert "hello needle" in text
|
||
assert "base64" not in text and "QUJDRA" not in text
|
||
assert result_search_text({"answer": "yes"}) == '{"answer": "yes"}'
|
||
assert result_search_text(None) == ""
|
||
|
||
|
||
# --- Tool surface ---------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_execute_requires_query_or_range(engine):
|
||
tool = SearchHistoryTool()
|
||
out = await tool.execute(ctx=SimpleNamespace(data_engine=engine), query="", step_range=None)
|
||
assert "needs a query and/or a step_range" in out
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_execute_searches_and_never_returns_images(engine):
|
||
out = await search_history.execute(
|
||
ctx=SimpleNamespace(data_engine=engine), query="verification code"
|
||
)
|
||
assert isinstance(out, str)
|
||
assert "Step 2" in out
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_execute_without_engine_degrades():
|
||
out = await search_history.execute(ctx=SimpleNamespace(data_engine=None), query="x")
|
||
assert "no active execution history" in out
|
||
|
||
|
||
def test_availability_requires_engine_and_config():
|
||
assert search_history.is_available(None) is False
|
||
assert search_history.is_available(SimpleNamespace(data_engine=None)) is False
|
||
assert search_history.is_available(SimpleNamespace(data_engine=object())) is True
|
||
assert search_history_available(SimpleNamespace(data_engine=object())) is True
|