# 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. import base64 import json from pathlib import Path from typing import Any from langchain_core.messages import ToolMessage from langchain_core.tools import BaseTool from langgraph.types import Command from pydantic import BaseModel from artemis.constants import VALIDATOR_MESSAGES_KEY from artemis.context import ArtemisContext from artemis.data_engine.trace import trace_langchain_tool from artemis.drivers.base import BaseDeviceDriver from artemis.graph.state import State from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry from artemis.tools.tool_wrapper import ToolWrapper from artemis.utils.logger import get_logger from artemis.utils.ocr_api import perform_ocr class OcrArgs(BaseModel): """Arguments schema for OCR recognition.""" logger = get_logger(__name__) OCR_DOCSTRING = ( "[EXPLORER] Performs OCR on the current screen and returns the detected" " text along with their coordinates." ) # pylint: disable=too-many-branches,too-many-locals,too-many-statements async def run_ocr_core( ctx: ArtemisContext | None = None, state: State | None = None, driver: BaseDeviceDriver | None = None, ) -> str: """Universal OCR core function: reads screenshot from state or driver and returns results.""" screenshot_b64 = None screenshot_path = getattr(state, "latest_screenshot", None) if state else None if screenshot_path: if not Path(screenshot_path).exists(): logger.error(f"Screenshot file does not exist: {screenshot_path}") raise FileNotFoundError(f"Screenshot file does not exist: {screenshot_path}") try: with open(screenshot_path, "rb") as f: screenshot_b64 = base64.b64encode(f.read()).decode("utf-8") logger.info(f"OCR loaded screenshot from file: {screenshot_path}") except Exception as e: # pylint: disable=broad-exception-caught logger.error(f"Failed to read screenshot from {screenshot_path}: {e}") raise e elif driver is not None and hasattr(driver, "get_screen_data"): try: screen_data = await driver.get_screen_data() if hasattr(screen_data, "screenshot_base64") and screen_data.screenshot_base64: screenshot_b64 = screen_data.screenshot_base64 elif hasattr(screen_data, "screenshot_bytes") and screen_data.screenshot_bytes: screenshot_b64 = base64.b64encode(screen_data.screenshot_bytes).decode("utf-8") elif hasattr(screen_data, "base64") and screen_data.base64: screenshot_b64 = screen_data.base64 except Exception as e: # pylint: disable=broad-exception-caught logger.error(f"Failed to capture screenshot from driver: {e}") raise e else: logger.error("No screenshot path found in state.latest_screenshot and no driver available") raise ValueError( "No screenshot path found in state.latest_screenshot and no driver available" ) if not screenshot_b64: raise ValueError("Failed to obtain screenshot for OCR.") width = 1080 height = 2400 try: if ctx and hasattr(ctx, "device") and ctx.device: w = getattr(ctx.device, "device_width", 1080) h = getattr(ctx.device, "device_height", 2400) if isinstance(w, int): width = w if isinstance(h, int): height = h elif driver and hasattr(driver, "screen_size") and driver.screen_size: width, height = driver.screen_size except Exception as e: # pylint: disable=broad-exception-caught logger.warning(f"Failed to get device resolution for OCR normalization: {e}") results = await perform_ocr(screenshot_b64) if results: normalized_results = [] for item in results: text = item.get("text") vertices = item.get("position") or [] # Calculate center point x_coords = [v.get("x", 0) for v in vertices] y_coords = [v.get("y", 0) for v in vertices] if x_coords and y_coords: min_x, max_x = min(x_coords), max(x_coords) min_y, max_y = min(y_coords), max(y_coords) center_x = (min_x + max_x) / 2 center_y = (min_y + max_y) / 2 # Normalize to 0-1000 norm_x = int(max(0, min(1000, center_x * 1000 / width))) norm_y = int(max(0, min(1000, center_y * 1000 / height))) normalized_results.append({"text": text, "coordinates": [norm_x, norm_y]}) else: normalized_results.append({"text": text, "coordinates": None}) return json.dumps(normalized_results, ensure_ascii=False, indent=2) return "No text detected on the screen." class OcrRecognitionTool(ArtemisTool): """Universal tool for performing OCR on the current screen.""" def __init__(self, category: ToolCategory = "explorer"): super().__init__( name="ocr_recognition", description=OCR_DOCSTRING, args_schema=OcrArgs, category=category, ) # pylint: disable=too-many-arguments,too-many-positional-arguments async def execute( self, driver: BaseDeviceDriver | None = None, ctx: ArtemisContext | None = None, tool_call_id: str | None = None, state: State | None = None, **kwargs: Any, ) -> Any: tcid = tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id") st = state if state is not None else kwargs.get("state") has_failed = False output = "" error_message = None try: output = await run_ocr_core(ctx=ctx, state=st, driver=driver) except Exception as e: # pylint: disable=broad-exception-caught has_failed = True error_message = f"Failed to perform OCR: {str(e)}" if st and callable(getattr(st, "asanitize_update", None)): tool_message = ToolMessage( tool_call_id=tcid or "", content=ocr_recognition_wrapper.on_failure_fn(error_message) if has_failed else ocr_recognition_wrapper.on_success_fn(output), additional_kwargs={"error": error_message} if has_failed else {}, status="error" if has_failed else "success", ) return Command( update=await st.asanitize_update( ctx=ctx, update={ VALIDATOR_MESSAGES_KEY: [tool_message], }, agent="validator", ) ) if has_failed: return ocr_recognition_wrapper.on_failure_fn(error_message) return ocr_recognition_wrapper.on_success_fn(output) # Universal tool instance & aliases ocr_recognition = OcrRecognitionTool() OcrRecognition = OcrRecognitionTool OCRRecognition = OcrRecognitionTool OcrTool = OcrRecognitionTool ToolRegistry.register(ocr_recognition) def get_ocr_tool(ctx: ArtemisContext) -> BaseTool: """Exports ocr_recognition as a LangChain BaseTool.""" return trace_langchain_tool(ocr_recognition.to_langchain_tool(ctx), ctx) ocr_recognition_wrapper = ToolWrapper( tool_fn_getter=get_ocr_tool, on_success_fn=lambda output: f"OCR Recognition successful. Results:\n{output}", on_failure_fn=lambda error: f"OCR Recognition failed: {error}", )