185 lines
7.1 KiB
Python
185 lines
7.1 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.
|
||
|
||
"""Screen observation: capture, fuse, index, and persist -- without LangGraph State.
|
||
|
||
This is the act-then-observe half that used to be welded into every executor action
|
||
(``capture_screenshot_and_parse_ui``). Split out so that actuation and observation are
|
||
separate calls: the Flash adapter observes after each action, while the
|
||
Validator -- which runs its own UI-change polling -- simply never calls it.
|
||
"""
|
||
|
||
import asyncio
|
||
import base64
|
||
from io import BytesIO
|
||
import time
|
||
from PIL import Image
|
||
|
||
from artemis.config import get_temp_dir
|
||
from artemis.config.settings import settings
|
||
from artemis.context import ArtemisContext
|
||
from artemis.controllers.unified_controller import UnifiedMobileController
|
||
from artemis.mcp.action_types import ActionCode, ObserveResult
|
||
from artemis.utils.logger import get_logger
|
||
from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml
|
||
from artemis.utils.visualization import format_minimal_list_with_elements
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
__all__ = ["observe", "require_observation", "DeviceObservationError"]
|
||
|
||
|
||
class DeviceObservationError(RuntimeError):
|
||
"""A live device could not be observed after bounded recovery."""
|
||
|
||
|
||
async def require_observation(
|
||
ctx: ArtemisContext, controller: UnifiedMobileController, settle_ms: int = 400
|
||
) -> tuple[ObserveResult, bytes | None]:
|
||
"""Recover observation without spending model calls on unavailable screens."""
|
||
last = "设备未返回有效画面"
|
||
for attempt in range(settings.ARTEMIS_OBSERVATION_MAX_ATTEMPTS):
|
||
observation, image = await observe(ctx, controller, settle_ms=settle_ms)
|
||
if observation.ok:
|
||
return observation, image
|
||
last = observation.message
|
||
logger.warning(
|
||
f"设备观测恢复 {attempt + 1}/{settings.ARTEMIS_OBSERVATION_MAX_ATTEMPTS}：{last}"
|
||
)
|
||
raise DeviceObservationError(
|
||
f"设备观测失败（已尝试 {settings.ARTEMIS_OBSERVATION_MAX_ATTEMPTS} 次）：{last}"
|
||
)
|
||
|
||
|
||
async def observe(
|
||
ctx: ArtemisContext,
|
||
controller: UnifiedMobileController,
|
||
settle_ms: int = 400,
|
||
ocr_results: list | None = None,
|
||
) -> tuple[ObserveResult, bytes | None]:
|
||
"""Captures and indexes the current screen.
|
||
|
||
Returns the structured observation plus the raw screenshot bytes (an in-process
|
||
convenience; the serialized ``ObserveResult`` carries only the persisted path).
|
||
|
||
Mirrors the historical ``capture_screenshot_and_parse_ui`` behavior exactly:
|
||
a settling delay, screen capture, XML fusion and element indexing, persistence via
|
||
the data engine with a temp-file fallback -- and the same failure semantics: a
|
||
failed capture yields ``ok=False``; a failed hierarchy parse still returns the
|
||
screenshot with ``hierarchy_ok=False``.
|
||
"""
|
||
try:
|
||
if settle_ms > 0:
|
||
logger.info(f"Delaying {settle_ms / 1000:.1f}s before capturing screen state...")
|
||
await asyncio.sleep(settle_ms / 1000.0)
|
||
else:
|
||
logger.info("Retrieving screen data directly (skip settling)...")
|
||
device_data = await asyncio.wait_for(
|
||
controller.get_screen_data(), timeout=settings.ARTEMIS_OBSERVATION_TIMEOUT_SECONDS
|
||
)
|
||
if getattr(device_data, "platform", "android") == "computer":
|
||
screenshot = base64.b64decode(device_data.base64) if device_data.base64 else None
|
||
path = None
|
||
if screenshot and ctx.data_engine:
|
||
path = str(
|
||
ctx.data_engine.get_image_path(ctx.data_engine.get_or_create_image(screenshot))
|
||
)
|
||
return ObserveResult(
|
||
ok=True,
|
||
message="已读取电脑插件观察结果",
|
||
screenshot_path=path,
|
||
width=device_data.width,
|
||
height=device_data.height,
|
||
elements_text=device_data.observation_text,
|
||
hierarchy_ok=True,
|
||
), screenshot
|
||
screenshot_bytes = base64.b64decode(device_data.base64)
|
||
with Image.open(BytesIO(screenshot_bytes)) as image:
|
||
if min(image.size) <= 1:
|
||
raise ValueError("设备截图为占位图，当前画面不可用")
|
||
image.verify()
|
||
except Exception as e:
|
||
if isinstance(e, TimeoutError):
|
||
e = TimeoutError(
|
||
f"设备屏幕观测超时（{settings.ARTEMIS_OBSERVATION_TIMEOUT_SECONDS:g} 秒）"
|
||
)
|
||
logger.error(f"Failed to capture screen state: {e}")
|
||
return (
|
||
ObserveResult(
|
||
ok=False,
|
||
code=ActionCode.DEVICE_ERROR,
|
||
message=f"Failed to capture screen state: {e}",
|
||
hierarchy_ok=False,
|
||
),
|
||
None,
|
||
)
|
||
|
||
xml_hierarchy = None
|
||
ocr_results = ocr_results or []
|
||
elements: list[dict] = []
|
||
hierarchy_ok = True
|
||
width = height = None
|
||
try:
|
||
xml_hierarchy = device_data.elements
|
||
width = device_data.width or 1080
|
||
height = device_data.height or 2400
|
||
if ctx.device:
|
||
ctx.device.device_width = width
|
||
ctx.device.device_height = height
|
||
|
||
fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results)
|
||
minimal_list_str, elements, _ = format_minimal_list_with_elements(fused_xml, width, height)
|
||
except Exception as e:
|
||
logger.error(f"Failed to fetch or fuse UI hierarchy: {e}")
|
||
minimal_list_str = "Not available due to hierarchy error."
|
||
hierarchy_ok = False
|
||
|
||
screenshot_path = None
|
||
if ctx.data_engine:
|
||
try:
|
||
image_name = ctx.data_engine.get_or_create_image(
|
||
screenshot_bytes,
|
||
ui_tree=xml_hierarchy,
|
||
ocr_result=ocr_results,
|
||
)
|
||
screenshot_path = str(ctx.data_engine.get_image_path(image_name))
|
||
except Exception as e:
|
||
logger.error(f"Failed to save image in DataEngine: {e}")
|
||
|
||
if not screenshot_path:
|
||
try:
|
||
temp_dir = get_temp_dir("screenshots")
|
||
temp_file = temp_dir / f"screenshot_{int(time.time())}.jpg"
|
||
temp_file.write_bytes(screenshot_bytes)
|
||
screenshot_path = str(temp_file)
|
||
except Exception as e:
|
||
logger.error(f"Failed to save fallback screenshot: {e}")
|
||
|
||
from pathlib import Path
|
||
|
||
return (
|
||
ObserveResult(
|
||
ok=True,
|
||
message="Screen observed.",
|
||
screenshot_path=screenshot_path,
|
||
image_name=Path(screenshot_path).stem if screenshot_path else None,
|
||
width=width,
|
||
height=height,
|
||
elements_text=minimal_list_str,
|
||
elements=elements,
|
||
hierarchy_ok=hierarchy_ok,
|
||
),
|
||
screenshot_bytes,
|
||
)
|