249 lines
11 KiB
Python
249 lines
11 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.
|
|
|
|
"""Asynchronous Visual Step Summarizer for Artemis Flash profile.
|
|
|
|
Processes Before/After screenshot pairs and action metadata in the background
|
|
to generate high-density, strictly objective visual state transition summaries.
|
|
Ensures zero blocking latency for the main runner and independently retries failed jobs.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from jinja2 import Template
|
|
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
|
|
|
|
from artemis.context import ArtemisContext
|
|
from artemis.services.llm import get_google_llm, get_llm
|
|
from artemis.utils.logger import get_logger
|
|
from artemis.utils.visualization import draw_action_overlay_on_image
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class VisualStepSummarizer:
|
|
"""Non-blocking background worker generating objective visual step summaries.
|
|
|
|
Key design properties:
|
|
1. Zero-blocking dispatch: The main Flash runner dispatches and proceeds immediately.
|
|
2. Lossless pending state: The compressor retains the source image until its summary is ready.
|
|
3. Independent retry: Every action owns a retry loop and does not depend on later actions arriving.
|
|
"""
|
|
|
|
def __init__(self, ctx: ArtemisContext, model_name: str | None = None):
|
|
self.ctx = ctx
|
|
self._summaries: dict[int | str, str] = {}
|
|
self._pending_tasks: dict[int | str, asyncio.Task] = {}
|
|
self._step_inputs: dict[int | str, dict[str, Any]] = {}
|
|
self._retry_counts: dict[int | str, int] = {}
|
|
self._retry_delays = (0.0, 0.5, 1.0, 2.0, 3.0)
|
|
|
|
# Initialize lightweight VLM: prioritize explicit model_name
|
|
target_model = model_name or "gemini-2.5-flash-lite"
|
|
try:
|
|
if model_name:
|
|
self._llm = get_google_llm(model_name=target_model, temperature=0.0)
|
|
else:
|
|
self._llm = get_llm(ctx, name="summarizer", is_utils=True)
|
|
except Exception:
|
|
self._llm = get_google_llm(model_name=target_model, temperature=0.0)
|
|
|
|
# Load system prompt template
|
|
prompt_path = Path(__file__).parent / "flash_summarizer.md"
|
|
if prompt_path.exists():
|
|
self._prompt_template = prompt_path.read_text(encoding="utf-8")
|
|
else:
|
|
self._prompt_template = (
|
|
"You are the Step Summarizer for an Android UI automation agent.\n"
|
|
"Synthesize the physical action and visual delta between BEFORE and AFTER screens in exactly ONE "
|
|
"continuous first-person paragraph using 'I' (e.g., 'In Step {{ step_number }}, I tapped... and observed...').\n"
|
|
"Strictly avoid subjective validation words: successfully, completed, failed, achieved, navigated to."
|
|
)
|
|
|
|
def dispatch(
|
|
self,
|
|
step_number: int,
|
|
action_name: str,
|
|
action_args: dict[str, Any],
|
|
pre_img_bytes: bytes | None,
|
|
post_img_bytes: bytes | None,
|
|
exec_outcome: str,
|
|
*,
|
|
action_key: str | None = None,
|
|
data_engine_step_id: UUID | str | None = None,
|
|
) -> None:
|
|
"""Dispatches an asynchronous summarization task without blocking the caller."""
|
|
key: int | str = action_key or step_number
|
|
self._step_inputs[key] = {
|
|
"step_number": step_number,
|
|
"action_name": action_name,
|
|
"action_args": action_args,
|
|
"pre_img_bytes": pre_img_bytes,
|
|
"post_img_bytes": post_img_bytes,
|
|
"exec_outcome": exec_outcome,
|
|
"data_engine_step_id": data_engine_step_id,
|
|
}
|
|
|
|
task = asyncio.create_task(self._run_summary_until_ready(key))
|
|
self._pending_tasks[key] = task
|
|
|
|
async def _run_summary_until_ready(self, action_key: int | str) -> None:
|
|
"""Retry one action independently until a non-empty summary is committed."""
|
|
while action_key not in self._summaries:
|
|
if await self._run_summary_once(action_key):
|
|
return
|
|
|
|
retry_count = self._retry_counts.get(action_key, 0) + 1
|
|
self._retry_counts[action_key] = retry_count
|
|
delay = self._retry_delays[min(retry_count - 1, len(self._retry_delays) - 1)]
|
|
input_data = self._step_inputs.get(action_key)
|
|
display_step = input_data.get("step_number") if input_data else action_key
|
|
logger.info(
|
|
f"VisualStepSummarizer: Retrying Step {display_step} "
|
|
f"(attempt {retry_count + 1}) after {delay:.1f}s"
|
|
)
|
|
if delay:
|
|
await asyncio.sleep(delay)
|
|
|
|
async def _run_summary_once(self, action_key: int | str) -> bool:
|
|
"""Execute one lightweight VLM attempt for a step transition."""
|
|
input_data = self._step_inputs.get(action_key)
|
|
if not input_data:
|
|
return False
|
|
|
|
step_number = input_data["step_number"]
|
|
action_name = input_data["action_name"]
|
|
action_args = input_data["action_args"]
|
|
pre_bytes = input_data["pre_img_bytes"]
|
|
post_bytes = input_data["post_img_bytes"]
|
|
exec_outcome = input_data["exec_outcome"]
|
|
|
|
try:
|
|
rendered_prompt = Template(self._prompt_template).render(step_number=step_number)
|
|
content_blocks: list[dict[str, Any]] = [
|
|
{
|
|
"type": "text",
|
|
"text": (
|
|
f"Step {step_number} Physical Action: {action_name}({action_args})\n"
|
|
f"Controller Outcome: {exec_outcome}"
|
|
),
|
|
}
|
|
]
|
|
|
|
if pre_bytes:
|
|
# 🎨 Visually mark the exact action (tap ripple, sequence numbers, swipe arrow) on the BEFORE screenshot
|
|
annotated_pre_bytes = draw_action_overlay_on_image(
|
|
image_bytes=pre_bytes,
|
|
action_name=action_name,
|
|
action_args=action_args,
|
|
)
|
|
b64_pre = base64.b64encode(annotated_pre_bytes).decode("utf-8")
|
|
content_blocks.append(
|
|
{
|
|
"type": "text",
|
|
"text": "--- [1] BEFORE ACTION SCREEN (Action Marked Visually in Red) ---",
|
|
}
|
|
)
|
|
content_blocks.append(
|
|
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_pre}"}}
|
|
)
|
|
|
|
if post_bytes:
|
|
b64_post = base64.b64encode(post_bytes).decode("utf-8")
|
|
content_blocks.append({"type": "text", "text": "--- [2] AFTER ACTION SCREEN ---"})
|
|
content_blocks.append(
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/jpeg;base64,{b64_post}"},
|
|
}
|
|
)
|
|
|
|
messages: list[BaseMessage] = [
|
|
SystemMessage(content=rendered_prompt),
|
|
HumanMessage(content=content_blocks),
|
|
]
|
|
|
|
response = await asyncio.wait_for(self._llm.ainvoke(messages), timeout=25.0)
|
|
summary_raw = response.content if isinstance(response.content, str) else ""
|
|
if isinstance(response.content, list):
|
|
summary_raw = "".join(
|
|
b.get("text", "")
|
|
for b in response.content
|
|
if isinstance(b, dict) and "text" in b
|
|
)
|
|
|
|
summary_text = summary_raw.strip()
|
|
if summary_text:
|
|
self._summaries[action_key] = summary_text
|
|
logger.info(
|
|
f"VisualStepSummarizer: Generated summary for Step {step_number}: {summary_text[:80]}..."
|
|
)
|
|
|
|
# Free binary image buffers from memory once summary is secured
|
|
if action_key in self._step_inputs:
|
|
self._step_inputs[action_key]["pre_img_bytes"] = None
|
|
self._step_inputs[action_key]["post_img_bytes"] = None
|
|
|
|
# Update DataEngine telemetry if active
|
|
if self.ctx.data_engine:
|
|
try:
|
|
target_step = input_data.get("data_engine_step_id") or step_number
|
|
self.ctx.data_engine.update_step_summary(target_step, summary_text)
|
|
except Exception as de_err:
|
|
logger.debug(f"DataEngine step summary update skipped: {de_err}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"VisualStepSummarizer: Error generating summary for step {step_number}: {e}"
|
|
)
|
|
return False
|
|
|
|
def get_summary(
|
|
self, action_key: int | str, fallback_text: str | None = None
|
|
) -> str | None:
|
|
"""Retrieve the summary for a stable action key, or fallback text."""
|
|
return self._summaries.get(action_key, fallback_text)
|
|
|
|
def has_summary(self, action_key: int | str) -> bool:
|
|
"""Check whether an action summary has completed."""
|
|
return action_key in self._summaries
|
|
|
|
def has_job(self, action_key: int | str) -> bool:
|
|
"""Check whether an action has been submitted for visual summarization."""
|
|
return action_key in self._step_inputs
|
|
|
|
def is_pending(self, action_key: int | str) -> bool:
|
|
"""Check whether an action still needs a summary."""
|
|
return self.has_job(action_key) and not self.has_summary(action_key)
|
|
|
|
async def flush(self, timeout_seconds: float = 30.0) -> None:
|
|
"""Wait for active jobs up to a bound, then cancel remaining retry loops."""
|
|
tasks = [t for t in self._pending_tasks.values() if not t.done()]
|
|
if tasks:
|
|
logger.info(f"VisualStepSummarizer: Flushing {len(tasks)} pending summary tasks...")
|
|
_, pending = await asyncio.wait(tasks, timeout=timeout_seconds)
|
|
if pending:
|
|
logger.warning(
|
|
f"VisualStepSummarizer: Cancelling {len(pending)} unfinished summary tasks "
|
|
f"after {timeout_seconds:.1f}s flush timeout."
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
await asyncio.gather(*pending, return_exceptions=True)
|