somew1nd
12d7902a82
Fix queue cancellation, concurrency limits and worker status updates. Use direct Explorer imports, split the Diagnoser loop into helpers and remove unused tool registration and prompt-builder code. Correct model caching, endpoint selection and retry handling.
467 lines
18 KiB
Python
467 lines
18 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.
|
|
|
|
"""Explorer in-memory screen search, label registry, bounds enrichment and pre-pass."""
|
|
|
|
import json
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from artemis.agents.explorer.explorer import Explorer
|
|
from artemis.agents.explorer.screen_index import ScreenElement, ScreenIndex
|
|
from artemis.context import ArtemisContext
|
|
from artemis.graph.state import State
|
|
|
|
RUN_SETUP = "artemis.agents.explorer.run_setup"
|
|
PERCEPTION_TOOLS = "artemis.agents.explorer.perception_tools"
|
|
W, H = 1080, 2400
|
|
|
|
#: One "Settings" button, two "Send" buttons and two OCR-only texts.
|
|
HIERARCHY = [
|
|
{
|
|
"text": "Settings",
|
|
"bounds": "[500,400][600,560]",
|
|
"class": "android.widget.Button",
|
|
"resource-id": "app:id/settings",
|
|
"clickable": "true",
|
|
},
|
|
{"text": "Send", "bounds": "[100,2000][300,2100]", "class": "android.widget.Button"},
|
|
{"text": "Send", "bounds": "[700,2000][900,2100]", "class": "android.widget.Button"},
|
|
{
|
|
"bounds": "[0,0][1080,2400]",
|
|
"class": "android.widget.FrameLayout",
|
|
"ocr_elements": [
|
|
{"text": "Settings page", "bounds": "[100,100][300,160]"},
|
|
{"text": "Welcome", "bounds": "[100,200][400,260]"},
|
|
],
|
|
},
|
|
]
|
|
|
|
#: "Settings" [500,400][600,560] on a 1080x2400 screen, in normalized units.
|
|
SETTINGS_COORDS = [509, 200]
|
|
SETTINGS_BOUNDS = [462, 166, 555, 233]
|
|
|
|
|
|
def _context() -> MagicMock:
|
|
ctx = MagicMock(spec=ArtemisContext)
|
|
ctx.device = MagicMock()
|
|
ctx.device.device_width = W
|
|
ctx.device.device_height = H
|
|
ctx.data_engine = MagicMock()
|
|
ctx.data_engine.base_dir = None
|
|
ctx.llm_config = MagicMock()
|
|
ctx.llm_config.explorer = MagicMock()
|
|
ctx.llm_config.explorer.model = "gemini-3.8-flash"
|
|
ctx.agent_config = MagicMock()
|
|
ctx.agent_config.denylisted_tools = {}
|
|
return ctx
|
|
|
|
|
|
def _state(screenshot: str, hierarchy=None) -> MagicMock:
|
|
state = MagicMock(spec=State)
|
|
state.latest_screenshot = screenshot
|
|
state.latest_ui_hierarchy = HIERARCHY if hierarchy is None else hierarchy
|
|
state.operator_raw_data = {}
|
|
return state
|
|
|
|
|
|
def _indexed_explorer(hierarchy=None) -> Explorer:
|
|
"""An Explorer whose screen is already indexed (no ``run`` involved)."""
|
|
explorer = Explorer(_context())
|
|
explorer.width, explorer.height = W, H
|
|
explorer.screenshot_path = "/tmp/screen.jpg"
|
|
explorer.screen_index = ScreenIndex.from_hierarchy(
|
|
HIERARCHY if hierarchy is None else hierarchy, W, H
|
|
)
|
|
return explorer
|
|
|
|
|
|
def _storage(record=None) -> MagicMock:
|
|
storage = MagicMock()
|
|
storage.get_image.return_value = record
|
|
return storage
|
|
|
|
|
|
def _outcome(candidates=None, fallback: str = "") -> str:
|
|
return json.dumps({"candidates": candidates or [], "fallback_message": fallback})
|
|
|
|
|
|
@pytest.fixture
|
|
def screenshot(tmp_path):
|
|
path = tmp_path / "screen.jpg"
|
|
path.write_bytes(b"fake-jpeg-bytes")
|
|
return str(path)
|
|
|
|
|
|
@pytest.fixture
|
|
def no_drawing():
|
|
"""Keeps the annotation helpers away from the file system."""
|
|
draw_dots = MagicMock()
|
|
with (
|
|
patch(f"{RUN_SETUP}.draw_dots", draw_dots),
|
|
patch(f"{PERCEPTION_TOOLS}.draw_dots", draw_dots),
|
|
patch("pathlib.Path.mkdir"),
|
|
patch("glob.glob", return_value=[]),
|
|
):
|
|
yield draw_dots
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Screen index construction inside run()
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_is_built_from_state_hierarchy_when_record_is_missing(screenshot, no_drawing):
|
|
explorer = Explorer(_context())
|
|
explorer._run_flash = AsyncMock(return_value=_outcome(fallback="nothing"))
|
|
|
|
with (
|
|
patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)),
|
|
patch(f"{RUN_SETUP}.perform_ocr", new_callable=AsyncMock) as perform_ocr,
|
|
):
|
|
await explorer.run("gear icon", "", screenshot, _state(screenshot), version="flash")
|
|
text = (await explorer._search_ui_helper("Settings", prefix="X"))["text"]
|
|
|
|
perform_ocr.assert_not_awaited()
|
|
assert len(explorer.screen_index) == 5
|
|
assert "not found in Data Engine" not in text
|
|
assert "[X1] 'Settings' at [509,200]" in text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_ignores_state_hierarchy_for_another_screenshot(screenshot):
|
|
explorer = Explorer(_context())
|
|
explorer._run_flash = AsyncMock(return_value=_outcome())
|
|
state = _state("/tmp/other.jpg")
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)):
|
|
await explorer.run("Settings", "", screenshot, state, version="flash")
|
|
|
|
assert len(explorer.screen_index) == 0
|
|
explorer._run_flash.assert_awaited_once_with("Settings", screenshot)
|
|
assert (await explorer._search_ui_helper("Settings"))["text"] == (
|
|
"No UI-tree data is available for this screen; rely on visual detection."
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Perception helpers answer from the index
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_ui_helper_merges_lenient_matches_without_duplicates(no_drawing):
|
|
hierarchy = [*HIERARCHY, {"text": "Settled", "bounds": "[100,300][300,360]"}]
|
|
explorer = _indexed_explorer(hierarchy)
|
|
spy = MagicMock(wraps=explorer.screen_index.search_text)
|
|
explorer.screen_index.search_text = spy
|
|
|
|
result = await explorer._search_ui_helper("Settings", prefix="X", color="green")
|
|
|
|
assert [call.args[1] for call in spy.call_args_list] == [0.7, 0.4]
|
|
parts = result["text"].split(" | ")
|
|
assert parts[0] == "[X1] 'Settings' at [509,200]"
|
|
assert parts[1] == "[O2] 'Settings page' at [185,54]"
|
|
# "Settled" (similarity ~0.53) only arrives through the lenient pass...
|
|
assert any(part.startswith("[X") and "'Settled'" in part for part in parts)
|
|
# ...and the strict hits are not listed a second time.
|
|
assert sum("'Settings'" in part for part in parts) == 1
|
|
assert explorer.label_registry["X1"].text == "Settings"
|
|
assert explorer.label_registry["O2"].source == "ocr"
|
|
assert no_drawing.call_args.kwargs["color"] == "green"
|
|
assert result["image_path"] is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_ui_helper_skips_lenient_pass_with_enough_strict_hits(no_drawing):
|
|
explorer = _indexed_explorer()
|
|
strict = explorer.screen_index.search_text("Settings", 0.4) # 5 hits, all "strict"
|
|
explorer.screen_index.search_text = MagicMock(return_value=strict)
|
|
|
|
await explorer._search_ui_helper("Settings", prefix="X")
|
|
|
|
explorer.screen_index.search_text.assert_called_once_with("Settings", 0.7)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_by_coords_helper_reports_innermost_first(no_drawing):
|
|
explorer = _indexed_explorer(
|
|
[
|
|
{"text": "Card", "bounds": "[0,0][1080,1200]", "class": "android.widget.FrameLayout"},
|
|
{"text": "Title", "bounds": "[500,400][600,560]", "class": "android.widget.TextView"},
|
|
]
|
|
)
|
|
|
|
hit = await explorer._search_by_coords_helper(550, 200, prefix="X", color="blue")
|
|
miss = await explorer._search_by_coords_helper(10, 990, prefix="X", color="blue")
|
|
|
|
assert hit["text"] == (
|
|
"Matched element at [550,200]: [X1]"
|
|
" | 'Title' class=android.widget.TextView source=ui-tree"
|
|
" | 'Card' class=android.widget.FrameLayout source=ui-tree"
|
|
)
|
|
assert explorer.label_registry["X1"].text == "Title"
|
|
assert miss["text"] == "No elements found at [10,990]."
|
|
assert "X2" not in explorer.label_registry
|
|
# The blue dot is drawn for both probes at the pixel position.
|
|
assert no_drawing.call_count == 2
|
|
assert no_drawing.call_args_list[0].args[1] == [[594, 480]]
|
|
assert no_drawing.call_args_list[0].args[2] == ["X1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_ocr_list_reads_the_index(no_drawing):
|
|
explorer = _indexed_explorer()
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager") as storage_cls:
|
|
result = await explorer.exec_get_ocr_list()
|
|
|
|
storage_cls.assert_not_called()
|
|
assert result["text"].splitlines() == [
|
|
"[O1] 'Settings page' coords: [185,54]",
|
|
"[O2] 'Welcome' coords: [231,95]",
|
|
]
|
|
assert result["image_path"] is not None
|
|
assert {k: v.source for k, v in explorer.label_registry.items()} == {"O1": "ocr", "O2": "ocr"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_ocr_list_without_ocr_elements():
|
|
explorer = _indexed_explorer([{"text": "Only XML", "bounds": "[0,0][100,100]"}])
|
|
result = await explorer.exec_get_ocr_list()
|
|
assert result == {"text": "No text elements detected on the screen.", "image_path": None}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Label registry from the initial numbered list
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_initial_list_registers_numbered_labels(no_drawing):
|
|
explorer = _indexed_explorer()
|
|
|
|
minimal_list, _marked = explorer._annotate_initial_screenshot(
|
|
HIERARCHY, "/tmp/screen.jpg", None, ""
|
|
)
|
|
|
|
assert minimal_list.splitlines()[0].startswith("[1] Text: 'Settings'")
|
|
assert sorted(explorer.label_registry) == ["1", "2", "3", "4", "5"]
|
|
settings = explorer.label_registry["1"]
|
|
assert settings.bounds == (500, 400, 600, 560)
|
|
assert settings.source == "xml" and settings.interactive # taken from the index
|
|
assert explorer.label_registry["4"].source == "ocr"
|
|
assert explorer.label_registry["4"].text == "Settings page"
|
|
assert explorer.global_label_idx == 6
|
|
assert no_drawing.call_args.args[1][0] == [550, 480]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Submit enrichment
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_enrich_candidates_adds_normalized_bounds_only_when_label_and_coords_agree():
|
|
explorer = _indexed_explorer()
|
|
settings = explorer.screen_index.exact_matches("Settings")[0]
|
|
explorer._register_label("X1", settings)
|
|
|
|
inside = {"label": "X1", "coords": SETTINGS_COORDS, "description": "gear"}
|
|
# 614px is 14px right of the element, inside the 2% (21.6px) tolerance.
|
|
near = {"label": "X1", "coords": [569, 200], "description": "gear"}
|
|
far = {"label": "X1", "coords": [100, 900], "description": "gear"}
|
|
unknown = {"label": "D7", "coords": SETTINGS_COORDS, "description": "gear"}
|
|
invalid = {"label": "X1", "coords": [5000, 5], "description": "gear"}
|
|
|
|
enriched = explorer._enrich_candidates([inside, near, far, unknown, invalid, "junk"])
|
|
|
|
assert enriched[0] == {**inside, "bounds": SETTINGS_BOUNDS, "source": "xml"}
|
|
assert enriched[1] == {**near, "bounds": SETTINGS_BOUNDS, "source": "xml"}
|
|
assert enriched[2:] == [far, unknown, invalid, "junk"]
|
|
assert "bounds" not in far and "bounds" not in inside # inputs are never mutated
|
|
|
|
|
|
def test_enrich_candidates_withholds_degenerate_bounds():
|
|
explorer = _indexed_explorer()
|
|
explorer._register_label("X9", ScreenElement("dot", (10, 10, 11, 11), "xml"))
|
|
cand = {"label": "X9", "coords": [9, 4]}
|
|
assert explorer._enrich_candidates([cand]) == [cand]
|
|
|
|
|
|
def test_native_submit_enriches_registered_candidates():
|
|
explorer = _indexed_explorer()
|
|
explorer._register_label("O2", explorer.screen_index.ocr_elements[0])
|
|
call = MagicMock()
|
|
call.name = "submit_answer"
|
|
call.args = {
|
|
"candidates": [
|
|
{"label": "O2", "coords": [185, 54], "description": "header"},
|
|
{"label": "D3", "coords": [900, 100], "description": "icon"},
|
|
],
|
|
"fallback_message": "",
|
|
}
|
|
turn_record = {"tool_calls": []}
|
|
|
|
raw = explorer._process_submit_answer([call], turn_record, [])
|
|
|
|
data = json.loads(raw)
|
|
assert list(data) == ["candidates", "fallback_message"]
|
|
assert data["candidates"][0] == {
|
|
"label": "O2",
|
|
"coords": [185, 54],
|
|
"description": "header",
|
|
"bounds": [92, 41, 277, 66],
|
|
"source": "ocr",
|
|
}
|
|
assert data["candidates"][1] == call.args["candidates"][1]
|
|
|
|
|
|
def test_universal_submit_enriches_registered_candidates():
|
|
explorer = _indexed_explorer()
|
|
explorer._register_label("1", explorer.screen_index.exact_matches("Settings")[0])
|
|
tool_calls = [
|
|
{
|
|
"name": "submit_answer",
|
|
"args": {
|
|
"candidates": [
|
|
{"label": "1", "coords": SETTINGS_COORDS, "description": "gear"},
|
|
{"label": "1", "coords": [5, 5000], "description": "bad"},
|
|
],
|
|
"fallback_message": "note",
|
|
},
|
|
"id": "c1",
|
|
}
|
|
]
|
|
|
|
data = json.loads(explorer._universal_submit_outcome(tool_calls))
|
|
|
|
assert data["fallback_message"] == "note"
|
|
assert data["candidates"] == [
|
|
{
|
|
"label": "1",
|
|
"coords": SETTINGS_COORDS,
|
|
"description": "gear",
|
|
"bounds": SETTINGS_BOUNDS,
|
|
"source": "xml",
|
|
}
|
|
]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Structural pre-pass
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.parametrize("version", ["flash", "pro"])
|
|
@pytest.mark.asyncio
|
|
async def test_prepass_answers_unique_exact_match_without_the_engine(screenshot, version):
|
|
explorer = Explorer(_context())
|
|
explorer._run_flash = AsyncMock()
|
|
explorer._run_loop = AsyncMock()
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)):
|
|
raw = await explorer.run(" settings ", "", screenshot, _state(screenshot), version=version)
|
|
|
|
explorer._run_flash.assert_not_called()
|
|
explorer._run_loop.assert_not_called()
|
|
assert json.loads(raw) == {
|
|
"candidates": [
|
|
{
|
|
"label": "T1",
|
|
"coords": SETTINGS_COORDS,
|
|
"bounds": SETTINGS_BOUNDS,
|
|
"source": "xml",
|
|
"description": "Settings (exact text match)",
|
|
}
|
|
],
|
|
"fallback_message": "",
|
|
}
|
|
assert explorer.label_registry["T1"].text == "Settings"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepass_leaves_ambiguous_matches_to_the_engine(screenshot):
|
|
explorer = Explorer(_context())
|
|
engine = _outcome([{"label": "D1", "coords": [185, 854], "description": "Send"}])
|
|
explorer._run_flash = AsyncMock(return_value=engine)
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)):
|
|
raw = await explorer.run("Send", "", screenshot, _state(screenshot), version="flash")
|
|
|
|
explorer._run_flash.assert_awaited_once_with("Send", screenshot)
|
|
assert raw == engine
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepass_without_pipe_or_match_passes_the_query_through(screenshot):
|
|
explorer = Explorer(_context())
|
|
engine = _outcome(fallback="Failed to detect: gear icon")
|
|
explorer._run_flash = AsyncMock(return_value=engine)
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)):
|
|
raw = await explorer.run("gear icon", "", screenshot, _state(screenshot), version="flash")
|
|
|
|
explorer._run_flash.assert_awaited_once_with("gear icon", screenshot)
|
|
assert raw == engine
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepass_merges_partial_results_with_the_engine_outcome(screenshot):
|
|
explorer = Explorer(_context())
|
|
detected = {"label": "D1", "coords": [185, 854], "description": "Send"}
|
|
explorer._run_loop = AsyncMock(return_value=_outcome([detected], fallback="two Send buttons"))
|
|
|
|
with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)):
|
|
raw = await explorer.run(
|
|
"Settings | Send | Welcome", "", screenshot, _state(screenshot), version="pro"
|
|
)
|
|
|
|
# Only the unresolved part reaches the engine; "Welcome" is an OCR exact match.
|
|
assert explorer._run_loop.await_args.args[1] == "Send"
|
|
data = json.loads(raw)
|
|
assert [c["label"] for c in data["candidates"]] == ["T1", "T2", "D1"]
|
|
assert data["candidates"][1]["source"] == "ocr"
|
|
assert data["candidates"][1]["description"] == "Welcome (exact text match)"
|
|
assert data["candidates"][2] == detected
|
|
assert data["fallback_message"] == "two Send buttons"
|
|
|
|
|
|
def test_merge_prepass_keeps_structural_hits_over_an_unreadable_engine_answer():
|
|
prepass = [{"label": "T1", "coords": [1, 2]}]
|
|
merged = json.loads(Explorer._merge_prepass(prepass, "Error: max iterations"))
|
|
assert merged == {"candidates": prepass, "fallback_message": "Error: max iterations"}
|
|
assert Explorer._merge_prepass([], "raw") == "raw"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepass_never_runs_on_the_fly_ocr_for_flash(screenshot):
|
|
record = SimpleNamespace(ui_tree=HIERARCHY[:3], ocr_result=None)
|
|
explorer = Explorer(_context())
|
|
explorer._run_flash = AsyncMock(return_value=_outcome())
|
|
explorer._run_loop = AsyncMock(return_value=_outcome())
|
|
|
|
with (
|
|
patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(record)),
|
|
patch(f"{RUN_SETUP}.is_ocr_configured", return_value=True),
|
|
patch(f"{RUN_SETUP}.perform_ocr", new_callable=AsyncMock, return_value=[]) as perform_ocr,
|
|
):
|
|
await explorer.run("gear icon", "", screenshot, _state(screenshot, []), version="flash")
|
|
perform_ocr.assert_not_awaited()
|
|
assert len(explorer.screen_index) == 3 # the record still feeds the index
|
|
|
|
await explorer.run("gear icon", "", screenshot, _state(screenshot, []), version="pro")
|
|
perform_ocr.assert_awaited_once()
|