From 358f4623bab34bc3e2ca9c9db4cc57cd2b28efa5 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Thu, 16 Jul 2026 10:25:20 +0800 Subject: [PATCH] feat(planner): add execution context to prompts --- .../device-host-agent/host_agent/execution.py | 15 +++++++ .../device-host-agent/tests/test_execution.py | 18 ++++++++ runtime/ai_planner.py | 1 + runtime/context.py | 1 + runtime/planner_prompts.py | 24 +++++++++++ runtime/task.py | 25 ++++++++++- tests/test_ai_planner.py | 16 +++++++ tests/test_ai_planner_task_runner.py | 20 ++++++++- tests/test_planner_prompts.py | 42 +++++++++++++++++++ 9 files changed, 159 insertions(+), 3 deletions(-) create mode 100644 tests/test_planner_prompts.py diff --git a/apps/device-host-agent/host_agent/execution.py b/apps/device-host-agent/host_agent/execution.py index 736e947..6c58ef3 100644 --- a/apps/device-host-agent/host_agent/execution.py +++ b/apps/device-host-agent/host_agent/execution.py @@ -50,6 +50,9 @@ def create_execution_factories( timeline=timeline, planner=_host_agent_planner(resolved_host_agent_config), planner_config=_host_agent_planner_config(), + device_platform_provider=lambda device_id: _device_platform( + manager, device_id + ), ) def create_workflow_runner() -> WorkflowRunner: @@ -103,3 +106,15 @@ def _host_agent_planner( client=CloudProxyToolCallingClient(resolved_config), config=planner_config, ) + + +def _device_platform(manager: DeviceManager, device_id: str) -> str | None: + for device in manager.list_devices(): + if device.id != device_id: + continue + if device.driver_type == "wda": + return "ios" + if device.driver_type == "uiautomator2": + return "android" + return None + return None diff --git a/apps/device-host-agent/tests/test_execution.py b/apps/device-host-agent/tests/test_execution.py index a3bb323..8807b88 100644 --- a/apps/device-host-agent/tests/test_execution.py +++ b/apps/device-host-agent/tests/test_execution.py @@ -141,6 +141,24 @@ def test_created_task_runner_observer_uses_configured_manager( assert seen == {"device_id": "phone-1", "manager": manager} +@pytest.mark.parametrize( + ("driver_type", "expected_platform"), + [("wda", "ios"), ("uiautomator2", "android")], +) +def test_created_task_runner_resolves_platform_from_configured_driver( + tmp_path, driver_type: str, expected_platform: str +) -> None: + manager = DeviceManager() + manager.register_device("phone-1", lambda: FakeDriver(), driver_type=driver_type) + factories = create_execution_factories( + manager, workflow_store=WorkflowStore(tmp_path / "workflows.sqlite3") + ) + task_runner = factories.task_runner_factory() + + assert task_runner.device_platform_provider is not None + assert task_runner.device_platform_provider("phone-1") == expected_platform + + def test_created_task_runner_defaults_to_ai_planner( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/runtime/ai_planner.py b/runtime/ai_planner.py index 8efc224..856d16f 100644 --- a/runtime/ai_planner.py +++ b/runtime/ai_planner.py @@ -40,6 +40,7 @@ class AIPlanner(Planner): goal=goal, scene_json=scene.to_dict(), history_summary=_history_summary(world), + device_platform=context.device_platform, ) decision = self.client.decide( system_prompt=PLANNER_SYSTEM_PROMPT, diff --git a/runtime/context.py b/runtime/context.py index b5eefdc..31f3651 100644 --- a/runtime/context.py +++ b/runtime/context.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: class TaskContext: task_id: str goal: str + device_platform: str | None = None scenes: list[Scene] = field(default_factory=list) step_results: list["StepResult"] = field(default_factory=list) world: "WorldState | None" = None diff --git a/runtime/planner_prompts.py b/runtime/planner_prompts.py index be8b0b6..549c435 100644 --- a/runtime/planner_prompts.py +++ b/runtime/planner_prompts.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from datetime import datetime from typing import Any PLANNER_SYSTEM_PROMPT = """You are the planning brain of a mobile device automation agent. @@ -66,8 +67,18 @@ def planner_user_prompt( goal: str, scene_json: dict[str, Any], history_summary: list[dict[str, Any]], + device_platform: str | None = None, + now: datetime | None = None, ) -> str: + current_time = now or datetime.now().astimezone() + if current_time.tzinfo is None: + current_time = current_time.astimezone() + timezone_name = current_time.tzname() or str(current_time.tzinfo) or "unknown" return ( + "Execution context:\n" + f"Current date and time: {current_time.isoformat(timespec='seconds')}\n" + f"Time zone: {timezone_name}\n" + f"Device type: {_device_type(scene_json, device_platform)}\n\n" "Goal:\n" f"{goal}\n\n" "Current Scene (JSON):\n" @@ -76,3 +87,16 @@ def planner_user_prompt( f"{json.dumps(history_summary, ensure_ascii=False, sort_keys=True)}\n\n" "Call exactly one tool for this turn." ) + + +def _device_type(scene_json: dict[str, Any], device_platform: str | None) -> str: + platform = device_platform + if platform is None: + app = scene_json.get("app") + if isinstance(app, dict): + raw_platform = app.get("platform") + platform = raw_platform if isinstance(raw_platform, str) else None + if not isinstance(platform, str): + return "unknown" + normalized = platform.strip().lower() + return normalized if normalized in {"ios", "android"} else "unknown" diff --git a/runtime/task.py b/runtime/task.py index 242ae48..398ddee 100644 --- a/runtime/task.py +++ b/runtime/task.py @@ -42,6 +42,7 @@ TaskSucceededHook = Callable[[str, str, Timeline], None] StopRequested = Callable[[], bool] StopReason = Callable[[], "str | None"] StepProgressCallback = Callable[[int, str, str], None] +DevicePlatformProvider = Callable[[str], str | None] def is_cancellation_reason(reason: str | None) -> bool: @@ -68,6 +69,7 @@ class TaskRunner: skill_embedding_client: EmbeddingClient | None = None, planner_config: PlannerConfig | None = None, on_step_progress: StepProgressCallback | None = None, + device_platform_provider: DevicePlatformProvider | None = None, ) -> None: self.planner_config = planner_config or load_planner_config() self.planner = planner or self._default_planner() @@ -98,6 +100,7 @@ class TaskRunner: else: self.on_task_succeeded = None self.on_step_progress = on_step_progress + self.device_platform_provider = device_platform_provider def run( self, @@ -108,7 +111,11 @@ class TaskRunner: ) -> Task: if self.metadata_store: self.metadata_store.create_task(task) - context = TaskContext(task_id=task.id, goal=task.goal) + context = TaskContext( + task_id=task.id, + goal=task.goal, + device_platform=self._device_platform(task.device_id), + ) world_handle = self._start_world_view(task.id) if world_handle is not None: context.world = world_handle.state @@ -201,7 +208,9 @@ class TaskRunner: ) return task - def _interrupt_task(self, task: Task, stop_reason: StopReason | None = None) -> Task: + def _interrupt_task( + self, task: Task, stop_reason: StopReason | None = None + ) -> Task: reason = stop_reason() if stop_reason is not None else None message = reason or "execution interrupted" status = "cancelled" if is_cancellation_reason(reason) else "failed" @@ -307,6 +316,18 @@ class TaskRunner: return AIPlanner(config=self.planner_config) return Planner() + def _device_platform(self, device_id: str) -> str | None: + if self.device_platform_provider is None: + return None + try: + return self.device_platform_provider(device_id) + except Exception: + logger.warning( + "device platform lookup failed; continuing without configured platform", + exc_info=True, + ) + return None + def _plan( self, goal: str, diff --git a/tests/test_ai_planner.py b/tests/test_ai_planner.py index 0850a60..ea52290 100644 --- a/tests/test_ai_planner.py +++ b/tests/test_ai_planner.py @@ -146,6 +146,22 @@ def test_ai_planner_forwards_tools_screenshot_and_timeout_to_client() -> None: assert "send a message" in call["user_prompt"] +def test_ai_planner_includes_device_platform_from_task_context() -> None: + client = FakeToolCallingClient( + ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2}) + ) + planner = AIPlanner(client=client) + context = TaskContext( + task_id="task-1", + goal="send a message", + device_platform="android", + ) + + planner.plan(goal="send a message", scene=_scene(), context=context) + + assert "Device type: android" in client.calls[0]["user_prompt"] + + def test_ai_planner_populates_step_prompt_from_user_prompt() -> None: """PlannedStep.prompt should carry the actual user prompt sent to the LLM, not the bare task goal.""" diff --git a/tests/test_ai_planner_task_runner.py b/tests/test_ai_planner_task_runner.py index 8323acd..f9884db 100644 --- a/tests/test_ai_planner_task_runner.py +++ b/tests/test_ai_planner_task_runner.py @@ -46,9 +46,11 @@ class NarrowSignaturePlanner(Planner): class ScreenshotRecordingPlanner(Planner): def __init__(self) -> None: self.screenshots: list[bytes | None] = [] + self.device_platforms: list[str | None] = [] def plan(self, *, goal, scene, context, screenshot=None): self.screenshots.append(screenshot) + self.device_platforms.append(context.device_platform) if context.step_results: return [] return [PlannedStep(action="tap", description="tap")] @@ -69,7 +71,9 @@ def _scene() -> Scene: ) -def _runner(*, planner=None, planner_config=None, observer=None) -> TaskRunner: +def _runner( + *, planner=None, planner_config=None, observer=None, device_platform_provider=None +) -> TaskRunner: return TaskRunner( planner=planner, planner_config=planner_config, @@ -80,6 +84,7 @@ def _runner(*, planner=None, planner_config=None, observer=None) -> TaskRunner: config=TaskRunnerConfig(max_steps=5), observer=observer or (lambda device_id: _scene()), screenshot_provider=lambda device_id: PNG_10X20, + device_platform_provider=device_platform_provider, ) @@ -126,6 +131,19 @@ def test_task_runner_passes_screenshot_to_planner_that_declares_it() -> None: assert planner.screenshots == [PNG_10X20, PNG_10X20] +def test_task_runner_passes_configured_device_platform_to_planner_context() -> None: + planner = ScreenshotRecordingPlanner() + runner = _runner( + planner=planner, + device_platform_provider=lambda device_id: "ios", + ) + + result = runner.run(Task(goal="inspect", device_id="phone")) + + assert result.status == "completed" + assert planner.device_platforms == ["ios", "ios"] + + def test_task_runner_default_planner_is_stub_when_ai_planner_disabled() -> None: runner = _runner(planner=None, planner_config=PlannerConfig(enabled=False)) diff --git a/tests/test_planner_prompts.py b/tests/test_planner_prompts.py new file mode 100644 index 0000000..5311064 --- /dev/null +++ b/tests/test_planner_prompts.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +from runtime.planner_prompts import planner_user_prompt + + +def test_planner_user_prompt_includes_time_zone_and_configured_device_type() -> None: + prompt = planner_user_prompt( + goal="open settings", + scene_json={"screen": {"width": 1, "height": 1}, "elements": []}, + history_summary=[], + device_platform="ios", + now=datetime( + 2026, + 7, + 16, + 9, + 8, + 7, + tzinfo=timezone(timedelta(hours=8), "Asia/Shanghai"), + ), + ) + + assert "Current date and time: 2026-07-16T09:08:07+08:00" in prompt + assert "Time zone: Asia/Shanghai" in prompt + assert "Device type: ios" in prompt + + +def test_planner_user_prompt_uses_scene_platform_when_context_is_unavailable() -> None: + prompt = planner_user_prompt( + goal="open settings", + scene_json={ + "screen": {"width": 1, "height": 1}, + "elements": [], + "app": {"platform": "android"}, + }, + history_summary=[], + now=datetime(2026, 7, 16, tzinfo=timezone.utc), + ) + + assert "Device type: android" in prompt