Tests / Test apps.device-host-agent.tests.test_mcp_token.test_load_or_create_concurrent_calls_do_not_corrupt failed
117 lines
4.1 KiB
Python
117 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from core.errors import TaskFailedError
|
|
from core.models import Scene
|
|
from runtime.context import TaskContext
|
|
from runtime.planner import PlannedStep, Planner
|
|
from runtime.planner_config import PlannerConfig, load_config
|
|
from runtime.planner_prompts import PLANNER_SYSTEM_PROMPT, planner_user_prompt
|
|
from runtime.tool_calling_client import ToolCallingClient, build_client
|
|
from runtime.tool_specs import ALL_TOOL_SPECS
|
|
|
|
if TYPE_CHECKING:
|
|
from world.models import WorldState
|
|
|
|
FINISH_TASK_TOOL = "finish_task"
|
|
|
|
|
|
class AIPlanner(Planner):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
client: ToolCallingClient | None = None,
|
|
config: PlannerConfig | None = None,
|
|
) -> None:
|
|
self.config = config or load_config()
|
|
self.client = client or build_client(self.config)
|
|
|
|
def plan(
|
|
self,
|
|
*,
|
|
goal: str,
|
|
scene: Scene,
|
|
context: TaskContext,
|
|
world: "WorldState | None" = None,
|
|
screenshot: bytes | None = None,
|
|
) -> list[PlannedStep]:
|
|
scene_json = _without_ocr(scene.to_dict()) if self.config.multimodal else scene.to_dict()
|
|
user_prompt = planner_user_prompt(
|
|
goal=goal,
|
|
scene_json=scene_json,
|
|
history_summary=_history_summary(world),
|
|
device_platform=context.device_platform,
|
|
)
|
|
if context.step_results and not context.step_results[-1].success:
|
|
user_prompt += (
|
|
"\n\nThe previous tool call failed. Diagnose the failure and choose a "
|
|
"corrected action or finish the task if it cannot proceed.\n"
|
|
f"Previous failure: {context.step_results[-1].error or 'unknown error'}"
|
|
)
|
|
decision = self.client.decide(
|
|
system_prompt=PLANNER_SYSTEM_PROMPT,
|
|
user_prompt=user_prompt,
|
|
screenshot=screenshot,
|
|
tools=ALL_TOOL_SPECS,
|
|
timeout=self.config.timeout,
|
|
)
|
|
|
|
if decision.tool_name == FINISH_TASK_TOOL:
|
|
if decision.arguments.get("success"):
|
|
return []
|
|
raise TaskFailedError(decision.arguments.get("reason") or "task failed")
|
|
|
|
return [
|
|
PlannedStep(
|
|
action=decision.tool_name,
|
|
description=(
|
|
f"AI planner: {decision.purpose}"
|
|
if decision.purpose
|
|
else f"AI planner: {decision.tool_name}({decision.arguments})"
|
|
),
|
|
args=dict(decision.arguments),
|
|
expected_text=decision.expected_outcome,
|
|
purpose=decision.purpose,
|
|
expected_outcome=decision.expected_outcome,
|
|
prompt=decision.user_prompt or user_prompt,
|
|
rationale=decision.text_output,
|
|
thinking=decision.thinking,
|
|
)
|
|
]
|
|
|
|
def goal_reached(self, *, goal: str, scene: Scene, context: TaskContext) -> bool:
|
|
# Completion is signaled exclusively via the finish_task tool call
|
|
# (mapped to an empty plan above), never via this hook.
|
|
return False
|
|
|
|
|
|
def _history_summary(world: "WorldState | None") -> list[dict[str, Any]]:
|
|
if world is None:
|
|
return []
|
|
return [
|
|
{
|
|
"page": event.page,
|
|
"action": event.action,
|
|
"arguments": dict(event.arguments),
|
|
"rationale": event.rationale,
|
|
"purpose": event.purpose,
|
|
"expected_outcome": event.expected_outcome,
|
|
"success": event.success,
|
|
}
|
|
for event in world.history
|
|
]
|
|
|
|
|
|
def _without_ocr(scene_json: dict[str, Any]) -> dict[str, Any]:
|
|
cleaned = dict(scene_json)
|
|
elements = cleaned.get("elements")
|
|
if isinstance(elements, list):
|
|
cleaned["elements"] = [
|
|
{key: value for key, value in element.items() if key not in {"source", "confidence", "foreground_color", "background_color"}}
|
|
for element in elements
|
|
if isinstance(element, dict) and element.get("source") != "ocr"
|
|
]
|
|
cleaned.pop("ocr_elements", None)
|
|
return cleaned
|