Artemis-Eapil/artemis/agents/flash/summarizer.py

# 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 includes out-of-order self-healing retries.
"""

import asyncio
import base64
from pathlib import Path
from typing import Any

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. Graceful fallback: If a step summary is not ready, it cleanly falls back to the original action text.
    3. Self-healing re-trigger: If a previous step's summary is stalled or lost while subsequent steps
       complete, it automatically re-dispatches the stalled step.
    """

    def __init__(self, ctx: ArtemisContext, model_name: str | None = None):
        self.ctx = ctx
        self._summaries: dict[int, str] = {}
        self._pending_tasks: dict[int, asyncio.Task] = {}
        self._step_inputs: dict[int, dict[str, Any]] = {}
        self._retry_counts: dict[int, int] = {}
        self._max_retries = 2

        # 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,
    ) -> None:
        """Dispatches an asynchronous summarization task without blocking the caller."""
        self._step_inputs[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,
        }

        task = asyncio.create_task(self._run_summary(step_number))
        self._pending_tasks[step_number] = task

        # Check for and re-trigger any previous stalled steps
        self._check_and_retrigger_stalled_steps(current_step=step_number)

    def _check_and_retrigger_stalled_steps(self, current_step: int) -> None:
        """Re-dispatches any previous steps that remain unsummarized while later steps arrive."""
        for past_step in range(1, current_step):
            if past_step not in self._summaries and past_step in self._step_inputs:
                task = self._pending_tasks.get(past_step)
                retries = self._retry_counts.get(past_step, 0)
                if (task is None or task.done()) and retries < self._max_retries:
                    self._retry_counts[past_step] = retries + 1
                    logger.info(
                        f"VisualStepSummarizer: Retriggering stalled summary for Step {past_step} "
                        f"(attempt {self._retry_counts[past_step] + 1})"
                    )
                    self._pending_tasks[past_step] = asyncio.create_task(
                        self._run_summary(past_step)
                    )

    async def _run_summary(self, step_number: int) -> None:
        """Executes the lightweight VLM call to summarize the step transition."""
        input_data = self._step_inputs.get(step_number)
        if not input_data:
            return

        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[step_number] = 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 step_number in self._step_inputs:
                    self._step_inputs[step_number]["pre_img_bytes"] = None
                    self._step_inputs[step_number]["post_img_bytes"] = None

                # Update DataEngine telemetry if active
                if self.ctx.data_engine:
                    try:
                        self.ctx.data_engine.update_step_summary(step_number, summary_text)
                    except Exception as de_err:
                        logger.debug(f"DataEngine step summary update skipped: {de_err}")

        except Exception as e:
            logger.warning(
                f"VisualStepSummarizer: Error generating summary for step {step_number}: {e}"
            )
            if self._retry_counts.get(step_number, 0) >= self._max_retries:
                # Max retries reached, free memory buffer
                if step_number in self._step_inputs:
                    self._step_inputs[step_number]["pre_img_bytes"] = None
                    self._step_inputs[step_number]["post_img_bytes"] = None

    def get_summary(self, step_number: int, fallback_text: str | None = None) -> str | None:
        """Retrieves summarized state for step_number, or fallback_text if not ready yet."""
        return self._summaries.get(step_number, fallback_text)

    def has_summary(self, step_number: int) -> bool:
        """Checks if a step summary has completed."""
        return step_number in self._summaries

    async def flush(self) -> None:
        """Flushes and awaits all pending background summary tasks upon completion."""
        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...")
            await asyncio.gather(*tasks, return_exceptions=True)