feat: checkpoint device agent runtime milestones
This commit is contained in:
@@ -0,0 +1,275 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
from core.models import Bounds, Scene, SceneElement, Task
|
||||
from runtime.context import TaskContext
|
||||
from runtime.executor import Executor, ExecutorConfig
|
||||
from runtime.planner import PlannedStep, Planner
|
||||
from runtime.task import TaskRunner, TaskRunnerConfig
|
||||
from semantic.models import SemanticScene, SemanticWidget
|
||||
from world.config import WorldConfig
|
||||
from world.model import WorldModel
|
||||
|
||||
|
||||
class RecordingPlanner(Planner):
|
||||
def __init__(self, steps: list[PlannedStep]) -> None:
|
||||
self.steps = steps
|
||||
self.worlds: list[Any] = []
|
||||
self.world_snapshots: list[dict[str, Any] | None] = []
|
||||
|
||||
def plan(self, *, goal, scene, context, world=None):
|
||||
self.worlds.append(world)
|
||||
self.world_snapshots.append(world.to_dict() if world is not None else None)
|
||||
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(
|
||||
result.success for result in context.step_results
|
||||
)
|
||||
|
||||
|
||||
class KwargRecordingPlanner(Planner):
|
||||
def __init__(self) -> None:
|
||||
self.received_world_kwarg: list[bool] = []
|
||||
|
||||
def plan(self, *, goal, scene, context, **kwargs):
|
||||
self.received_world_kwarg.append("world" in kwargs)
|
||||
if context.step_results:
|
||||
return []
|
||||
return [PlannedStep(action="tap", description="tap")]
|
||||
|
||||
def goal_reached(self, *, goal, scene, context):
|
||||
return bool(context.step_results)
|
||||
|
||||
|
||||
class NoWorldPlanner(Planner):
|
||||
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)
|
||||
|
||||
|
||||
def _scene(text: str = "Send") -> Scene:
|
||||
return Scene(
|
||||
width=10,
|
||||
height=20,
|
||||
elements=[
|
||||
SceneElement(
|
||||
id="send",
|
||||
type="button",
|
||||
text=text,
|
||||
bounds=Bounds(1, 2, 3, 4),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _semantic_scene(page: str) -> SemanticScene:
|
||||
return SemanticScene(
|
||||
page=page,
|
||||
intents=["send a message"],
|
||||
widgets=[SemanticWidget(element_id="send", purpose="send message")],
|
||||
)
|
||||
|
||||
|
||||
def _runner(
|
||||
*,
|
||||
planner: Planner,
|
||||
executor: Executor,
|
||||
scenes: list[Scene] | None = None,
|
||||
world_model: WorldModel | None = None,
|
||||
world_config: WorldConfig | None = None,
|
||||
) -> TaskRunner:
|
||||
scene_iter: Iterator[Scene] = iter(scenes or [_scene(), _scene(), _scene()])
|
||||
|
||||
def observer(device_id: str) -> Scene:
|
||||
return next(scene_iter, _scene())
|
||||
|
||||
return TaskRunner(
|
||||
planner=planner,
|
||||
executor=executor,
|
||||
config=TaskRunnerConfig(max_steps=5),
|
||||
observer=observer,
|
||||
world_model=world_model,
|
||||
world_config=world_config,
|
||||
)
|
||||
|
||||
|
||||
def test_planner_plan_accepts_calls_with_and_without_world() -> None:
|
||||
scene = _scene()
|
||||
planner = Planner()
|
||||
context = TaskContext(
|
||||
task_id="task",
|
||||
goal="inspect",
|
||||
)
|
||||
world_model = WorldModel(config=WorldConfig(history_size=2))
|
||||
|
||||
without_world = planner.plan(goal="inspect", scene=scene, context=context)
|
||||
with_world = planner.plan(
|
||||
goal="inspect",
|
||||
scene=scene,
|
||||
context=context,
|
||||
world=world_model.state,
|
||||
)
|
||||
|
||||
assert with_world == without_world
|
||||
|
||||
|
||||
def test_task_runner_populates_world_after_step_when_enabled_by_default() -> None:
|
||||
planner = RecordingPlanner([PlannedStep(action="tap", description="tap")])
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={"tap": lambda **kwargs: {"ok": True}},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
)
|
||||
|
||||
result = runner.run(Task(goal="inspect", device_id="phone"))
|
||||
|
||||
assert result.status == "completed"
|
||||
assert planner.worlds[0] is not None
|
||||
assert planner.world_snapshots[1] is not None
|
||||
assert len(planner.world_snapshots[1]["history"]) == 1
|
||||
|
||||
|
||||
def test_task_runner_keeps_existing_planner_subclasses_without_world_working() -> None:
|
||||
planner = NoWorldPlanner()
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={"tap": lambda **kwargs: {"ok": True}},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
)
|
||||
|
||||
result = runner.run(Task(goal="inspect", device_id="phone"))
|
||||
|
||||
assert result.status == "completed"
|
||||
assert planner.calls == 2
|
||||
|
||||
|
||||
def test_task_runner_disabled_world_config_does_not_pass_world_kwarg() -> None:
|
||||
planner = KwargRecordingPlanner()
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={"tap": lambda **kwargs: {"ok": True}},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
world_config=WorldConfig(enabled=False),
|
||||
)
|
||||
|
||||
result = runner.run(Task(goal="inspect", device_id="phone"))
|
||||
|
||||
assert result.status == "completed"
|
||||
assert planner.received_world_kwarg == [False, False]
|
||||
|
||||
|
||||
def test_task_runner_passes_current_world_to_planner_by_second_step() -> None:
|
||||
planner = RecordingPlanner(
|
||||
[
|
||||
PlannedStep(action="tap", description="first"),
|
||||
PlannedStep(action="tap", description="second"),
|
||||
]
|
||||
)
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={"tap": lambda **kwargs: {"ok": True}},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
)
|
||||
|
||||
runner.run(Task(goal="inspect", device_id="phone"))
|
||||
|
||||
assert planner.worlds[1] is not None
|
||||
assert planner.world_snapshots[1] is not None
|
||||
assert planner.world_snapshots[2] is not None
|
||||
assert len(planner.world_snapshots[1]["history"]) == 1
|
||||
assert len(planner.world_snapshots[2]["history"]) == 2
|
||||
|
||||
|
||||
def test_task_runner_world_tracks_app_page_and_history_across_steps() -> None:
|
||||
world_model = WorldModel(config=WorldConfig(history_size=3))
|
||||
planner = RecordingPlanner(
|
||||
[
|
||||
PlannedStep(
|
||||
action="launch_app",
|
||||
description="launch chat",
|
||||
args={"app_id": "com.example.chat"},
|
||||
),
|
||||
PlannedStep(action="tap", description="open chat"),
|
||||
]
|
||||
)
|
||||
|
||||
def launch_app(**kwargs):
|
||||
return {"ok": True, "semantic_scene": _semantic_scene("Home")}
|
||||
|
||||
def tap(**kwargs):
|
||||
return {"ok": True, "semantic_scene": _semantic_scene("Chat")}
|
||||
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={"launch_app": launch_app, "tap": tap},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
world_model=world_model,
|
||||
)
|
||||
|
||||
result = runner.run(Task(goal="open chat", device_id="phone"))
|
||||
|
||||
assert result.status == "completed"
|
||||
assert world_model.state.current_app == "com.example.chat"
|
||||
assert world_model.state.current_page == "Chat"
|
||||
assert [event.action for event in world_model.state.history] == [
|
||||
"launch_app",
|
||||
"tap",
|
||||
]
|
||||
assert planner.world_snapshots[1]["current_page"] == "Home"
|
||||
assert planner.world_snapshots[2]["current_page"] == "Chat"
|
||||
|
||||
|
||||
def test_task_runner_world_tracking_failures_do_not_interrupt_loop() -> None:
|
||||
world_model = WorldModel(config=WorldConfig(history_size=3))
|
||||
planner = RecordingPlanner(
|
||||
[
|
||||
PlannedStep(action="launch_app", description="missing app args"),
|
||||
PlannedStep(
|
||||
action="tap",
|
||||
description="bad remember",
|
||||
args={"remember": "not-a-dict"},
|
||||
),
|
||||
]
|
||||
)
|
||||
runner = _runner(
|
||||
planner=planner,
|
||||
executor=Executor(
|
||||
tools={
|
||||
"launch_app": lambda **kwargs: {"ok": True},
|
||||
"tap": lambda **kwargs: {"ok": True},
|
||||
},
|
||||
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
|
||||
),
|
||||
world_model=world_model,
|
||||
)
|
||||
|
||||
result = runner.run(Task(goal="inspect", device_id="phone"))
|
||||
|
||||
assert result.status == "completed"
|
||||
assert world_model.state.current_app is None
|
||||
assert world_model.state.current_page is None
|
||||
assert world_model.state.variables == {}
|
||||
assert len(world_model.state.history) == 2
|
||||
Reference in New Issue
Block a user