refactor(explorer): split explorer into facade and phase modules

artemis/agents/explorer/run_setup.py

@@ -0,0 +1,384 @@
+# 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.
+
+"""Setup phases of ``Explorer.run``.
+
+Split out of ``artemis.agents.explorer.explorer``: the flash-mode fast path
+and the named preparation phases executed before the reasoning loop (image
+record resolution, version limits, screen dimensions, image pool, GenAI
+client, model parameters, prompt construction, and initial screenshot
+annotation), packaged as a mixin consumed by ``Explorer``.  Patched
+collaborators (``settings``, ``StorageManager``, ``genai``,
+``is_ocr_configured``, ``perform_ocr``, ``draw_dots``,
+``_run_object_detection``, ``logger``) are resolved through the facade module
+at call time; see ``artemis.agents.explorer._facade``.
+"""
+
+import base64
+import glob
+import hashlib
+import json
+import os
+from pathlib import Path
+import re
+
+from artemis.agents.explorer._facade import facade
+from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS
+from artemis.agents.explorer.perception_tools import load_detector_templates
+from artemis.graph.state import State
+from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml
+from artemis.utils.visualization import format_minimal_list_with_points
+
+
+class RunSetupMixin:
+    """Flash mode and pre-loop setup phases of :class:`Explorer`."""
+
+    async def _run_flash(self, query: str, screenshot_path: str) -> str:
+        """Flash mode: single-shot object detection without a reasoning loop."""
+        _ex = facade()
+        templates, global_timeout = load_detector_templates()
+
+        self.screenshot_path = screenshot_path
+        self.image_name = None
+
+        queries = [q.strip() for q in query.split("|") if q.strip()]
+
+        try:
+            result = await _ex._run_object_detection(
+                self.ctx,
+                screenshot_path,
+                queries,
+                templates,
+                global_timeout=global_timeout,
+            )
+            detected_items = result.get("detected", [])
+            candidates = []
+            for idx, item in enumerate(detected_items):
+                pos = item.get("point")
+                if pos and isinstance(pos, list) and len(pos) == 2:
+                    candidates.append(
+                        {
+                            "label": f"D{idx + 1}",
+                            "coords": pos,
+                            "description": item.get("label", query),
+                        }
+                    )
+            fallback_message = "" if candidates else f"Failed to detect: {query}"
+            return json.dumps(
+                {
+                    "candidates": candidates,
+                    "fallback_message": fallback_message,
+                },
+                ensure_ascii=False,
+            )
+        except Exception as e:
+            _ex.logger.error(f"Flash mode object detection failed: {e}")
+            return json.dumps(
+                {
+                    "candidates": [],
+                    "fallback_message": f"Flash mode detection error: {e}",
+                },
+                ensure_ascii=False,
+            )
+
+    def _resolve_image_record(self, screenshot_path: str) -> tuple:
+        """Computes the screenshot hash and looks it up in the Data Engine DB."""
+        _ex = facade()
+        image_name = None
+        record = None
+        try:
+            sha256_hash = hashlib.sha256()
+            with open(screenshot_path, "rb") as f:
+                for byte_block in iter(lambda: f.read(4096), b""):
+                    sha256_hash.update(byte_block)
+            computed_hash = sha256_hash.hexdigest()
+            _ex.logger.info(f"Computed screenshot hash: {computed_hash}")
+
+            # Check if it exists in DB
+            db_path = _ex.settings.DATA_ENGINE_DB_PATH
+            base_dir = _ex.settings.TRACES_PATH
+            storage = _ex.StorageManager(db_path, base_dir)
+            record = storage.get_image(computed_hash)
+
+            if record:
+                image_name = computed_hash
+            else:
+                _ex.logger.warning(
+                    f"Image hash {computed_hash} not found in Data Engine DB."
+                    " Data Engine might not be synced yet."
+                )
+
+        except Exception as e:
+            _ex.logger.warning(f"Failed to compute hash or check DB: {e}")
+
+        return image_name, record
+
+    def _resolve_version_limits(self, version: str) -> tuple[str, int]:
+        """Resolves mode label, iteration budget, and the effective tool denylist."""
+        _ex = facade()
+        mode = version.capitalize()
+        version_denylist = set()
+        if version == "ultra":
+            max_iterations = 8
+        elif version == "pro":
+            max_iterations = 3
+            version_denylist = {
+                "ask_image_processor",
+                "get_ocr_list",
+                "inspect_region",
+                "detect_objects",
+            }
+        else:
+            max_iterations = 3
+            version_denylist = {
+                "ask_image_processor",
+                "get_ocr_list",
+                "inspect_region",
+            }
+
+        if not _ex.is_ocr_configured():
+            version_denylist.add("get_ocr_list")
+
+        try:
+            denylisted_config = (
+                self.ctx.agent_config.denylisted_tools.get("explorer", [])
+                if self.ctx.agent_config
+                else []
+            )
+            self.denylisted_tools = version_denylist.union(set(denylisted_config))
+        except (TypeError, AttributeError):
+            self.denylisted_tools = version_denylist
+
+        return mode, max_iterations
+
+    def _resolve_screen_dimensions(self, state: State) -> None:
+        """Resolves device screen dimensions from operator data or the device."""
+        ctx = self.ctx
+        self.width = 1080
+        self.height = 2400
+        operator_raw_data = getattr(state, "operator_raw_data", {}) or {}
+        w_raw = operator_raw_data.get("width")
+        h_raw = operator_raw_data.get("height")
+        if isinstance(w_raw, int) and isinstance(h_raw, int):
+            self.width = w_raw
+            self.height = h_raw
+        else:
+            if ctx.device and getattr(ctx.device, "device_width", None):
+                self.width = ctx.device.device_width
+            if ctx.device and getattr(ctx.device, "device_height", None):
+                self.height = ctx.device.device_height
+
+    def _init_image_pool(self, screenshot_path: str) -> None:
+        """Seeds the Image Pool with the original screenshot as img_0."""
+        self.image_pool = {
+            "img_0": {
+                "path": screenshot_path,
+                "transform": {
+                    "offset_x": 0.0,
+                    "offset_y": 0.0,
+                    "scale_x": 1.0,
+                    "scale_y": 1.0,
+                },
+                "description": "Original complete screenshot",
+            }
+        }
+        self.next_img_id = 1
+
+    def _ensure_genai_client(self):
+        """Initializes or reuses the context-level GenAI client for connection pooling."""
+        _ex = facade()
+        ctx = self.ctx
+        client = getattr(ctx, "_genai_client", None)
+        if client is None:
+            _ex.logger.info(
+                "Initializing new GenAI client on context for connection pooling (Explorer)..."
+            )
+            client = _ex.genai.Client(
+                api_key=_ex.settings.GOOGLE_API_KEY.get_secret_value()
+                if _ex.settings.GOOGLE_API_KEY
+                else None
+            )
+            ctx._genai_client = client
+        return client
+
+    def _resolve_model_params(self) -> tuple:
+        """Resolves model name, temperature, thinking level, and fallback model."""
+        llm_cfg = self.ctx.llm_config.explorer
+        temperature = 0.1
+        fallback_model = None
+        thinking_level = None
+        if llm_cfg:
+            model_name = llm_cfg.model
+            if "/" in model_name:
+                model_name = model_name.split("/")[-1]
+            if getattr(llm_cfg, "temperature", None) is not None:
+                temperature = llm_cfg.temperature
+            if getattr(llm_cfg, "thinking_level", None) is not None:
+                thinking_level = llm_cfg.thinking_level
+            if getattr(llm_cfg, "fallback", None):
+                fallback_model = llm_cfg.fallback.model
+                if "/" in fallback_model:
+                    fallback_model = fallback_model.split("/")[-1]
+        else:
+            model_name = "gemini-3.7-flash"
+
+        return model_name, temperature, thinking_level, fallback_model
+
+    def _build_prompt_template(
+        self, version: str, mode: str, max_iterations: int
+    ) -> tuple[str | None, str | None]:
+        """Builds the system prompt template; returns (prompt, error_message)."""
+        prompt_path = Path(__file__).parent / "explorer.json"
+        if not prompt_path.exists():
+            return None, "Error: Explorer prompt template not found."
+
+        try:
+            content = prompt_path.read_text(encoding="utf-8")
+            content_no_comments = re.sub(r"(?<!:)\/\/.*", "", content)
+            content_no_comments = re.sub(r"/\*.*?\*/", "", content_no_comments, flags=re.DOTALL)
+            content_no_comments = re.sub(r",\s*([\]}])", r"\1", content_no_comments)
+            data = json.loads(content_no_comments)
+        except Exception as e:
+            return None, f"Error loading or parsing explorer.json: {e}"
+
+        prompt_parts = []
+        for section, content_val in data.items():
+            prompt_parts.append(f"# {section}")
+            if isinstance(content_val, list):
+                for bullet in content_val:
+                    prompt_parts.append(f"- {bullet}")
+            else:
+                prompt_parts.append(content_val)
+            prompt_parts.append("")
+
+        prompt_template = "\n".join(prompt_parts)
+        prompt_template += "\n{denylist}"
+        if self.denylisted_tools:
+            denylist_section = self.DENYLIST_TEMPLATE.format(
+                tools=", ".join(sorted(self.denylisted_tools))
+            )
+        else:
+            denylist_section = ""
+        prompt_template = prompt_template.replace("{denylist}", denylist_section)
+
+        explore_info = EXPLORE_DESCRIPTIONS.get(version, EXPLORE_DESCRIPTIONS["pro"])
+        version_prompt = explore_info.get("version_prompt")
+        if version_prompt:
+            version_prompt = version_prompt.format(mode=mode, max_iterations=max_iterations)
+            prompt_template += f"\n# EXECUTION CONSTRAINT\n- {version_prompt}\n"
+
+        return prompt_template, None
+
+    async def _build_fused_hierarchy(self, record, state: State, screenshot_path: str):
+        """Loads (and if needed OCR-fuses) the UI hierarchy for the screenshot."""
+        _ex = facade()
+        fused_xml = []
+        if record and getattr(record, "ui_tree", None):
+            ui_tree = record.ui_tree
+            ocr_results = getattr(record, "ocr_result", None)
+            if ocr_results is None:
+                if _ex.is_ocr_configured():
+                    try:
+                        _ex.logger.info(
+                            "Previous screenshot OCR is missing. Running OCR on-the-fly..."
+                        )
+                        with open(screenshot_path, "rb") as img_file:
+                            img_b64 = base64.b64encode(img_file.read()).decode("utf-8")
+                        ocr_results = await _ex.perform_ocr(img_b64, client=self.http_client)
+                    except Exception as ocr_err:
+                        _ex.logger.error(
+                            f"On-the-fly OCR failed for previous screenshot: {ocr_err}"
+                        )
+                        ocr_results = []
+                else:
+                    ocr_results = []
+
+            fused_xml = fuse_ocr_with_xml(ui_tree, ocr_results or [])
+            _ex.logger.info("Successfully loaded and fused UI hierarchy for previous screenshot.")
+        elif hasattr(state, "latest_ui_hierarchy"):
+            if screenshot_path == getattr(state, "latest_screenshot", None):
+                fused_xml = state.latest_ui_hierarchy
+        return fused_xml
+
+    def _annotate_initial_screenshot(
+        self, fused_xml, screenshot_path: str, image_name, minimal_list: str
+    ) -> tuple[str, object]:
+        """Draws labeled dots for the fused hierarchy; returns (minimal_list, marked_path)."""
+        _ex = facade()
+        ctx = self.ctx
+        marked_path = None
+        try:
+            _ex.logger.info(
+                "Explorer self-annotating initial screenshot using latest_ui_hierarchy..."
+            )
+            formatted_list, points, labels = format_minimal_list_with_points(
+                fused_xml, self.width, self.height
+            )
+            minimal_list = formatted_list
+            self.global_label_idx = len(points) + 1
+
+            base_dir = (
+                Path(ctx.data_engine.base_dir)
+                if ctx.data_engine and getattr(ctx.data_engine, "base_dir", None)
+                else None
+            )
+            if not base_dir:
+                db_path = _ex.settings.DATA_ENGINE_DB_PATH
+                base_dir = _ex.settings.TRACES_PATH
+            images_dir = base_dir / "images"
+            initial_marked_dir = images_dir / "initial_marked"
+            initial_marked_dir.mkdir(parents=True, exist_ok=True)
+
+            existing_files = glob.glob(
+                str(initial_marked_dir / f"{image_name or 'temp_image'}_*.jpg")
+            )
+            max_seq = 0
+            for f in existing_files:
+                match = re.search(r"_(\d+)\.jpg$", f)
+                if match:
+                    max_seq = max(max_seq, int(match.group(1)))
+            seq = max_seq + 1
+            marked_path = initial_marked_dir / f"{image_name or 'temp_image'}_{seq}.jpg"
+
+            # Draw dots on the raw screenshot
+            _ex.draw_dots(screenshot_path, points, labels, str(marked_path))
+            _ex.logger.info(
+                f"Successfully drew {len(points)} dots and saved marked"
+                f" image to {marked_path}"
+            )
+        except Exception as e:
+            _ex.logger.error(f"Failed to self-annotate initial screenshot: {e}")
+        return minimal_list, marked_path
+
+    async def _prepare_initial_annotation(
+        self, record, state: State, screenshot_path: str, minimal_list: str, image_name
+    ) -> tuple[str, str]:
+        """Generates initial visual annotations when no minimal list is provided.
+
+        Returns the (possibly updated) minimal list and the image path the
+        model should read (marked screenshot when available, raw otherwise).
+        """
+        marked_path = None
+        if not minimal_list:
+            fused_xml = await self._build_fused_hierarchy(record, state, screenshot_path)
+            if fused_xml:
+                minimal_list, marked_path = self._annotate_initial_screenshot(
+                    fused_xml, screenshot_path, image_name, minimal_list
+                )
+
+        if marked_path and os.path.exists(str(marked_path)):
+            img_to_read = str(marked_path)
+        else:
+            img_to_read = screenshot_path
+        return minimal_list, img_to_read