290 lines
9.1 KiB
Python
290 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from core.models import Bounds, Scene, SceneElement, Task
|
|
from runtime.ai_planner import AIPlanner
|
|
from runtime.executor import Executor, ExecutorConfig
|
|
from runtime.planner import PlannedStep, Planner
|
|
from runtime.planner_config import PlannerConfig
|
|
from runtime.task import TaskRunner, TaskRunnerConfig
|
|
from runtime.tool_calling_client import ToolCallDecision
|
|
from storage.artifact_store import ArtifactStore
|
|
from storage.timeline import Timeline
|
|
from tests.fakes import PNG_10X20
|
|
|
|
|
|
class RaisingPlanner(Planner):
|
|
def __init__(self, error: Exception) -> None:
|
|
self.error = error
|
|
self.calls = 0
|
|
|
|
def plan(self, *, goal, scene, context):
|
|
self.calls += 1
|
|
raise self.error
|
|
|
|
def goal_reached(self, *, goal, scene, context):
|
|
return False
|
|
|
|
|
|
class NarrowSignaturePlanner(Planner):
|
|
"""Predates the `screenshot` parameter added to the base Planner.plan()."""
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def plan(self, *, goal, scene, context):
|
|
self.calls += 1
|
|
if context.step_results:
|
|
return []
|
|
return [PlannedStep(action="tap", description="tap")]
|
|
|
|
def goal_reached(self, *, goal, scene, context):
|
|
return bool(context.step_results)
|
|
|
|
|
|
class ScreenshotRecordingPlanner(Planner):
|
|
def __init__(self) -> None:
|
|
self.screenshots: list[bytes | None] = []
|
|
|
|
def plan(self, *, goal, scene, context, screenshot=None):
|
|
self.screenshots.append(screenshot)
|
|
if context.step_results:
|
|
return []
|
|
return [PlannedStep(action="tap", description="tap")]
|
|
|
|
def goal_reached(self, *, goal, scene, context):
|
|
return bool(context.step_results)
|
|
|
|
|
|
def _scene() -> Scene:
|
|
return Scene(
|
|
width=10,
|
|
height=20,
|
|
elements=[
|
|
SceneElement(
|
|
id="send", type="button", text="Send", bounds=Bounds(1, 2, 3, 4)
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
def _runner(*, planner=None, planner_config=None, observer=None) -> TaskRunner:
|
|
return TaskRunner(
|
|
planner=planner,
|
|
planner_config=planner_config,
|
|
executor=Executor(
|
|
tools={"tap": lambda **kwargs: {"ok": True}},
|
|
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
|
),
|
|
config=TaskRunnerConfig(max_steps=5),
|
|
observer=observer or (lambda device_id: _scene()),
|
|
screenshot_provider=lambda device_id: PNG_10X20,
|
|
)
|
|
|
|
|
|
def test_task_runner_marks_task_failed_when_planner_raises() -> None:
|
|
planner = RaisingPlanner(RuntimeError("boom"))
|
|
runner = _runner(planner=planner)
|
|
|
|
result = runner.run(Task(goal="inspect", device_id="phone"))
|
|
|
|
assert result.status == "failed"
|
|
assert result.failure_reason == "RuntimeError: boom"
|
|
assert planner.calls == 1
|
|
|
|
|
|
def test_task_runner_marks_task_failed_when_observer_raises() -> None:
|
|
def failing_observer(device_id: str) -> Scene:
|
|
raise RuntimeError("no device")
|
|
|
|
runner = _runner(planner=Planner(), observer=failing_observer)
|
|
|
|
result = runner.run(Task(goal="inspect", device_id="phone"))
|
|
|
|
assert result.status == "failed"
|
|
assert result.failure_reason == "RuntimeError: no device"
|
|
|
|
|
|
def test_task_runner_omits_screenshot_kwarg_for_narrow_signature_planner() -> None:
|
|
planner = NarrowSignaturePlanner()
|
|
runner = _runner(planner=planner)
|
|
|
|
result = runner.run(Task(goal="inspect", device_id="phone"))
|
|
|
|
assert result.status == "completed"
|
|
assert planner.calls == 2
|
|
|
|
|
|
def test_task_runner_passes_screenshot_to_planner_that_declares_it() -> None:
|
|
planner = ScreenshotRecordingPlanner()
|
|
runner = _runner(planner=planner)
|
|
|
|
result = runner.run(Task(goal="inspect", device_id="phone"))
|
|
|
|
assert result.status == "completed"
|
|
assert planner.screenshots == [PNG_10X20, PNG_10X20]
|
|
|
|
|
|
def test_task_runner_default_planner_is_stub_when_ai_planner_disabled() -> None:
|
|
runner = _runner(planner=None, planner_config=PlannerConfig(enabled=False))
|
|
|
|
assert type(runner.planner) is Planner
|
|
|
|
|
|
def test_task_runner_default_planner_is_ai_planner_when_enabled() -> None:
|
|
runner = _runner(
|
|
planner=None,
|
|
planner_config=PlannerConfig(
|
|
enabled=True, provider="anthropic", model="test-model"
|
|
),
|
|
)
|
|
|
|
assert isinstance(runner.planner, AIPlanner)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# D9: per-step prompt recording
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class ScriptedToolCallingClient:
|
|
"""Returns a sequence of decisions, capturing the actual user_prompt each call."""
|
|
|
|
def __init__(self, decisions: list[ToolCallDecision]) -> None:
|
|
self._decisions = list(decisions)
|
|
self.calls: list[dict[str, Any]] = []
|
|
|
|
def decide(
|
|
self,
|
|
*,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
screenshot: bytes | None,
|
|
tools: list[Any],
|
|
timeout: float,
|
|
) -> ToolCallDecision:
|
|
index = len(self.calls)
|
|
self.calls.append(
|
|
{
|
|
"system_prompt": system_prompt,
|
|
"user_prompt": user_prompt,
|
|
}
|
|
)
|
|
return self._decisions[index]
|
|
|
|
|
|
def _multi_step_scene(element_text: str = "Send") -> Scene:
|
|
return Scene(
|
|
width=10,
|
|
height=20,
|
|
elements=[
|
|
SceneElement(
|
|
id="btn",
|
|
type="button",
|
|
text=element_text,
|
|
bounds=Bounds(1, 2, 3, 4),
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
def test_multi_step_timeline_records_actual_per_step_prompts(tmp_path) -> None:
|
|
"""When AIPlanner is used, each timeline step's prompt is the real
|
|
per-step user prompt (containing scene JSON), not the bare task goal."""
|
|
decisions = [
|
|
ToolCallDecision(
|
|
tool_name="tap",
|
|
arguments={"x": 1, "y": 2},
|
|
purpose="Open the send control.",
|
|
expected_outcome="The message composer is focused.",
|
|
),
|
|
ToolCallDecision(tool_name="finish_task", arguments={"success": True}),
|
|
]
|
|
client = ScriptedToolCallingClient(decisions)
|
|
planner = AIPlanner(client=client)
|
|
executor = Executor(
|
|
tools={"tap": lambda **kwargs: {"ok": True, **kwargs}},
|
|
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
|
)
|
|
timeline = Timeline(ArtifactStore(tmp_path / "history"))
|
|
task = Task(goal="tap the button", device_id="phone")
|
|
|
|
runner = TaskRunner(
|
|
planner=planner,
|
|
executor=executor,
|
|
timeline=timeline,
|
|
config=TaskRunnerConfig(max_steps=5),
|
|
observer=lambda device_id: _multi_step_scene("Send"),
|
|
screenshot_provider=lambda device_id: PNG_10X20,
|
|
)
|
|
|
|
runner.run(task)
|
|
|
|
records = timeline.read(task.id)
|
|
# Step 1 was recorded (step 2 was finish_task, which returns empty plan
|
|
# and completes the task without a timeline append).
|
|
assert len(records) == 1
|
|
prompt = records[0]["prompt"]
|
|
# The per-step prompt is NOT the bare task goal.
|
|
assert prompt != "tap the button"
|
|
# It contains scene-specific content that only the real planner_user_prompt
|
|
# would include.
|
|
assert "Current Scene (JSON)" in prompt
|
|
assert "Call exactly one tool" in prompt
|
|
assert "tap the button" in prompt
|
|
assert records[0]["tool_call"]["args"] == {"x": 1, "y": 2}
|
|
assert records[0]["tool_call"]["purpose"] == "Open the send control."
|
|
assert (
|
|
records[0]["tool_call"]["expected_outcome"]
|
|
== "The message composer is focused."
|
|
)
|
|
|
|
|
|
def test_non_ai_planner_falls_back_to_task_goal_for_prompt(tmp_path) -> None:
|
|
"""A non-LLM planner (no step.prompt) keeps recording task.goal as the
|
|
timeline prompt — backward compat with pre-D9 behavior."""
|
|
scene = _multi_step_scene("Search")
|
|
planner = ScriptedPlannerForTimeline(
|
|
[PlannedStep(action="tap", description="tap", args={"x": 1, "y": 2})]
|
|
)
|
|
executor = Executor(
|
|
tools={"tap": lambda **kwargs: {"ok": True, **kwargs}},
|
|
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
|
)
|
|
timeline = Timeline(ArtifactStore(tmp_path / "history"))
|
|
task = Task(goal="search for something", device_id="phone")
|
|
|
|
runner = TaskRunner(
|
|
planner=planner,
|
|
executor=executor,
|
|
timeline=timeline,
|
|
config=TaskRunnerConfig(max_steps=5),
|
|
observer=lambda device_id: scene,
|
|
screenshot_provider=lambda device_id: PNG_10X20,
|
|
)
|
|
|
|
runner.run(task)
|
|
|
|
records = timeline.read(task.id)
|
|
assert len(records) >= 1
|
|
# Non-AI planner: prompt falls back to task goal.
|
|
assert records[0]["prompt"] == "search for something"
|
|
|
|
|
|
class ScriptedPlannerForTimeline(Planner):
|
|
"""Simple planner that returns a fixed list of steps then signals done."""
|
|
|
|
def __init__(self, steps: list[PlannedStep]) -> None:
|
|
self.steps = steps
|
|
|
|
def plan(self, *, goal, scene, context):
|
|
if len(context.step_results) >= len(self.steps):
|
|
return []
|
|
return [self.steps[len(context.step_results)]]
|
|
|
|
def goal_reached(self, *, goal, scene, context):
|
|
return len(context.step_results) >= len(self.steps) and all(
|
|
r.success for r in context.step_results
|
|
)
|